Skip to content

Commit

Permalink
Refactor psplups treatment in level_populations and add test for …
Browse files Browse the repository at this point in the history
…full database (#229)

* psplups refactoring

* add tests for full database

* test missing recombining ion codepath
  • Loading branch information
wtbarnes authored Mar 13, 2023
1 parent 66a0a3e commit bb4670e
Show file tree
Hide file tree
Showing 6 changed files with 66 additions and 23 deletions.
12 changes: 11 additions & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,16 @@ jobs:
cache-key: chianti-${{ github.event.number }}
secrets:
CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
test_full_database:
needs: [test]
uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@v1
with:
posargs: '--ascii-dbase-root ~/.chianti --include-all-files'
toxdeps: "'tox<4' tox-pypi-filter"
envs: |
- linux: py310
cache-path: ~/.chianti
cache-key: chianti-${{ github.event.number }}
precommit:
uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@v1
with:
Expand Down Expand Up @@ -57,7 +67,7 @@ jobs:
github.event_name == 'pull_request' &&
contains(github.event.pull_request.labels.*.name, 'Run publish')
)
needs: [test]
needs: [test,test_full_database]
uses: OpenAstronomy/github-actions-workflows/.github/workflows/publish_pure_python.yml@main
with:
test_extras: 'dev'
Expand Down
2 changes: 1 addition & 1 deletion conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,6 @@ def pytest_addoption(parser):
help='Disable MD5 hash checks on test files')
parser.addoption('--idl-executable', action='store', default=None)
parser.addoption('--idl-codebase-root', action='store', default=None)
parser.addoption('--include-all-files', action='store', default=False)
parser.addoption('--include-all-files', action='store_true', default=False)
parser.addoption('--skip-version-check', action='store_true', default=False,
help='Do not check CHIANTI version')
2 changes: 2 additions & 0 deletions fiasco/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
'h_1.wgfa': 'ed2a561ecdba5ee4b05ea56c297724ba',
'h_1.scups': 'de180c9e1b4f50a503efad8c83e714ab',
'h_1.diparams': 'f9be79794cc092c75733ece7aba9f29f',
'h_1.fblvl': 'af45fd1ee7af4b4e71aba288d0b16bc8',
'h_2.rrparams': '05eb5044dc1ad070338d1ba3745dc4de',
'he_1.elvlc': '577245da46cfc4d27a05f40147f17610',
'he_2.elvlc': '58eee740f6842850f1f35a4f95b3e12c',
Expand Down Expand Up @@ -71,6 +72,7 @@
'fe_5.wgfa': '6f8e4d41760d5a0540008e15aca1038c',
'fe_6.elvlc': '081519f986b8a8ed99a34ecf813f1358',
'fe_6.fblvl': '9478aeab2479e9eefebfd1e552e27ddf',
'fe_7.fblvl': 'ad0264f06018b0b9cca1eba502809708',
'fe_9.elvlc': 'b7d04f65a87a8de1c2bfc77331e438f3',
'fe_12.elvlc': 'ee9beb8b2ff03ba8fc046bd722992b21',
'fe_15.elvlc': 'cb8e6291bac17325130ea8564e118d48',
Expand Down
50 changes: 30 additions & 20 deletions fiasco/ions.py
Original file line number Diff line number Diff line change
Expand Up @@ -411,37 +411,36 @@ def level_populations(self,
"""
Energy level populations as a function of temperature and density.
Compute the level populations of the given ion as a function of temperature and
density. This is done by solving the homogeneous linear system of equations
describing the processes that populate and depopulate each energy level of each
ion. Section 3 of :cite:t:`young_chianti_2016` provides a brief description of
this set of equations.
Parameters
----------
density : `~astropy.units.Quantity`
include_protons : `bool`, optional
If True (default), include proton excitation and de-excitation rates.
include_level_resolved_rate_correction: `bool`, optional
If True (default), include the level-resolved ionization and recombination rate
correction in the resulting level populations as described in Section 2.3 of
:cite:t:`landi_chianti-atomic_2006`.
couple_density_to_temperature: `bool`, optional
If True, the density will vary along the same axis as temperature
in the computed level populations. The number of densities must be the same as the number of temperatures. This is useful, for
example, when computing the level populations at constant
pressure and is also much faster than computing the level
populations along an independent density axis. By default, this
is set to False.
in the computed level populations and the number of densities must be the same as
the number of temperatures. This is useful, for example, when computing the level
populations at constant pressure and is also much faster than computing the level
populations along an independent density axis. By default, this is set to False.
Returns
-------
`~astropy.units.Quantity`
A ``(l, m, n)`` shaped quantity, where ``l`` is the number of
temperatures, ``m`` is the number of densities, and ``n``
is the number of energy levels. If
``couple_density_to_temperature=True``, then ``m=1`` and ``l``
temperatures, ``m`` is the number of densities, and ``n`` is the number of energy
levels. If ``couple_density_to_temperature=True``, then ``m=1`` and ``l``
represents the number of temperatures and densities.
"""
# NOTE: Cannot include protons if psplups data not available
try:
_ = self._psplups
except KeyError:
self.log.warning(
f'No proton data available for {self.ion_name}. '
'Not including proton excitation and de-excitation in level populations calculation.')
include_protons = False

density = np.atleast_1d(density)
if couple_density_to_temperature:
if density.shape != self.temperature.shape:
Expand All @@ -467,7 +466,20 @@ def level_populations(self,
lower_level, level, ex_rate_e.value.T, 0).T * ex_rate_e.unit
dex_diagonal_e = vectorize_where_sum(
upper_level, level, dex_rate_e.value.T, 0).T * dex_rate_e.unit

# Collisional--protons
if include_protons:
# NOTE: Cannot include protons if psplups data not available for this ion
try:
ex_rate_p = self.proton_collision_excitation_rate
dex_rate_p = self.proton_collision_deexcitation_rate
except MissingDatasetException:
self.log.warning(
f'No proton data available for {self.ion_name}. '
'Not including proton excitation and de-excitation in level populations calculation.')
include_protons = False
# NOTE: Having two of the same if blocks here is ugly, but necessary. We cannot continue
# with the proton rate calculation if the data is not available.
if include_protons:
lower_level_p = self._psplups['lower_level']
upper_level_p = self._psplups['upper_level']
Expand All @@ -476,8 +488,6 @@ def level_populations(self,
proton_density = (pe_ratio * density)[:, np.newaxis, np.newaxis]
else:
proton_density = np.outer(pe_ratio, density)[:, :, np.newaxis]
ex_rate_p = self.proton_collision_excitation_rate
dex_rate_p = self.proton_collision_deexcitation_rate
ex_diagonal_p = vectorize_where_sum(
lower_level_p, level, ex_rate_p.value.T, 0).T * ex_rate_p.unit
dex_diagonal_p = vectorize_where_sum(
Expand Down Expand Up @@ -508,7 +518,7 @@ def level_populations(self,
# Excitation from lower states
c_matrix[:, upper_level-1, lower_level-1] += d*ex_rate_e
# Same processes as above, but for protons
if include_protons and self._psplups is not None:
if include_protons:
d_p = proton_density[:, i, :]
c_matrix[:, level-1, level-1] -= d_p*(ex_diagonal_p + dex_diagonal_p)
c_matrix[:, lower_level_p-1, upper_level_p-1] += d_p * dex_rate_p
Expand Down
2 changes: 1 addition & 1 deletion fiasco/tests/test_collections.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def test_free_bound(another_collection, wavelength):
assert fb.shape == temperature.shape + wavelength.shape if wavelength.shape else (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'))
assert u.allclose(fb[index_t, index_w], 3.057781475607237e-36 * u.Unit('erg cm3 s-1 Angstrom-1'))


def test_radiative_los(collection):
Expand Down
21 changes: 21 additions & 0 deletions fiasco/tests/test_ion.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,11 @@ def another_ion(hdf5_dbase_root):
return fiasco.Ion('Fe 6', temperature, hdf5_dbase_root=hdf5_dbase_root)


@pytest.fixture
def h1(hdf5_dbase_root):
return fiasco.Ion('H 1', temperature, hdf5_dbase_root=hdf5_dbase_root)


@pytest.fixture
def fe10(hdf5_dbase_root):
return fiasco.Ion('Fe 10', temperature, hdf5_dbase_root=hdf5_dbase_root)
Expand Down Expand Up @@ -180,6 +185,13 @@ def test_level_populations(ion):
assert u.allclose(pop.squeeze().sum(axis=1), 1, atol=None, rtol=1e-15)


def test_level_populations_proton_data_toggle(ion):
# Fe V has no psplups data so the toggle should have no effect
lp_protons = ion.level_populations(1e9*u.cm**(-3), include_protons=True)
lp_no_protons = ion.level_populations(1e9*u.cm**(-3), include_protons=False)
assert u.allclose(lp_protons, lp_no_protons, atol=0, rtol=0)


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
Expand Down Expand Up @@ -307,6 +319,15 @@ def test_free_bound(ion):
assert u.allclose(emission[0, 0], 9.7902609e-26 * u.cm**3 * u.erg / u.Angstrom / u.s)


def test_free_bound_no_recombining(h1):
# This is test the case where there is no data available for the recombining
# ion (H 2)
emission = h1.free_bound(200 * u.Angstrom)
assert emission.shape == h1.temperature.shape + (1, )
# This value has not been tested for correctness
assert u.allclose(emission[0, 0], 1.9611545671496785e-28 * u.cm**3 * u.erg / u.Angstrom / u.s)


def test_add_ions(ion, another_ion):
collection = ion + another_ion
assert isinstance(collection, fiasco.IonCollection)
Expand Down

0 comments on commit bb4670e

Please sign in to comment.