Fast X: level populations edition

fiasco/base.py

@@ -4,8 +4,10 @@
 These classes are not meant to be instantiated directly by the user.
 """
 import astropy.units as u
+import pathlib
 import plasmapy.particles
 
+from functools import cache
 from packaging.version import Version
 from plasmapy.utils import roman
 
@@ -21,6 +23,21 @@
 __all__ = ['IonBase']
 
 
+@cache
+def _ion_names_in_dbase(hdf5_dbase_root, mtime_ns, size):
+    # Memoized set of available ion names per database file. The modification
+    # time and size of the file are part of the key such that the cache is
+    # invalidated if the database is rebuilt.
+    return frozenset(fiasco.list_ions(hdf5_dbase_root, sort=False))
+
+
+@cache
+def _dbase_fiasco_version(hdf5_dbase_root, mtime_ns, size):
+    # Memoized fiasco version of the database file (see _ion_names_in_dbase)
+    # to avoid opening the database on every ion instantiation.
+    return DataIndexer(hdf5_dbase_root, '/').fiasco_version
+
+
 class IonBase:
     """
     Base class for accessing data attached to a particular ion.
@@ -56,7 +73,9 @@ def __init__(self, ion_name, hdf5_dbase_root=None, **kwargs):
             # base rep is a tuple of integers (atomic_number, ionization_stage)
             self._base_rep = parse_ion_name(ion_name)
             check_database(self.hdf5_dbase_root, **kwargs)
-            if self.ion_name not in fiasco.list_ions(self.hdf5_dbase_root, sort=False):
+            path = pathlib.Path(self.hdf5_dbase_root).resolve()
+            stat = path.stat()
+            if self.ion_name not in _ion_names_in_dbase(str(path), stat.st_mtime_ns, stat.st_size):
                 raise MissingIonError(f'{self.ion_name} not found in {self.hdf5_dbase_root}')
             # Warn users if the database they are using is potentially stale.
             self._check_dbase_fiasco_version()
@@ -65,7 +84,9 @@ def __init__(self, ion_name, hdf5_dbase_root=None, **kwargs):
 
     def _check_dbase_fiasco_version(self):
         "Warn if database was generated with an earlier version of fiasco."
-        dbase_version = DataIndexer(self.hdf5_dbase_root, '/').fiasco_version
+        path = pathlib.Path(self.hdf5_dbase_root).resolve()
+        stat = path.stat()
+        dbase_version = _dbase_fiasco_version(str(path), stat.st_mtime_ns, stat.st_size)
         if dbase_version is not None:
             current_version = Version(fiasco.__version__)
             if dbase_version < current_version and not current_version.is_devrelease:

fiasco/elements.py

@@ -10,6 +10,7 @@
 import fiasco
 
 from fiasco.util import parse_ion_name
+from fiasco.util.util import _atomic_number
 
 __all__ = ['Element']
 
@@ -37,7 +38,7 @@ class Element(fiasco.IonCollection):
     def __init__(self, element_name, temperature: u.K, **kwargs):
         if isinstance(element_name, str):
             element_name = element_name.capitalize()
-        Z = plasmapy.particles.atomic_number(element_name)
+        Z = _atomic_number(element_name)
         ion_list = []
         if (pe_ratio := kwargs.pop('proton_electron_ratio', None)) is None:
             pe_ratio = fiasco.proton_electron_ratio(temperature, **kwargs)

fiasco/fiasco.py

@@ -3,7 +3,6 @@
 """
 import astropy.units as u
 import numpy as np
-import plasmapy.particles
 
 from plasmapy.particles.exceptions import InvalidParticleError
 from scipy.interpolate import interp1d
@@ -12,6 +11,7 @@
 
 from fiasco.io import DataIndexer
 from fiasco.util import parse_ion_name
+from fiasco.util.util import _atomic_number, _atomic_symbol
 
 __all__ = [
     'list_elements',
@@ -40,11 +40,11 @@ def list_elements(hdf5_dbase_root=None, sort=True):
     root = DataIndexer.create_indexer(hdf5_dbase_root, '/')
     for f in root.fields:
         try:
-            elements.append(plasmapy.particles.atomic_symbol(f.capitalize()))
+            elements.append(_atomic_symbol(f.capitalize()))
         except InvalidParticleError:
             continue
     if sort:
-        elements = sorted(elements, key=lambda x: plasmapy.particles.atomic_number(x))
+        elements = sorted(elements, key=lambda x: _atomic_number(x))
     return elements
 
 
@@ -73,15 +73,15 @@ def list_ions(hdf5_dbase_root=None, sort=True, base_rep=False):
         ions = []
         for f in root.fields:
             try:
-                el = plasmapy.particles.atomic_symbol(f.capitalize())
+                el = _atomic_symbol(f.capitalize())
                 for i in root[f].fields:
                     if f == i.split('_')[0]:
                         ions.append(f"{el} {i.split('_')[1]}")
             except InvalidParticleError:
                 continue
     # Optional because adds significant overhead
     if sort:
-        ions = sorted(ions, key=lambda x: (plasmapy.particles.atomic_number(x.split()[0]),
+        ions = sorted(ions, key=lambda x: (_atomic_number(x.split()[0]),
                                            int(x.split()[1])))
     # NOTE: when grabbing straight from the index and not sorting, the result will be
     # a numpy array. Cast to a list to make sure the return type is consistent for
@@ -90,7 +90,7 @@ def list_ions(hdf5_dbase_root=None, sort=True, base_rep=False):
     if base_rep:
         # NOTE: Explicitly not using parse_ion_name here as it can be slow.
         elements = list_elements(hdf5_dbase_root=hdf5_dbase_root, sort=False)
-        el_map = {el: plasmapy.particles.atomic_number(el) for el in elements}
+        el_map = {el: _atomic_number(el) for el in elements}
         ion_list = [(el_map[el], int(ion)) for el, ion in map(lambda x: x.split(), ion_list)]
     return ion_list
 
@@ -140,7 +140,7 @@ def get_isoelectronic_sequence(element, hdf5_dbase_root=None):
     hdf5_dbase_root: path-like, optional
         If not specified, will default to that specified in ``fiasco.defaults``.
     """
-    Z_iso = plasmapy.particles.atomic_number(element)
+    Z_iso = _atomic_number(element)
     all_ions = list_ions(hdf5_dbase_root=hdf5_dbase_root)
 
     def _is_in_sequence(ion):

fiasco/io/datalayer.py

@@ -114,12 +114,16 @@ def __contains__(self, key):
     def __getitem__(self, key):
         if isinstance(key, int):
             raise NotImplementedError('Iteration not supported.')
-        if key not in self:
-            raise KeyError(f'{key} not found in {self.top_level_path}')
+        # NOTE: The membership check, indexer creation, and read are all done with
+        # a single open of the HDF5 file as opening and closing the file on every
+        # access carries a non-negligible overhead.
         with h5py.File(self.hdf5_dbase_root, 'r') as hf:
-            ds = hf[self.top_level_path][key]
+            grp = hf[self.top_level_path]
+            if key not in grp:
+                raise KeyError(f'{key} not found in {self.top_level_path}')
+            ds = grp[key]
             if isinstance(ds, h5py.Group):
-                data = DataIndexer.create_indexer(
+                data = DataIndexerHDF5(
                     self.hdf5_dbase_root, '/'.join([self.top_level_path, key]))
             else:
                 if ds.attrs['unit'] == 'SKIP' or ds.dtype == 'object':

fiasco/ions.py

@@ -234,6 +234,8 @@ def proton_electron_ratio(self, value):
         if value is None:
             value = proton_electron_ratio(self.temperature, **self._instance_kwargs)
         self._proton_electron_ratio = u.Quantity(value*np.ones(self.temperature.shape))
+        self.__dict__.pop('_next_ion', None)
+        self.__dict__.pop('_previous_ion', None)
 
     def next_ion(self):
         """
@@ -257,6 +259,18 @@ def previous_ion(self):
                            self.temperature,
                            **self._instance_kwargs)
 
+    @cached_property
+    def _next_ion(self):
+        # Cached instance of the next ion for internal use such that expensive
+        # quantities computed on the next ion (e.g. rate matrices) are only
+        # computed once rather than on every call to next_ion().
+        return self.next_ion()
+
+    @cached_property
+    def _previous_ion(self):
+        # Cached instance of the previous ion for internal use (see _next_ion).
+        return self.previous_ion()
+
     @property
     @needs_dataset('elvlc', 'wgfa')
     def transitions(self):
@@ -292,6 +306,8 @@ def ionization_fraction(self, ionization_fraction):
             ionization_fraction = np.atleast_1d(ionization_fraction) * np.ones(self.temperature.shape)
             self._dset_names['ionization_fraction'] = None
             self._ionization_fraction = ionization_fraction
+        self.__dict__.pop('_next_ion', None)
+        self.__dict__.pop('_previous_ion', None)
 
     @staticmethod
     def _interpolate_ionization_fraction(temperature, temperature_data, ionization_data):
@@ -711,22 +727,19 @@ def _build_two_ion_coefficient_matrix(self, electron_density, include_protons=Fa
         c_matrix_recombined = self._build_coefficient_matrix(electron_density, include_protons=include_protons)
         # Get coefficient matrix of recombining ion
         try:
-            c_matrix_recombining = self.next_ion()._build_coefficient_matrix(electron_density,
-                                                                             include_protons=include_protons)
+            c_matrix_recombining = self._next_ion._build_coefficient_matrix(electron_density,
+                                                                            include_protons=include_protons)
         except MissingDatasetException:
             self.log.warning(
-                f'No rate data available for recombining ion {self.next_ion().ion_name}. '
+                f'No rate data available for recombining ion {self._next_ion.ion_name}. '
                 f'Using single-ion model for {self.ion_name}.'
             )
             return c_matrix_recombined
         rate_matrix_total = self._empty_rate_matrix(unit='s-1')
         # Add terms that include both ions
         d_e = electron_density[:, np.newaxis, np.newaxis]
         rate_matrix_total += self._rate_matrix_autoionization
-        rate_matrix_total += d_e * (self._rate_matrix_ionization
-                                    + self._rate_matrix_radiative_recombination
-                                    + self._rate_matrix_dielectronic_capture
-                                    + self._rate_matrix_dielectronic_recombination)
+        rate_matrix_total += d_e * self._rate_matrix_two_ion_density_dependent
         # Add depopulating terms
         # NOTE: By summing over the rows, we are computing the processes that depopulate
         # that level by summing up all of the processes that populate *from* that level.
@@ -782,12 +795,23 @@ def _rate_matrix_collisional_proton(self) -> u.Unit('cm3 s-1'):
         return rate_matrix
 
     def _empty_rate_matrix(self, temperature_dependent=True, unit='cm3 s-1'):
-        n_levels = self.n_levels + self.next_ion().n_levels
+        n_levels = self.n_levels + self._next_ion.n_levels
         shape = (n_levels, n_levels)
         if temperature_dependent:
             shape = self.temperature.shape + shape
         return u.Quantity(np.zeros(shape), unit)
 
+    @cached_property
+    @u.quantity_input
+    def _rate_matrix_two_ion_density_dependent(self) -> u.Unit('cm3 s-1'):
+        # Sum of the density-dependent two-ion rate matrices. This sum is
+        # independent of density so it is cached to avoid recomputing it for
+        # every density value in the level populations calculation.
+        return (self._rate_matrix_ionization
+                + self._rate_matrix_radiative_recombination
+                + self._rate_matrix_dielectronic_capture
+                + self._rate_matrix_dielectronic_recombination)
+
     @cached_property
     @u.quantity_input
     def _rate_matrix_ionization(self) -> u.Unit('cm3 s-1'):
@@ -802,7 +826,7 @@ def _rate_matrix_radiative_recombination(self) -> u.Unit('cm3 s-1'):
         rate_matrix = self._empty_rate_matrix()
         try:
             # NOTE: Using copy to avoid in-place modification of cached property
-            rr_rate_ground = self.next_ion().radiative_recombination_rate.copy()
+            rr_rate_ground = self._next_ion.radiative_recombination_rate.copy()
         except MissingDatasetException:
             rr_rate_ground = u.Quantity(np.zeros(self.temperature.shape), 'cm3 s-1')
         if self._has_dataset('rrlvl'):
@@ -849,7 +873,7 @@ def _rate_matrix_autoionization(self) -> u.Unit('s-1'):
 
     def _dielectronic_capture_rate(self, level_lower, level_upper, A_auto):
         # See Eq. A4 of Dere et al. (2019)
-        next_ion = self.next_ion()
+        next_ion = self._next_ion
         levels_recombined = self[vectorize_where(self.levels.level, level_upper)]
         levels_recombining = next_ion[vectorize_where(next_ion.levels.level, level_lower)]
         g_ratio = levels_recombined.weight / levels_recombining.weight
@@ -888,7 +912,7 @@ def _rate_matrix_dielectronic_recombination(self) -> u.Unit('cm3 s-1'):
         rate_matrix = self._empty_rate_matrix()
         # Compute ground-ground dielectronic recombination rate
         try:
-            dr_rate_ground = self.next_ion().dielectronic_recombination_rate
+            dr_rate_ground = self._next_ion.dielectronic_recombination_rate
         except MissingDatasetException:
             dr_rate_ground = u.Quantity(np.zeros(self.temperature.shape), 'cm3 s-1')
         # NOTE: Explicitly not using a decorator here in order to return an empty matrix
@@ -1081,14 +1105,14 @@ def _population_correction(self, population, density, rate_matrix):
         numerator = np.zeros(population.shape)
         try:
             upper_level_ionization, ionization_rate = self._level_resolved_ionization_rate
-            ionization_fraction_previous = self.previous_ion().ionization_fraction.value[:, np.newaxis]
+            ionization_fraction_previous = self._previous_ion.ionization_fraction.value[:, np.newaxis]
             upper_index_ionization = upper_level_ionization-1
             numerator[:, upper_index_ionization] += (ionization_rate * ionization_fraction_previous).to_value('cm3 s-1')
         except MissingDatasetException:
             pass
         try:
             upper_level_recombination, recombination_rate = self._level_resolved_recombination_rate
-            ionization_fraction_next = self.next_ion().ionization_fraction.value[:, np.newaxis]
+            ionization_fraction_next = self._next_ion.ionization_fraction.value[:, np.newaxis]
             upper_index_recombination = upper_level_recombination-1
             numerator[:, upper_index_recombination] += (recombination_rate * ionization_fraction_next).to_value('cm3 s-1')
         except MissingDatasetException:

fiasco/tests/test_base.py

@@ -4,6 +4,7 @@
 import h5py
 import logging
 import numpy as np
+import os
 import pytest
 
 from packaging.version import Version
@@ -73,3 +74,30 @@ def test_old_fiasco_version_warning(tmpdir_factory, caplog):
         assert caplog.text == ''
     else:
         assert 'was produced with an earlier version of fiasco' in caplog.text
+
+
+@pytest.mark.parametrize(('second_ion', 'second_version'), [
+    ('C 1', '0.9.0'),
+    ('H 1', '0.8.0'),
+])
+def test_database_cache_with_relative_paths(tmp_path, monkeypatch, caplog, second_ion, second_version):
+    monkeypatch.setattr(fiasco, '__version__', '0.9.0')
+    paths = []
+    for directory, ion_name, version in [('first', 'H 1', '0.9.0'), ('second', second_ion, second_version)]:
+        path = tmp_path / directory / 'dbase.h5'
+        path.parent.mkdir()
+        with h5py.File(path, 'w') as hf:
+            hf.attrs['fiasco_version'] = version
+            dataset = hf.create_dataset('ion_index', data=np.array([ion_name], dtype='S3'))
+            dataset.attrs['unit'] = 'SKIP'
+        os.utime(path, ns=(1_000_000_000, 1_000_000_000))
+        paths.append(path)
+    assert paths[0].stat().st_size == paths[1].stat().st_size
+    monkeypatch.chdir(paths[0].parent)
+    fiasco.base.IonBase('H 1', hdf5_dbase_root='dbase.h5')
+    monkeypatch.chdir(paths[1].parent)
+    caplog.clear()
+    with caplog.at_level(logging.WARN):
+        ion = fiasco.base.IonBase(second_ion, hdf5_dbase_root='dbase.h5')
+    assert ion.ion_name == second_ion
+    assert ('was produced with an earlier version' in caplog.text) == (second_version == '0.8.0')

fiasco/tests/test_ion.py

@@ -625,6 +625,30 @@ def test_proton_electron_ratio_setter(ion, value):
     assert u.allclose(new_ion.proton_electron_ratio, value)
 
 
+@pytest.mark.requires_dbase_version('>= 9')
+def test_two_ion_matrix_after_changing_proton_electron_ratio(c5):
+    ion = c5._new_instance(temperature=c5.temperature[:2], proton_electron_ratio=0.83)
+    density = [1e9] * u.cm**-3
+    before = ion._build_two_ion_coefficient_matrix(density, include_protons=True)
+    ion.proton_electron_ratio = 0.0
+    after = ion._build_two_ion_coefficient_matrix(density, include_protons=True)
+    expected = ion._new_instance()._build_two_ion_coefficient_matrix(density, include_protons=True)
+    np.testing.assert_array_equal(after, expected)
+    assert np.any(before != after)
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_level_populations_after_changing_ionization_fraction(fe20):
+    ion = fe20._new_instance(temperature=[1e7, 1.5e7] * u.K)
+    density = [1e9] * u.cm**-3
+    before = ion.level_populations(density, use_two_ion_model=False)
+    ion.ionization_fraction = 0.123 * u.dimensionless_unscaled
+    after = ion.level_populations(density, use_two_ion_model=False)
+    expected = ion._new_instance().level_populations(density, use_two_ion_model=False)
+    np.testing.assert_array_equal(after, expected)
+    assert np.any(before != after)
+
+
 @pytest.mark.requires_dbase_version('>= 8')
 def test_emissivity_uses_proton_electron_ratio(ion):
     # Setting the ratio to 0 should zero the emissivity, which is only the case

fiasco/util/tests/test_util.py

@@ -1,6 +1,7 @@
 """
 Tests for util functions
 """
+import numpy as np
 import packaging.version
 import pytest
 
@@ -41,3 +42,9 @@ def test_get_chianti_catalog(ascii_dbase_root):
 def test_chianti_version(ascii_dbase_root):
     version = read_chianti_version(ascii_dbase_root)
     assert isinstance(version, packaging.version.Version)
+
+
+def test_parse_ion_name_cache_preserves_input_type_validation():
+    assert parse_ion_name((np.int64(26), 21)) == (26, 21)
+    with pytest.raises(TypeError):
+        parse_ion_name((26.0, 21))