Skip to content

Fast X: level populations edition - #466

Open
nabobalis wants to merge 4 commits into
wtbarnes:mainfrom
nabobalis:speed_3
Open

nabobalis wants to merge 4 commits into
wtbarnes:mainfrom
nabobalis:speed_3

Conversation

@nabobalis

@nabobalis nabobalis commented Aug 10, 2026 •

Copy link
Copy Markdown
Collaborator

Calls to level_populations with multiple densities repeatedly create adjacent ions and rebuild their rate matrices. This PR reuses those ions and the summed two-ion rate coefficients across density values. It also reduces HDF5 opens and caches database metadata and atomic lookups.

The caches preserve public setter behavior: changing an ion's proton/electron ratio or ionization fraction refreshes its adjacent ions. Resolved database paths avoid cache collisions across directories, and typed atomic-number caching preserves input validation.

These measurements compare a7e8ccb with upstream/main at 668833a, which already includes the proton/electron-ratio improvements from #471. They replace the earlier measurements against the older baseline.

Case Upstream main This PR Speedup
Fe XXIII, two-ion model 5.868 s 4.877 s 1.20x
Ca XVIII, two-ion model 9.263 s 8.371 s 1.11x
O VI, default options 9.289 s 7.879 s 1.18x
Fe XII, single-ion model 3.014 s 2.774 s 1.09x
Fe XXIII, coupled temperature/density 1.724 s 1.432 s 1.20x
proton_electron_ratio 0.763 s 0.553 s 1.38x

Medians of three timed calls after one warm-up per case, using the full CHIANTI 11.0.2 database on an AMD Ryzen 9 9950X with one linear-algebra thread. Every timed population calculation constructs a fresh Ion; construction is included. Database/module caches remain warm. All six benchmark outputs are bit-for-bit identical to upstream.

Temperature: 21 logarithmically spaced points from 10^5.5 to 10^7.5 K. Density: nine points from 10^8 to 10^12 cm^-3. The coupled case uses density = 1e15 K cm^-3 / temperature and the two-ion model.

Both versions used Python 3.14.5, NumPy 2.4.6, SciPy 1.17.0, Astropy 8.1.0.dev684+g2b5b20b6d, h5py 3.16.0, and PlasmaPy 2026.2.0.

Validation: 63 focused tests pass, including numerical regression tests for setter changes and tests for database-cache isolation and input validation. Changed files pass Ruff.

Reproduce the timing comparison

Run the following on each checkout with the same full CHIANTI database configured and OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1:

import gc
from statistics import median
from time import perf_counter

import astropy.units as u
import fiasco
import numpy as np

fiasco.log.setLevel("ERROR")
temperature = np.logspace(5.5, 7.5, 21) * u.K
density = np.logspace(8, 12, 9) * u.cm**-3
cases = [
    ("Fe 23", density, {"use_two_ion_model": True}),
    ("Ca 18", density, {"use_two_ion_model": True}),
    ("O 6", density, {}),
    ("Fe 12", density, {"use_two_ion_model": False}),
    ("Fe 23", 1e15 * u.K * u.cm**-3 / temperature,
     {"use_two_ion_model": True, "couple_density_to_temperature": True}),
    (None, None, {}),
]
for ion_name, case_density, options in cases:
    def run():
        if ion_name is None:
            return fiasco.proton_electron_ratio(temperature)
        ion = fiasco.Ion(ion_name, temperature)
        return ion.level_populations(case_density, **options)

    reference = run()
    timings = []
    for _ in range(3):
        gc.collect()
        start = perf_counter()
        result = run()
        timings.append(perf_counter() - start)
        np.testing.assert_array_equal(result, reference)
    print(ion_name or "proton_electron_ratio", options, median(timings))

@nabobalis nabobalis changed the title Speed 3 - Needs Speed 2 Fast X: level populations edition Aug 11, 2026
@nabobalis

nabobalis commented Aug 11, 2026 •

Copy link
Copy Markdown
Collaborator Author

LOOK AT THE CI FLY. ALSO AS GOOD AS MY PIDGEOT WITH HM FLY

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.
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.
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.
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant