From bb4670e2df538c674149285235d9555f736d0960 Mon Sep 17 00:00:00 2001 From: Will Barnes Date: Mon, 13 Mar 2023 12:43:28 -0400 Subject: [PATCH] Refactor `psplups` treatment in `level_populations` and add test for full database (#229) * psplups refactoring * add tests for full database * test missing recombining ion codepath --- .github/workflows/ci.yml | 12 +++++++- conftest.py | 2 +- fiasco/conftest.py | 2 ++ fiasco/ions.py | 50 +++++++++++++++++++------------- fiasco/tests/test_collections.py | 2 +- fiasco/tests/test_ion.py | 21 ++++++++++++++ 6 files changed, 66 insertions(+), 23 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4f024796..c9583ba8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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: @@ -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' diff --git a/conftest.py b/conftest.py index 2e7d7ce4..1703d14f 100644 --- a/conftest.py +++ b/conftest.py @@ -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') diff --git a/fiasco/conftest.py b/fiasco/conftest.py index 15519c7f..365be2dc 100644 --- a/fiasco/conftest.py +++ b/fiasco/conftest.py @@ -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', @@ -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', diff --git a/fiasco/ions.py b/fiasco/ions.py index 6dd30dbe..e0a21819 100644 --- a/fiasco/ions.py +++ b/fiasco/ions.py @@ -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: @@ -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'] @@ -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( @@ -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 diff --git a/fiasco/tests/test_collections.py b/fiasco/tests/test_collections.py index 507ae1f4..f8e47541 100644 --- a/fiasco/tests/test_collections.py +++ b/fiasco/tests/test_collections.py @@ -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): diff --git a/fiasco/tests/test_ion.py b/fiasco/tests/test_ion.py index fa7ea952..f743daf7 100644 --- a/fiasco/tests/test_ion.py +++ b/fiasco/tests/test_ion.py @@ -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) @@ -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 @@ -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)