diff --git a/fiasco/base.py b/fiasco/base.py index fa64647a..9bb02505 100644 --- a/fiasco/base.py +++ b/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: diff --git a/fiasco/elements.py b/fiasco/elements.py index 7917e6aa..12a5fbf8 100644 --- a/fiasco/elements.py +++ b/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) diff --git a/fiasco/fiasco.py b/fiasco/fiasco.py index 1a50ca2b..ec14a3ab 100644 --- a/fiasco/fiasco.py +++ b/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,7 +73,7 @@ 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]}") @@ -81,7 +81,7 @@ def list_ions(hdf5_dbase_root=None, sort=True, base_rep=False): 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): diff --git a/fiasco/io/datalayer.py b/fiasco/io/datalayer.py index 9d05de86..01d567ff 100644 --- a/fiasco/io/datalayer.py +++ b/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': diff --git a/fiasco/ions.py b/fiasco/ions.py index 75fcd440..f6d4b8e6 100644 --- a/fiasco/ions.py +++ b/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,11 +727,11 @@ 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 @@ -723,10 +739,7 @@ def _build_two_ion_coefficient_matrix(self, electron_density, include_protons=Fa # 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: diff --git a/fiasco/tests/test_base.py b/fiasco/tests/test_base.py index 70ea34da..bd922498 100644 --- a/fiasco/tests/test_base.py +++ b/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') diff --git a/fiasco/tests/test_ion.py b/fiasco/tests/test_ion.py index 0ecf1629..8232632e 100644 --- a/fiasco/tests/test_ion.py +++ b/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 diff --git a/fiasco/util/tests/test_util.py b/fiasco/util/tests/test_util.py index 5ffbd01c..54949d1a 100644 --- a/fiasco/util/tests/test_util.py +++ b/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)) diff --git a/fiasco/util/util.py b/fiasco/util/util.py index 64868ea7..ad44ac1f 100644 --- a/fiasco/util/util.py +++ b/fiasco/util/util.py @@ -7,6 +7,7 @@ import plasmapy.particles import sys +from functools import cache, lru_cache from packaging.version import Version from plasmapy.utils import roman @@ -16,6 +17,19 @@ __all__ = ['setup_paths', 'get_chianti_catalog', 'read_chianti_version', 'parse_ion_name'] +@lru_cache(maxsize=None, typed=True) +def _atomic_number(element): + # Memoized version of plasmapy.particles.atomic_number as constructing + # a plasmapy Particle on every call is comparatively expensive. + return plasmapy.particles.atomic_number(element) + + +@cache +def _atomic_symbol(atomic_number): + # Memoized version of plasmapy.particles.atomic_symbol (see _atomic_number). + return plasmapy.particles.atomic_symbol(atomic_number) + + def parse_ion_name(ion_name): """ Parse the atomic number and ionization stage from representation of ion. @@ -40,7 +54,7 @@ def parse_ion_name(ion_name): # Parse element string if isinstance(element, str): element = element.capitalize() - element = plasmapy.particles.atomic_number(element) + element = _atomic_number(element) # Parse ion string if isinstance(ion, str): if '+' in ion: