Reactoring of quenching decorator + bugfixes

src/pst/models.py

@@ -44,8 +44,9 @@ def sfh_quenching_decorator(stellar_mass_formed):
     def wrapper_stellar_mass_formed(*args):
         quenching_time = getattr(args[0], "quenching_time", 20.0 << u.Gyr)
         stellar_mass = stellar_mass_formed(*args)
-        return np.clip(stellar_mass / stellar_mass[args[1] <= quenching_time].max(),
-                0, 1) * stellar_mass[-1]
+        final_mass = stellar_mass_formed(args[0], quenching_time)
+        stellar_mass[args[1] > quenching_time] = final_mass
+        return stellar_mass
     return wrapper_stellar_mass_formed
 
 
@@ -392,22 +393,27 @@ class LogNormalCEM(ChemicalEvolutionModel):
     metallicity : float
         Metallicity of the gas (constant).
     """
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.lnt0 = kwargs['lnt0']
+        self.t0 = check_unit(kwargs['t0'], u.Gyr)
         self.scale = kwargs['scale']
+        self.today = check_unit(kwargs['today'], u.Gyr)
         self.mass_today = check_unit(kwargs['mass_today'], u.Msun)
-        self.metallicity = kwargs.get('metallicity', np.nan)
+        self.mass_norm = 1
+        mtoday = self.stellar_mass_formed(self.today)
+        self.mass_norm = self.mass_today / mtoday
+        self.metallicity_today = kwargs.get('metallicity_today', np.nan)
 
     @u.quantity_input
     def stellar_mass_formed(self, times: u.Quantity):
-        z = - (np.log(times.to_value("Gyr")) - self.lnt0) / self.scale
+        z = - np.log(times / self.t0) / self.scale
         m = 0.5 * (1 - special.erf(z / SQRT_2))
-        return m / m.max() * self.mass_today
+        return m * self.mass_norm
 
     @u.quantity_input
     def ism_metallicity(self, time : u.Gyr):
-        return np.full(time.size, fill_value=self.metallicity)
+        return np.full(time.size, fill_value=self.metallicity_today)
 
 
 class LogNormalZPowerLawCEM(MassPropMetallicityMixin, LogNormalCEM):
@@ -429,6 +435,8 @@ class LogNormalQuenchedCEM(LogNormalZPowerLawCEM):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.quenching_time = kwargs["quenching_time"]
+        mtoday = self.stellar_mass_formed(self.today)
+        self.mass_norm *= self.mass_today / mtoday
 
     @sfh_quenching_decorator
     def stellar_mass_formed(self, times: u.Quantity):
@@ -535,7 +543,7 @@ def __init__(self, times, masses, alpha_powerlaw, ism_metallicity_today, **kwarg
         super().__init__(times, masses, **kwargs)
 
 
-class ParticleGridCEM(ChemicalEvolutionModel):
+class ParticleListCEM(ChemicalEvolutionModel):
     """
     Chemical Evolution Model using individual Simple Stellar Population (SSP) data.
 

tests/test_models.py

@@ -52,14 +52,15 @@ def test_exponential_quenched(self):
 
     def test_lognormal_zpowerlaw(self):
         model = models.LogNormalZPowerLawCEM(
-            lnt0=0.5, scale=1.0, mass_today=1.0,
+            t0=3.0, scale=1.0, mass_today=1.0,
+            today=13.7,
             ism_metallicity_today=0.02, alpha_powerlaw=2.0
         )
         mass = model.stellar_mass_formed(self.dummy_times)
         metals = model.ism_metallicity(self.dummy_times)
 
         self.assertEqual(mass[0], 0.0)
-        self.assertEqual(mass[-1], 1.0 * u.Msun)
+        self.assertTrue(np.isclose(mass[-1], 1.0 * u.Msun, rtol=1e-4))
         self.assertEqual(metals[0], 0.0)
         self.assertEqual(metals[-1], 0.02)
 
@@ -79,7 +80,7 @@ def test_particle_grid(self):
         particles_z = 10**(np.random.uniform(-4, 0.3, n_particles))
         particles_t_form = np.random.exponential(3, n_particles)
         particles_mass = 10**(np.random.uniform(5, 6, n_particles))
-        model = models.ParticleGridCEM(
+        model = models.ParticleListCEM(
             time_form=particles_t_form * u.Gyr,
             metallicities=particles_z * u.dimensionless_unscaled,
             masses=particles_mass * u.Msun)

tutorials/models/analytical.ipynb

@@ -65,7 +65,7 @@
        " pst.models.GaussianBurstCEM,\n",
        " pst.models.LogNormalCEM,\n",
        " pst.models.TabularCEM,\n",
-       " pst.models.ParticleGridCEM,\n",
+       " pst.models.ParticleListCEM,\n",
        " pst.fitting_module.Polynomial_MFH]"
       ]
      },
@@ -238,7 +238,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 4,
    "id": "dc9fbc91-da50-40b8-992b-3337e56eb98d",
    "metadata": {},
    "outputs": [],
@@ -247,9 +247,10 @@
     "mass_today = 1e11 << u.Msun\n",
     "ism_metallicity_today = 0.02 << u.dimensionless_unscaled\n",
     "alpha_powerlaw = 1.5\n",
-    "lnt0 = np.log(7.0)\n",
+    "t0 = 7.0 << u.Gyr\n",
     "scale = 2.0\n",
-    "quenching_time = 10 << u.Gyr"
+    "quenching_time = 10 << u.Gyr\n",
+    "today = 13.7 << u.Gyr"
    ]
   },
   {
@@ -262,13 +263,28 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 5,
    "id": "1b80ee87-53d4-4499-a667-0b4356204ef7",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "ename": "KeyError",
+     "evalue": "'today'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m model \u001b[38;5;241m=\u001b[39m models\u001b[38;5;241m.\u001b[39mLogNormalQuenchedCEM(\n\u001b[1;32m      2\u001b[0m     mass_today\u001b[38;5;241m=\u001b[39mmass_today, t0\u001b[38;5;241m=\u001b[39mt0, scale\u001b[38;5;241m=\u001b[39mscale, quenching_time\u001b[38;5;241m=\u001b[39mquenching_time,\n\u001b[1;32m      3\u001b[0m     alpha_powerlaw\u001b[38;5;241m=\u001b[39malpha_powerlaw, ism_metallicity_today\u001b[38;5;241m=\u001b[39mism_metallicity_today,\n\u001b[1;32m      4\u001b[0m )\n",
+      "File \u001b[0;32m~/Develop/population-synthesis-toolkit/src/pst/models.py:437\u001b[0m, in \u001b[0;36mLogNormalQuenchedCEM.__init__\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m    436\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m--> 437\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    438\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mquenching_time \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mquenching_time\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m    439\u001b[0m     mtoday \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstellar_mass_formed(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtoday)\n",
+      "File \u001b[0;32m~/Develop/population-synthesis-toolkit/src/pst/models.py:428\u001b[0m, in \u001b[0;36mLogNormalZPowerLawCEM.__init__\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m    427\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m--> 428\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    429\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mism_metallicity_today \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mism_metallicity_today\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m    430\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39malpha_powerlaw \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124malpha_powerlaw\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n",
+      "File \u001b[0;32m~/Develop/population-synthesis-toolkit/src/pst/models.py:401\u001b[0m, in \u001b[0;36mLogNormalCEM.__init__\u001b[0;34m(self, **kwargs)\u001b[0m\n\u001b[1;32m    399\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mt0 \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mt0\u001b[39m\u001b[38;5;124m'\u001b[39m]\n\u001b[1;32m    400\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mscale \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mscale\u001b[39m\u001b[38;5;124m'\u001b[39m]\n\u001b[0;32m--> 401\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtoday \u001b[38;5;241m=\u001b[39m check_unit(kwargs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtoday\u001b[39m\u001b[38;5;124m'\u001b[39m], u\u001b[38;5;241m.\u001b[39mGyr)\n\u001b[1;32m    402\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmass_today \u001b[38;5;241m=\u001b[39m check_unit(kwargs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmass_today\u001b[39m\u001b[38;5;124m'\u001b[39m], u\u001b[38;5;241m.\u001b[39mMsun)\n\u001b[1;32m    404\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmass_norm \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n",
+      "\u001b[0;31mKeyError\u001b[0m: 'today'"
+     ]
+    }
+   ],
    "source": [
     "model = models.LogNormalQuenchedCEM(\n",
-    "    mass_today=mass_today, lnt0=lnt0, scale=scale, quenching_time=quenching_time,\n",
+    "    mass_today=mass_today, t0=t0, scale=scale, quenching_time=quenching_time,\n",
     "    alpha_powerlaw=alpha_powerlaw, ism_metallicity_today=ism_metallicity_today,\n",
     ")"
    ]
@@ -283,7 +299,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": null,
    "id": "131ff42d",
    "metadata": {},
    "outputs": [
@@ -341,7 +357,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": null,
    "id": "656d3306",
    "metadata": {},
    "outputs": [
@@ -370,7 +386,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": null,
    "id": "f3231653-6673-4b96-925c-a43f32803d67",
    "metadata": {},
    "outputs": [],
@@ -380,7 +396,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": null,
    "id": "e9c5e622-6dc1-45ee-b80a-b281a8dcfc54",
    "metadata": {},
    "outputs": [
@@ -411,7 +427,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": null,
    "id": "1d035a95-dcae-4a49-bb37-2c172dfe6f06",
    "metadata": {},
    "outputs": [
@@ -443,7 +459,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 35,
+   "execution_count": null,
    "id": "a4b1dd2d",
    "metadata": {},
    "outputs": [
@@ -491,9 +507,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "test-pst",
+   "display_name": "base",
    "language": "python",
-   "name": "test-pst"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -505,7 +521,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,

tutorials/models/lognormal_quenched_sfh.py

@@ -15,14 +15,15 @@
 # Define the parameters of the Chemical Evolution Model
 alpha_powerlaw = -2.0
 ism_metallicity_today = 0.02 * u.dimensionless_unscaled
-lnt0 = np.log(7.0)
+t0 = 7.0 * u.Gyr
 scale = 2.0
 quenching_time = 10 * u.Gyr
 
 model = models.LogNormalQuenchedCEM(
     alpha_powerlaw=alpha_powerlaw, ism_metallicity_today=ism_metallicity_today,
-    mass_today=1 * u.Msun,
-    lnt0=lnt0, scale=scale, quenching_time=quenching_time)
+    mass_today=1234 * u.Msun,
+    today=13.7,
+    t0=t0, scale=scale, quenching_time=quenching_time)
 
 dummy_t = np.linspace(0, 13.7, 1000) * u.Gyr
 

tutorials/models/particles.ipynb

@@ -5,11 +5,11 @@
    "id": "61f042d2",
    "metadata": {},
    "source": [
-    "# Creating a ParticleGridCEM\n",
+    "# Creating a ParticleListCEM\n",
     "\n",
     "In this tutorial we are going to create a Chemical Evolution Model from a set of individual simple stellar populations (referred to as \"particles\"). This is a common way of producing synthetic observations from cosmological hydrodynamical simulations.\n",
     "\n",
-    "For that purpose, PST provides the `ParticleGridCEM` class. The three input arguments required for creating the model are the following:\n",
+    "For that purpose, PST provides the `ParticleListCEM` class. The three input arguments required for creating the model are the following:\n",
     "\n",
     "- `time_form`: cosmic time when each particle was formed\n",
     "- `metallicities`: metallicity of the particles\n",
@@ -130,7 +130,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sfh_model = models.ParticleGridCEM(time_form=particles_t_form, metallicities=particles_z, masses=particles_mass)"
+    "sfh_model = models.ParticleListCEM(time_form=particles_t_form, metallicities=particles_z, masses=particles_mass)"
    ]
   },
   {