Merge pull request #62 from paranoya/feature/add-mass-frac-time

Add mass frac time method

Author: PabloCorcho

src/pst/models.py

@@ -80,8 +80,19 @@ class ChemicalEvolutionModel(ABC):
     (Z-SFR) need to be implemented in a subclass.
     """
 
+    @property
+    def today(self):
+        return self._today
+
+    @today.setter
+    def today(self, value):
+        if value is not None:
+            self._today = check_unit(value, u.Gyr)
+        else:
+            self._today = None
+
     def __init__(self, **kwargs):
-        pass
+        self.today = kwargs.get("today")
 
     @abstractmethod
     def stellar_mass_formed(self, time) -> u.Quantity:
@@ -141,6 +152,48 @@ def interpolate_ssp_masses(self, ssp: SSPBase, t_obs: u.Gyr,
                                metallicities=bin_metallicity,
                                masses=bin_mass)
 
+    def time_at_stellar_mass_frac(self, frac, today=None, time_res=30 << u.Myr):
+        """
+        Get the cosmic time at which the model formed a fraction of the total stellar mass.
+
+        Parameters
+        ----------
+        frac : float
+            Input fraction.
+        today : float or astropy.units.Quantity
+            Age of the Universe at the time of the observation. If not provided,
+            the defult is to use ``self.today'' (returning an error if not set).
+        time_res : float or astropy.units.Quantity
+            Time-step resolution for sampling the SFH.
+        Returns
+        -------
+        frac_times : astropy.units.Quantity
+            Fraction formation times.
+        """
+        frac = np.atleast_1d(frac)
+        if today is None:
+            if self.today is None:
+                raise ValueError(
+                    "current CEM does not have attribute 'today' set,"
+                    + " it must be provided by the user")
+            else:
+                today = self.today
+        else:
+            today = check_unit(today, u.Gyr)
+
+        time_res = check_unit(time_res, u.Gyr)
+        dummy_time = np.arange(0, today.to_value("Gyr"),
+                               time_res.to_value("Gyr")) << u.Gyr
+        mass_history = self.stellar_mass_formed(dummy_time)
+        frac_history = mass_history / mass_history[-1]
+        idx = np.searchsorted(frac_history, frac).clip(
+            min=1, max=frac_history.size - 1)
+        t_frac = (
+            dummy_time[idx - 1] * (frac - frac_history[idx]) / (frac_history[idx] - frac_history[idx - 1])
+         + dummy_time[idx] * (1 - (frac - frac_history[idx]) / (frac_history[idx] - frac_history[idx - 1])
+        ))
+        return t_frac
+
     def compute_SED(self, ssp : SSPBase, t_obs : u.Quantity,
                     allow_negative=False) -> u.Quantity:
         """

tutorials/models/analytical.ipynb

@@ -166,6 +166,11 @@
     "sfr = np.diff(mfh) / np.diff(cosmic_time)\n",
     "t_sfr = (cosmic_time[1:] + cosmic_time[:-1]) / 2\n",
     "\n",
+    "# Estimate the times at which the model formed a given fraction of the present\n",
+    "# mass\n",
+    "mass_fractions = [0.1, 0.5, 0.9, 0.99]\n",
+    "times_frac = model.time_at_stellar_mass_frac(mass_fractions)\n",
+    "\n",
     "fig, axes = plt.subplots(nrows=3, sharex=True)\n",
     "fig.suptitle(\"Log-normal SFH with a quenching event\")\n",
     "fig.supxlabel(f\"cosmic time ({cosmic_time.unit})\")\n",
@@ -175,6 +180,12 @@
     "ax.set_yscale('log')\n",
     "ax.plot(cosmic_time, mfh.to_value(u.Msun), '.-')\n",
     "\n",
+    "for t, m in zip(times_frac, mass_fractions):\n",
+    "    ax.annotate(f\"{m * 100:.1f}%\", xy=(t.value, 1e9), va=\"top\", ha=\"center\",\n",
+    "                arrowprops=dict(facecolor='black', shrink=0.05))\n",
+    "    # ax.arrow(t.value, 1e9, dx=0, dy=3e9, width=0.03,\n",
+    "    #          head_width=0.1, head_length=0.1)\n",
+    "\n",
     "ax = axes[1]\n",
     "ax.set_ylabel(r\"SFR (M$_\\odot$/yr)\")\n",
     "ax.set_yscale('log')\n",
@@ -183,7 +194,7 @@
     "ax = axes[2]\n",
     "ax.set_ylabel(r\"Z\")\n",
     "ax.set_yscale('log')\n",
-    "ax.plot(cosmic_time, model.ism_metallicity(cosmic_time), '.-')\n",
+    "ax.plot(cosmic_time, model.ism_metallicity(cosmic_time), '.-')  \n",
     "\n",
     "for ax in axes:\n",
     "    ax.axvline(quenching_time.to_value(cosmic_time.unit), color='r', label=\"Quenching event\")\n",
@@ -243,7 +254,7 @@
     "ax.set_yscale('log')\n",
     "ax.set_ylim(5e-5, 0.05)\n",
     "\n",
-    "mappable = ax.pcolormesh(ssp.ages.to_value(\"Gyr\"), ssp.metallicities, ssp_weights.to_value(u.Msun), norm=LogNorm())\n",
+    "mappable = ax.pcolormesh(ssp.ages.to_value(\"Gyr\"), ssp.metallicities.value, ssp_weights.to_value(u.Msun), norm=LogNorm())\n",
     "\n",
     "plt.colorbar(mappable, ax=ax, label=r\"SSP mass (M$_\\odot$)\")\n",
     "\n",

tutorials/models/particles.ipynb

@@ -183,8 +183,10 @@
     "ax.set_yscale('log')\n",
     "ax.set_ylim(5e-5, 0.05)\n",
     "\n",
-    "mappable = ax.pcolormesh(ssp.ages.to_value(\"Gyr\"), ssp.metallicities, ssp_weights, norm=LogNorm())\n",
-    "ax.plot((today - particles_t_form).to_value(u.Gyr), particles_z, 'k,', alpha=.25)\n",
+    "mappable = ax.pcolormesh(ssp.ages.to_value(\"Gyr\"), ssp.metallicities.value,\n",
+    "                         ssp_weights.value, norm=LogNorm())\n",
+    "ax.plot((today - particles_t_form).to_value(u.Gyr), particles_z,\n",
+    "        'k,', alpha=.25)\n",
     "\n",
     "plt.colorbar(mappable, ax=ax, label=r\"SSP mass (M$_\\odot$)\")"
    ]
@@ -226,9 +228,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "test-pst",
+   "display_name": "base",
    "language": "python",
-   "name": "test-pst"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -240,7 +242,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,