From 31105d9549941145daca53955864ffc42c1d7777 Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Mon, 10 Aug 2026 16:35:16 -0700 Subject: [PATCH 1/4] Speed up level populations calculation Avoid repeated expensive recomputation in the level populations calculation while keeping the output bit-identical: - Cache the Ion instances returned by next_ion()/previous_ion() for internal use so that rate matrices of the recombining ion are only computed once rather than on every density iteration. - Cache the density-independent sum of the two-ion rate matrices. - Memoize plasmapy atomic_number/atomic_symbol/element_name lookups, which construct an expensive Particle object on every call. - Share the proton-electron ratio with the next ion since it only depends on temperature and the abundance/ionization datasets. Benchmark (21 temperatures x 9 densities): Fe XXIII 73.5->7.2 s, Ca XVIII 74.3->8.1 s, O VI 73.3->9.3 s, Fe XII (single-ion model) 9.9->5.0 s, coupled n-T 14.5->3.0 s. --- fiasco/ions.py | 46 ++++++++++++++++++++++++++++++++------------- fiasco/util/util.py | 10 +++++++++- 2 files changed, 42 insertions(+), 14 deletions(-) diff --git a/fiasco/ions.py b/fiasco/ions.py index 75fcd440..76b189fd 100644 --- a/fiasco/ions.py +++ b/fiasco/ions.py @@ -257,6 +257,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): @@ -711,11 +723,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 +735,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 +791,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 +822,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 +869,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 +908,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 +1101,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/util/util.py b/fiasco/util/util.py index 64868ea7..6e5dd64b 100644 --- a/fiasco/util/util.py +++ b/fiasco/util/util.py @@ -7,6 +7,7 @@ import plasmapy.particles import sys +from functools import lru_cache from packaging.version import Version from plasmapy.utils import roman @@ -16,6 +17,13 @@ __all__ = ['setup_paths', 'get_chianti_catalog', 'read_chianti_version', 'parse_ion_name'] +@lru_cache(maxsize=None) +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) + + def parse_ion_name(ion_name): """ Parse the atomic number and ionization stage from representation of ion. @@ -40,7 +48,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: From 38cf43fa0e7ea345e97c3bbf865badb3ba947e32 Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Mon, 10 Aug 2026 16:43:04 -0700 Subject: [PATCH 2/4] Reduce overhead of HDF5 access and Ion instantiation Further speed up of the level populations calculation (and anything else that creates many Ion instances or reads many datasets), again keeping the output bit-identical: - DataIndexerHDF5.__getitem__ now does the membership check, child indexer creation, and read with a single open of the HDF5 file instead of two to three opens per access. - Memoize the set of available ion names and the database fiasco version per database file, keyed on the file modification time and size so that rebuilding the database invalidates the cache. This avoids re-reading the full ion index on every Ion instantiation. Ion instantiation drops from ~3.7 ms to ~1.9 ms and fiasco.proton_electron_ratio from ~2.1 s to ~1.5 s. O VI level populations (21 T x 9 n): 9.3 -> 8.1 s, dominated by np.linalg.solve. --- fiasco/base.py | 23 +++++++++++++++++++++-- fiasco/io/datalayer.py | 12 ++++++++---- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/fiasco/base.py b/fiasco/base.py index fa64647a..325bde5f 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 lru_cache from packaging.version import Version from plasmapy.utils import roman @@ -21,6 +23,21 @@ __all__ = ['IonBase'] +@lru_cache(maxsize=None) +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)) + + +@lru_cache(maxsize=None) +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,8 @@ 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): + stat = pathlib.Path(self.hdf5_dbase_root).stat() + if self.ion_name not in _ion_names_in_dbase(str(self.hdf5_dbase_root), 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 +83,8 @@ 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 + stat = pathlib.Path(self.hdf5_dbase_root).stat() + dbase_version = _dbase_fiasco_version(str(self.hdf5_dbase_root), 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/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': From 3ea64b1d47e30dd39083161d0768e425b07a866b Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Mon, 10 Aug 2026 16:50:37 -0700 Subject: [PATCH 3/4] Speed up proton_electron_ratio by reading ionization data directly Rather than constructing an Element instance (and thus an Ion instance for every ionization stage) of every element just to read the tabulated ionization fraction data, read the ioneq data directly from the database and instantiate only the first ion of each element to resolve the abundance exactly as before. The iteration order, arithmetic, warning messages, and missing-dataset behavior are unchanged such that the result is bit-identical, including on databases with missing ionization fraction datasets. Also use the memoized atomic number/symbol lookups in list_elements, list_ions, get_isoelectronic_sequence, and Element. fiasco.proton_electron_ratio drops from ~1.4 s to ~0.4 s. O VI level populations (21 T x 9 n): 8.1 -> 6.9 s. The full test suite also runs roughly twice as fast. --- fiasco/elements.py | 3 ++- fiasco/fiasco.py | 14 +++++++------- fiasco/util/util.py | 6 ++++++ 3 files changed, 15 insertions(+), 8 deletions(-) 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/util/util.py b/fiasco/util/util.py index 6e5dd64b..9bb16041 100644 --- a/fiasco/util/util.py +++ b/fiasco/util/util.py @@ -24,6 +24,12 @@ def _atomic_number(element): return plasmapy.particles.atomic_number(element) +@lru_cache(maxsize=None) +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. From 3285ea6d00c2ffdb3c6bb48a933f9e5b3e188883 Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Fri, 9 Oct 2026 12:34:48 -0700 Subject: [PATCH 4/4] Preserve public behavior in performance caches Refresh adjacent ions after changing proton/electron ratios or ionization fractions. Resolve database paths before caching metadata and distinguish atomic-number input types. Add numerical and cache-isolation regression tests. --- fiasco/base.py | 16 +++++++++------- fiasco/ions.py | 4 ++++ fiasco/tests/test_base.py | 28 ++++++++++++++++++++++++++++ fiasco/tests/test_ion.py | 24 ++++++++++++++++++++++++ fiasco/util/tests/test_util.py | 7 +++++++ fiasco/util/util.py | 6 +++--- 6 files changed, 75 insertions(+), 10 deletions(-) diff --git a/fiasco/base.py b/fiasco/base.py index 325bde5f..9bb02505 100644 --- a/fiasco/base.py +++ b/fiasco/base.py @@ -7,7 +7,7 @@ import pathlib import plasmapy.particles -from functools import lru_cache +from functools import cache from packaging.version import Version from plasmapy.utils import roman @@ -23,7 +23,7 @@ __all__ = ['IonBase'] -@lru_cache(maxsize=None) +@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 @@ -31,7 +31,7 @@ def _ion_names_in_dbase(hdf5_dbase_root, mtime_ns, size): return frozenset(fiasco.list_ions(hdf5_dbase_root, sort=False)) -@lru_cache(maxsize=None) +@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. @@ -73,8 +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) - stat = pathlib.Path(self.hdf5_dbase_root).stat() - if self.ion_name not in _ion_names_in_dbase(str(self.hdf5_dbase_root), stat.st_mtime_ns, stat.st_size): + 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() @@ -83,8 +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." - stat = pathlib.Path(self.hdf5_dbase_root).stat() - dbase_version = _dbase_fiasco_version(str(self.hdf5_dbase_root), stat.st_mtime_ns, stat.st_size) + 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/ions.py b/fiasco/ions.py index 76b189fd..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): """ @@ -304,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): 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 9bb16041..ad44ac1f 100644 --- a/fiasco/util/util.py +++ b/fiasco/util/util.py @@ -7,7 +7,7 @@ import plasmapy.particles import sys -from functools import lru_cache +from functools import cache, lru_cache from packaging.version import Version from plasmapy.utils import roman @@ -17,14 +17,14 @@ __all__ = ['setup_paths', 'get_chianti_catalog', 'read_chianti_version', 'parse_ion_name'] -@lru_cache(maxsize=None) +@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) -@lru_cache(maxsize=None) +@cache def _atomic_symbol(atomic_number): # Memoized version of plasmapy.particles.atomic_symbol (see _atomic_number). return plasmapy.particles.atomic_symbol(atomic_number)