From 35fb6a25a11afa1beaa29e184b1b5e58c228535b Mon Sep 17 00:00:00 2001
From: Yago Ascasibar <yago.ascasibar@uam.es>
Date: Mon, 7 Oct 2024 13:20:22 +0200
Subject: [PATCH] A few check_units

---
 src/pst/models.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/src/pst/models.py b/src/pst/models.py
index 3885b3e..e6bf18b 100644
--- a/src/pst/models.py
+++ b/src/pst/models.py
@@ -216,16 +216,10 @@ class SingleBurstCEM(ChemicalEvolutionModel):
         Metallicity of the burst.
     """
     def __init__(self, **kwargs):
-        self.mass_burst = kwargs['mass_burst']
-        if not isinstance(self.mass_burst, u.Quantity):
-            self.mass_burst *= u.Msun
-        self.time_burst = kwargs['time_burst']
-        if not isinstance(self.time_burst, u.Quantity):
-            self.time_burst *= u.Gyr
-
+        self.mass_burst = check_units(kwargs['mass_burst'], u.Msun)
+        self.time_burst = check_units(kwargs['time_burst'], u.Gyr)
         self.burst_metallicity = kwargs.get("burst_metallicity",
                                             0.02 * u.dimensionless_unscaled)
-
         ChemicalEvolutionModel.__init__(self, **kwargs)
 
     @u.quantity_input
@@ -434,7 +428,7 @@ class LogNormalQuenchedCEM(LogNormalZPowerLawCEM):
     """
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.quenching_time = kwargs["quenching_time"]
+        self.quenching_time = check_unit(kwargs["quenching_time"], u.Gyr)
         mtoday = self.stellar_mass_formed(self.today)
         self.mass_norm *= self.mass_today / mtoday