From 66759af40cdbd749af63f51223624f7b8f5a7436 Mon Sep 17 00:00:00 2001
From: Will Barnes <will.t.barnes@gmail.com>
Date: Fri, 2 Feb 2024 17:34:59 +0000
Subject: [PATCH] Backport PR #262: Use total recombination data if available
 to calculate recombination rate

---
 fiasco/ions.py           | 36 ++++++++++++++++++++++++++++++++++++
 fiasco/tests/test_ion.py | 28 +++++++++++++++++++---------
 2 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/fiasco/ions.py b/fiasco/ions.py
index df9c9360..a6387e4f 100644
--- a/fiasco/ions.py
+++ b/fiasco/ions.py
@@ -1159,6 +1159,17 @@ def dielectronic_recombination_rate(self) -> u.cm**3 / u.s:
         else:
             raise ValueError(f"Unrecognized fit type {self._drparams['fit_type']}")
 
+    @cached_property
+    @needs_dataset('trparams')
+    @u.quantity_input
+    def _total_recombination_rate(self) -> u.cm**3 / u.s:
+        temperature_data = self._trparams['temperature'].to_value('K')
+        rate_data = self._trparams['recombination_rate'].to_value('cm3 s-1')
+        f_interp = interp1d(temperature_data, rate_data, fill_value='extrapolate', kind='cubic')
+        f_interp = PchipInterpolator(np.log10(temperature_data), np.log10(rate_data), extrapolate=True)
+        rate_interp = 10**f_interp(np.log10(self.temperature.to_value('K')))
+        return u.Quantity(rate_interp, 'cm3 s-1')
+
     @cached_property
     @u.quantity_input
     def recombination_rate(self) -> u.cm**3 / u.s:
@@ -1172,18 +1183,43 @@ def recombination_rate(self) -> u.cm**3 / u.s:
 
             \alpha_{R} = \alpha_{RR} + \alpha_{DR}
 
+        .. warning::
+
+            For most ions, this total recombination rate is computed by summing the
+            outputs of the `radiative_recombination_rate` and `dielectronic_recombination_rate` methods.
+            However, for some ions, total recombination rate data is available in the
+            so-called ``.trparams`` files. For these ions, the output of this method
+            will *not* be equal to the sum of the `dielectronic_recombination_rate` and
+            `radiative_recombination_rate` method. As such, when computing the total
+            recombination rate, this method should always be used.
+
         See Also
         --------
         radiative_recombination_rate
         dielectronic_recombination_rate
         """
+        # NOTE: If the trparams data is available, then it is prioritized over the sum
+        # of the dielectronic and radiative recombination rates. This is also how the
+        # total recombination rates are computed in IDL. The reasoning here is that the
+        # total recombination rate data, if available, is more reliable than the sum of
+        # the radiative and dielectronic recombination rates. According to P. Young, there
+        # is some slight controversy over this within some communities, but CHIANTI has chosen
+        # to prioritize this data if it exists.
+        try:
+            tr_rate = self._total_recombination_rate
+        except MissingDatasetException:
+            self.log.debug(f'No total recombination data available for {self.ion_name}.')
+        else:
+            return tr_rate
         try:
             rr_rate = self.radiative_recombination_rate
         except MissingDatasetException:
+            self.log.debug(f'No radiative recombination data available for {self.ion_name}.')
             rr_rate = u.Quantity(np.zeros(self.temperature.shape), 'cm3 s-1')
         try:
             dr_rate = self.dielectronic_recombination_rate
         except MissingDatasetException:
+            self.log.debug(f'No dielectronic recombination data available for {self.ion_name}.')
             dr_rate = u.Quantity(np.zeros(self.temperature.shape), 'cm3 s-1')
         return rr_rate + dr_rate
 
diff --git a/fiasco/tests/test_ion.py b/fiasco/tests/test_ion.py
index 37fe4450..ce3987d6 100644
--- a/fiasco/tests/test_ion.py
+++ b/fiasco/tests/test_ion.py
@@ -287,18 +287,28 @@ def test_intensity(ion, em):
     assert intens.shape == ion.temperature.shape + (1, ) + wave_shape
 
 
-def test_excitation_autoionization_rate(ion):
-    rate = ion.excitation_autoionization_rate
+@pytest.mark.parametrize(('rate_name','answer'), [
+    # NOTE: The expected values have not been tested for correctness and
+    # are only meant to indicate whether a value has changed or not.
+    ('direct_ionization_rate', 9.448935172152884e-13*u.cm**3 / u.s),
+    ('excitation_autoionization_rate', 1.14821255e-12 * u.cm**3 / u.s),
+    ('dielectronic_recombination_rate', 1.60593802e-11 * u.cm**3 / u.s),
+    ('radiative_recombination_rate', 1.6221634159408823e-12*u.cm**3 / u.s),
+])
+def test_rates(ion, rate_name, answer):
+    rate = getattr(ion, rate_name)
     assert rate.shape == ion.temperature.shape
-    # This value has not been tested for correctness
-    assert u.allclose(rate[0], 1.14821255e-12 * u.cm**3 / u.s)
+    assert u.allclose(rate[0], answer)
 
 
-def test_dielectronic_recombination_rate(ion):
-    rate = ion.dielectronic_recombination_rate
-    assert rate.shape == ion.temperature.shape
-    # This value has not been tested for correctness
-    assert u.allclose(rate[0], 1.60593802e-11 * u.cm**3 / u.s)
+def test_total_recombination_rate_priority(ion):
+    # This tests that in cases where the total recombination data in the trparams
+    # files is available that it is prioritized over the sum of the DR and RR rates.
+    # This is the case for this ion because Fe V has total recombination rate data
+    # available.
+    recomb_rate = ion.recombination_rate
+    tot_recomb_rate = ion._total_recombination_rate
+    assert np.all(recomb_rate == tot_recomb_rate)
 
 
 def test_free_free(ion):