Correct level populations for ionization and recombination processes

fiasco/conftest.py

@@ -84,6 +84,11 @@
     'fe_27.rrparams':                       '75383b0f1b167f862cfd26bbadd2a029',
     'fe_10.psplups':                        'dd34363f6daa81dbf106fbeb211b457d',
     'fe_10.elvlc':                          'f221d4c7167336556d57378ac368afc1',
+    'fe_20.elvlc':                          'bbddcf958dd41311ea24bf177c2b62de',
+    'fe_20.wgfa':                           'c991c30b98b03c9152ba5a2c71877149',
+    'fe_20.scups':                          'f0e375cad2ec8296efb2abcb8f02705e',
+    'fe_20.cilvl':                          'b71833c51a03c7073f1657ce60afcdbb',
+    'fe_20.reclvl':                         'cf28869709acef521fb6a1c9a2b59530',
 }
 
 

fiasco/ions.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 from functools import cached_property
-from scipy.interpolate import interp1d, splev, splrep
+from scipy.interpolate import interp1d, PchipInterpolator, splev, splrep
 from scipy.ndimage import map_coordinates
 
 from fiasco import proton_electron_ratio
@@ -173,18 +173,34 @@ def ioneq(self):
         ionization equilibrium outside of this temperature range, it is better to use the ionization
         and recombination rates.
 
+        Note
+        ----
+        The cubic interpolation is performed in log-log spaceusing a Piecewise Cubic Hermite
+        Interpolating Polynomial with `~scipy.interpolate.PchipInterpolator`. This helps to
+        ensure smoothness while reducing oscillations in the interpolated ionization fractions.
+
         See Also
         --------
         fiasco.Element.equilibrium_ionization
         """
-        f = interp1d(self._ioneq[self._dset_names['ioneq_filename']]['temperature'].to('MK').value,
-                     self._ioneq[self._dset_names['ioneq_filename']]['ionization_fraction'],
-                     kind='linear',
-                     bounds_error=False,
-                     fill_value=np.nan)
-        ioneq = f(self.temperature.to('MK').value)
-        isfinite = np.isfinite(ioneq)
-        ioneq[isfinite] = np.where(ioneq[isfinite] < 0., 0., ioneq[isfinite])
+        temperature = self.temperature.to_value('K')
+        temperature_data = self._ioneq[self._dset_names['ioneq_filename']]['temperature'].to_value('K')
+        ioneq_data = self._ioneq[self._dset_names['ioneq_filename']]['ionization_fraction'].value
+        # Perform PCHIP interpolation in log-space on only the non-zero ionization fractions.
+        # See https://github.com/wtbarnes/fiasco/pull/223 for additional discussion.
+        is_nonzero = ioneq_data > 0.0
+        f_interp = PchipInterpolator(np.log10(temperature_data[is_nonzero]),
+                                     np.log10(ioneq_data[is_nonzero]),
+                                     extrapolate=False)
+        ioneq = f_interp(np.log10(temperature))
+        ioneq = 10**ioneq
+        # This sets all entries that would have interpolated to zero ionization fraction to zero
+        ioneq = np.where(np.isnan(ioneq), 0.0, ioneq)
+        # Set entries that are truly out of bounds of the original temperature data back to NaN
+        out_of_bounds = np.logical_or(temperature<temperature_data.min(), temperature>temperature_data.max())
+        ioneq = np.where(out_of_bounds, np.nan, ioneq)
+        is_finite = np.isfinite(ioneq)
+        ioneq[is_finite] = np.where(ioneq[is_finite] < 0., 0., ioneq[is_finite])
         return u.Quantity(ioneq)
 
     @property
@@ -339,6 +355,7 @@ def proton_collision_excitation_rate(self) -> u.cm**3 / u.s:
 
         See Also
         --------
+        proton_collision_deexcitation_rate
         electron_collision_excitation_rate
         """
         # Create scaled temperature--these are not stored in the file
@@ -389,6 +406,7 @@ def proton_collision_deexcitation_rate(self) -> u.cm**3 / u.s:
     def level_populations(self,
                           density: u.cm**(-3),
                           include_protons=True,
+                          include_level_resolved_rate_correction=True,
                           couple_density_to_temperature=False) -> u.dimensionless_unscaled:
         """
         Energy level populations as a function of temperature and density.
@@ -507,10 +525,136 @@ def level_populations(self,
             # positivity
             np.fabs(pop, out=pop)
             np.divide(pop, pop.sum(axis=1)[:, np.newaxis], out=pop)
+            # Apply ionization/recombination correction
+            if include_level_resolved_rate_correction:
+                correction = self._population_correction(pop, d, c_matrix)
+                pop *= correction
+                np.divide(pop, pop.sum(axis=1)[:, np.newaxis], out=pop)
             populations[:, i, :] = pop
 
         return u.Quantity(populations)
 
+    def _level_resolved_rates_interpolation(self, temperature_table, rate_table,
+                                            extrapolate_above=False,
+                                            extrapolate_below=False):
+        # NOTE: According to CHIANTI Technical Report No. 20, Section 5,
+        # the interpolation of the level resolved recombination,
+        # the rates should be zero below the temperature range and above
+        # the temperature range, the last two points should be used to perform
+        # a linear extrapolation. For the ionization rates, the rates should be
+        # zero above the temperature range and below the temperature range, the
+        # last two points should be used. Thus, we need to perform two interpolations
+        # for each level.
+        # NOTE: In the CHIANTI IDL code, the interpolation is done using a cubic spline.
+        # Here, the rates are interpolated using a Piecewise Cubic Hermite Interpolating
+        # Polynomial (PCHIP) which balances smoothness and also reduces the oscillations
+        # that occur with higher order spline fits. This is needed mostly due to the wide
+        # range over which this data is fit.
+        temperature = self.temperature.to_value('K')
+        rates = []
+        for t, r in zip(temperature_table.to_value('K'), rate_table.to_value('cm3 s-1')):
+            rate_interp = PchipInterpolator(t, r, extrapolate=False)(temperature)
+            # NOTE: Anything outside of the temperature range will be set to NaN by the
+            # interpolation but we want these to be 0.
+            rate_interp = np.where(np.isnan(rate_interp), 0, rate_interp)
+            if extrapolate_above:
+                f_extrapolate = interp1d(t[-2:], r[-2:], kind='linear', fill_value='extrapolate')
+                i_extrapolate = np.where(temperature > t[-1])
+                rate_interp[i_extrapolate] = f_extrapolate(temperature[i_extrapolate])
+            if extrapolate_below:
+                f_extrapolate = interp1d(t[:2], r[:2], kind='linear', fill_value='extrapolate')
+                i_extrapolate = np.where(temperature < t[0])
+                rate_interp[i_extrapolate] = f_extrapolate(temperature[i_extrapolate])
+            rates.append(rate_interp)
+        # NOTE: Take transpose to maintain consistent ordering of temperature in the leading
+        # dimension and levels in the last dimension
+        rates =  u.Quantity(rates, 'cm3 s-1').T
+        # NOTE: The linear extrapolation at either end may return rates < 0 so we set these
+        # to zero.
+        rates = np.where(rates<0, 0, rates)
+        return rates
+
+    @cached_property
+    @needs_dataset('cilvl')
+    @u.quantity_input
+    def _level_resolved_ionization_rate(self):
+        ionization_rates = self._level_resolved_rates_interpolation(
+            self._cilvl['temperature'],
+            self._cilvl['ionization_rate'],
+            extrapolate_below=True,
+            extrapolate_above=False,
+        )
+        return self._cilvl['upper_level'], ionization_rates
+
+    @cached_property
+    @needs_dataset('reclvl')
+    @u.quantity_input
+    def _level_resolved_recombination_rate(self):
+        recombination_rates = self._level_resolved_rates_interpolation(
+            self._reclvl['temperature'],
+            self._reclvl['recombination_rate'],
+            extrapolate_below=False,
+            extrapolate_above=True,
+        )
+        return self._reclvl['upper_level'], recombination_rates
+
+    @u.quantity_input
+    def _population_correction(self, population, density, rate_matrix):
+        """
+        Correct level population to account for ionization and
+        recombination processes.
+
+        Parameters
+        ----------
+        population: `np.ndarray`
+        density: `~astropy.units.Quantity`
+        rate_matrix: `~astropy.units.Quantity`
+
+        Returns
+        -------
+        correction: `np.ndarray`
+            Correction factor to multiply populations by
+        """
+        # NOTE: These are done in separate try/except blocks because some ions have just a cilvl file,
+        # some have just a reclvl file, and some have both.
+        # NOTE: Ioneq values for surrounding ions are retrieved afterwards because first and last ions do
+        # not have previous or next ions but also do not have reclvl or cilvl files.
+        # NOTE: stripping the units off and adding them at the end because of some strange astropy
+        # Quantity behavior that does not allow for adding these two compatible shapes together.
+        numerator = np.zeros(population.shape)
+        try:
+            upper_level_ionization, ionization_rate = self._level_resolved_ionization_rate
+            ioneq_previous = self.previous_ion().ioneq.value[:, np.newaxis]
+            numerator[:, upper_level_ionization-1] += (ionization_rate * ioneq_previous).to_value('cm3 s-1')
+        except MissingDatasetException:
+            pass
+        try:
+            upper_level_recombination, recombination_rate = self._level_resolved_recombination_rate
+            ioneq_next = self.next_ion().ioneq.value[:, np.newaxis]
+            numerator[:, upper_level_recombination-1] += (recombination_rate * ioneq_next).to_value('cm3 s-1')
+        except MissingDatasetException:
+            pass
+        numerator *= density.to_value('cm-3')
+
+        c = rate_matrix.to_value('s-1').copy()
+        # This excludes processes that depopulate the level
+        i_diag, j_diag = np.diag_indices(c.shape[1])
+        c[:, i_diag, j_diag] = 0.0
+        # Sum of the population-weighted excitations from lower levels
+        # and cascades from higher levels
+        denominator = np.einsum('ijk,ik->ij', c, population)
+        denominator *= self.ioneq.value[:, np.newaxis]
+        # Set any zero entries to NaN to avoid divide by zero warnings
+        denominator = np.where(denominator==0.0, np.nan, denominator)
+
+        ratio = numerator / denominator
+        # Set ratio to zero where denominator is zero. This also covers the
+        # case of out-of-bounds ionization fractions (which will be NaN)
+        ratio = np.where(np.isfinite(ratio), ratio, 0.0)
+        # NOTE: Correction should not affect the ground state populations
+        ratio[:, 0] = 0.0
+        return 1.0 + ratio
+
     @needs_dataset('abundance', 'elvlc')
     @u.quantity_input
     def contribution_function(self, density: u.cm**(-3), **kwargs) -> u.cm**3 * u.erg / u.s:

fiasco/tests/idl/test_idl_ioneq.py

@@ -32,6 +32,7 @@ def ioneq_from_idl(idl_env, ascii_dbase_root):
     'C 2',
     'C 3',
     'Ca 2',
+    'Fe 20',
 ])
 def test_ioneq_from_idl(ion_name, ioneq_from_idl, hdf5_dbase_root):
     temperature = 10**ioneq_from_idl['ioneq_logt'] * u.K

fiasco/tests/test_collections.py

@@ -93,16 +93,18 @@ def test_length(collection):
 def test_free_free(another_collection, wavelength):
     ff = another_collection.free_free(wavelength)
     assert ff.shape == temperature.shape + wavelength.shape if wavelength.shape else (1,)
-    index = 50 if wavelength.shape else 0
-    assert u.allclose(ff[50, index], 3.19877384e-35 * u.Unit('erg cm3 s-1 Angstrom-1'))
+    index_w = 50 if wavelength.shape else 0
+    index_t = 24  # This is approximately where the ioneq for Fe V peaks
+    assert u.allclose(ff[index_t, index_w], 3.2914969734961024e-42 * u.Unit('erg cm3 s-1 Angstrom-1'))
 
 
 @pytest.mark.parametrize('wavelength', [wavelength, wavelength[50]])
 def test_free_bound(another_collection, wavelength):
     fb = another_collection.free_bound(wavelength)
     assert fb.shape == temperature.shape + wavelength.shape if wavelength.shape else (1,)
-    index = 50 if wavelength.shape else 0
-    assert u.allclose(fb[50, index], 3.2653516e-29 * u.Unit('erg cm3 s-1 Angstrom-1'))
+    index_w = 50 if wavelength.shape else 0
+    index_t = 24  # This is approximately where the ioneq for Fe V peaks
+    assert u.allclose(fb[index_t, index_w], 1.1573022245197259e-35 * u.Unit('erg cm3 s-1 Angstrom-1'))
 
 
 def test_radiative_los(collection):

fiasco/tests/test_ion.py

@@ -32,6 +32,13 @@ def c6(hdf5_dbase_root):
     return fiasco.Ion('C VI', temperature, hdf5_dbase_root=hdf5_dbase_root)
 
 
+@pytest.fixture
+def fe20(hdf5_dbase_root):
+    # NOTE: This ion was added because it has reclvl and cilvl files which
+    # we need to test the level-resolved rate correction factor
+    return fiasco.Ion('Fe XX', temperature, hdf5_dbase_root=hdf5_dbase_root)
+
+
 def test_new_instance(ion):
     abundance_filename = ion._instance_kwargs['abundance_filename']
     new_ion = ion._new_instance()
@@ -99,15 +106,7 @@ def test_scalar_temperature(hdf5_dbase_root):
     t_data = ion._ioneq[ion._dset_names['ioneq_filename']]['temperature']
     ioneq_data = ion._ioneq[ion._dset_names['ioneq_filename']]['ionization_fraction']
     i_t = np.where(t_data == ion.temperature)
-    np.testing.assert_allclose(ioneq, ioneq_data[i_t])
-
-
-def test_scalar_density(hdf5_dbase_root):
-    ion = fiasco.Ion('H 1', temperature, hdf5_dbase_root=hdf5_dbase_root)
-    pop = ion.level_populations(1e8 * u.cm**-3)
-    assert pop.shape == ion.temperature.shape + (1,) + ion._elvlc['level'].shape
-    # This value has not been checked for correctness
-    np.testing.assert_allclose(pop[0, 0, 0], 0.9965048292729177)
+    assert u.allclose(ioneq, ioneq_data[i_t])
 
 
 def test_no_elvlc_raises_index_error(hdf5_dbase_root):
@@ -116,13 +115,21 @@ def test_no_elvlc_raises_index_error(hdf5_dbase_root):
 
 
 def test_ioneq(ion):
-    assert ion.ioneq.shape == temperature.shape
     t_data = ion._ioneq[ion._dset_names['ioneq_filename']]['temperature']
     ioneq_data = ion._ioneq[ion._dset_names['ioneq_filename']]['ionization_fraction']
-    i_t = np.where(t_data == ion.temperature[0])
-    # Essentially test that we've done the interpolation to the data correctly
-    # for a single value
-    np.testing.assert_allclose(ion.ioneq[0], ioneq_data[i_t])
+    ion_at_nodes = ion._new_instance(temperature=t_data)
+    assert u.allclose(ion_at_nodes.ioneq, ioneq_data, rtol=1e-6)
+
+
+def test_ioneq_positive(ion):
+    assert np.all(ion.ioneq >= 0)
+
+
+def test_ioneq_out_bounds_is_nan(ion):
+    t_data = ion._ioneq[ion._dset_names['ioneq_filename']]['temperature']
+    t_out_of_bounds = t_data[[0,-1]] + [-100, 1e6] * u.K
+    ion_out_of_bounds = ion._new_instance(temperature=t_out_of_bounds)
+    assert np.isnan(ion_out_of_bounds.ioneq).all()
 
 
 def test_formation_temeprature(ion):
@@ -132,7 +139,7 @@ def test_formation_temeprature(ion):
 def test_abundance(ion):
     assert ion.abundance.dtype == np.dtype('float64')
     # This value has not been tested for correctness
-    np.testing.assert_allclose(ion.abundance, 0.0001258925411794166)
+    assert u.allclose(ion.abundance, 0.0001258925411794166)
 
 
 def test_proton_collision(fe10):
@@ -164,6 +171,15 @@ def test_missing_ip(hdf5_dbase_root):
         _ = ion.ip
 
 
+def test_level_populations(ion):
+    pop = ion.level_populations(1e8 * u.cm**-3)
+    assert pop.shape == ion.temperature.shape + (1,) + ion._elvlc['level'].shape
+    # This value has not been checked for correctness
+    assert u.allclose(pop[0, 0, 0], 0.011643747849652244)
+    # Check that the total populations are normalized to 1 for all temperatures
+    assert u.allclose(pop.squeeze().sum(axis=1), 1, atol=None, rtol=1e-15)
+
+
 def test_contribution_function(ion):
     cont_func = ion.contribution_function(1e7 * u.cm**-3)
     assert cont_func.shape == ion.temperature.shape + (1, ) + ion._wgfa['wavelength'].shape
@@ -204,6 +220,39 @@ def test_coupling_unequal_dimensions_exception(ion):
         _ = ion.level_populations([1e7, 1e8]*u.cm**(-3), couple_density_to_temperature=True)
 
 
+@pytest.fixture
+def pops_with_correction(fe20):
+    return fe20.level_populations(1e9*u.cm**(-3)).squeeze()
+
+
+@pytest.fixture
+def pops_no_correction(fe20):
+    return fe20.level_populations(1e9*u.cm**(-3),
+                                  include_level_resolved_rate_correction=False).squeeze()
+
+
+def test_level_populations_normalized(pops_no_correction, pops_with_correction):
+    assert u.allclose(pops_with_correction.sum(axis=1), 1, atol=None, rtol=1e-15)
+    assert u.allclose(pops_no_correction.sum(axis=1), 1, atol=None, rtol=1e-15)
+
+
+def test_level_populations_correction(fe20, pops_no_correction, pops_with_correction):
+    # Test level-resolved correction applied to correct levels
+    i_corrected = np.unique(np.concatenate([fe20._cilvl['upper_level'], fe20._reclvl['upper_level']]))
+    i_corrected -= 1
+    # This tests that, for at least some portion of the temperature axis, the populations are
+    # significantly different for each corrected level
+    pops_equal = u.isclose(pops_with_correction[:, i_corrected], pops_no_correction[:, i_corrected],
+                           atol=0.0, rtol=1e-5)
+    assert ~np.all(np.all(pops_equal, axis=0))
+    # All other levels should be unchanged (with some tolerance for renormalization)
+    is_uncorrected = np.ones(pops_no_correction.shape[-1], dtype=bool)
+    is_uncorrected[i_corrected] = False
+    i_uncorrected = np.where(is_uncorrected)
+    assert u.allclose(pops_with_correction[:, i_uncorrected], pops_no_correction[:, i_uncorrected],
+                      atol=0.0, rtol=1e-5)
+
+
 def test_emissivity(ion):
     emm = ion.emissivity(1e7 * u.cm**-3)
     assert emm.shape == ion.temperature.shape + (1, ) + ion._wgfa['wavelength'].shape