Add IonCollection.build_line_list

fiasco/collections.py

@@ -6,6 +6,8 @@
 
 from astropy.convolution import convolve, Model1DKernel
 from astropy.modeling.models import Gaussian1D
+from astropy.table import QTable, vstack
+from scipy.interpolate import interp1d
 
 import fiasco
 
@@ -214,6 +216,133 @@ def two_photon(self, wavelength: u.angstrom, electron_density: u.cm**-3, **kwarg
                 two_photon += tp * ion.abundance * ion.ionization_fraction[:, np.newaxis, np.newaxis]
         return two_photon
 
+    @u.quantity_input
+    def build_line_list(self,
+                        density: u.cm**(-3),
+                        minimum_abundance=0,
+                        wavelength_range: u.angstrom = None,
+                        emission_measure=None,
+                        **kwargs):
+        r"""
+        Build a table of the emission lines of all ions in the collection.
+
+        Each row of the table is one bound-bound transition of one ion in the
+        collection, together with the contribution function of that transition
+        as a function of temperature and density. Ions without the data needed to
+        compute the contribution function are skipped with a warning.
+
+        Parameters
+        ----------
+        density: `~astropy.units.Quantity`
+            Electron number density used to compute the contribution function.
+        minimum_abundance: `float`, optional
+            Ions of elements with an abundance below this value are not included.
+        wavelength_range: `~astropy.units.Quantity`, optional
+            Two-element quantity giving the lower and upper bounds on the wavelengths of
+            transitions to include. Ions with no transition in this range are skipped.
+            Default is to include all transitions.
+        emission_measure: `~astropy.table.QTable`, optional
+            Differential emission measure model in the form returned by
+            `fiasco.get_dem_model`, with ``temperature_bin_center`` and ``em`` columns.
+            If given, the line-of-sight intensity of each transition is included in the
+            table as well. The contribution function is interpolated (in :math:`\log T`)
+            to the temperature bin centers of the model and summed over the bins; bins
+            outside the temperature range of the collection contribute nothing.
+        kwargs: `dict`, optional
+            Any keyword argument accepted by `fiasco.Ion.contribution_function`, such as
+            ``couple_density_to_temperature``, is passed through to it.
+
+        Returns
+        -------
+        : `~astropy.table.QTable`
+            One row per transition, with the properties of the transition and of the ion
+            it belongs to as columns. The ``contribution_function`` column has an entry
+            of shape ``(l, m)`` per row, where ``l`` is the number of temperatures and
+            ``m`` is the number of densities (or 1 if ``couple_density_to_temperature=True``),
+            and ``intensity``, present if ``emission_measure`` is given, has an entry of
+            shape ``(m,)``. The temperature, density, dataset names, and database and
+            fiasco versions are stored in the ``meta`` of the table.
+
+        See Also
+        --------
+        fiasco.Ion.contribution_function
+        fiasco.get_dem_model
+        """
+        density = np.atleast_1d(density)
+        if wavelength_range is None:
+            wavelength_range = u.Quantity([0, np.inf], 'angstrom')
+        if emission_measure is not None:
+            log_temperature = np.log10(self.temperature.to_value('K'))
+            log_temperature_dem = np.log10(emission_measure['temperature_bin_center'].to_value('K'))
+            em = emission_measure['em']
+        tables = []
+        included = []
+        for ion in self:
+            try:
+                abundance = ion.abundance
+                transitions = ion.transitions
+            except MissingDatasetException as e:
+                self.log.warning(f'{ion.ion_name} not included in line list. {e}')
+                continue
+            if abundance < minimum_abundance:
+                self.log.info(f'{ion.ion_name} not included in line list. '
+                              f'Abundance {abundance:.3g} is below minimum_abundance={minimum_abundance:.3g}.')
+                continue
+            is_bound_bound = transitions.is_bound_bound
+            wavelength = transitions.wavelength[is_bound_bound]
+            in_range = np.logical_and(wavelength >= wavelength_range[0],
+                                      wavelength <= wavelength_range[1])
+            if not in_range.any():
+                self.log.info(f'{ion.ion_name} not included in line list. '
+                              f'No transitions between {wavelength_range[0]:.3f} and {wavelength_range[1]:.3f}.')
+                continue
+            try:
+                g = ion.contribution_function(density, **kwargs)
+            except MissingDatasetException as e:
+                self.log.warning(f'{ion.ion_name} not included in line list. {e}')
+                continue
+            # Put the transition axis first so that each row of the table holds the
+            # contribution function of one transition
+            g = np.moveaxis(g[..., in_range], -1, 0)
+            n_lines = in_range.sum()
+            table = QTable({
+                'element': [ion.atomic_symbol] * n_lines,
+                'atomic number': [ion.atomic_number] * n_lines,
+                'ion name': [ion.ion_name_roman] * n_lines,
+                'abundance': [abundance] * n_lines,
+                'ionization potential': u.Quantity([ion.ionization_potential] * n_lines),
+                'wavelength': wavelength[in_range],
+                'energy': transitions.delta_energy[is_bound_bound][in_range],
+                'A': transitions.A[is_bound_bound][in_range],
+                'upper level': transitions.upper_level[is_bound_bound][in_range],
+                'lower level': transitions.lower_level[is_bound_bound][in_range],
+                'upper label': transitions.upper_label[is_bound_bound][in_range],
+                'lower label': transitions.lower_label[is_bound_bound][in_range],
+                'is observed': transitions.is_observed[is_bound_bound][in_range],
+                'contribution function': g,
+            })
+            if emission_measure is not None:
+                f_interp = interp1d(log_temperature, g.value, axis=1, bounds_error=False, fill_value=0.0)
+                g_dem = f_interp(log_temperature_dem) * g.unit
+                intensity = (g_dem * em[np.newaxis, :, np.newaxis]).sum(axis=1)
+                table['intensity'] = intensity / (4*np.pi*u.steradian)
+            tables.append(table)
+            included.append(ion)
+        if not tables:
+            raise ValueError('No lines found for any ion in collection.')
+        line_list = vstack(tables)
+        line_list.meta['temperature'] = self.temperature
+        line_list.meta['density'] = density
+        # Each ion in a collection can use its own abundance and ionization fraction
+        # datasets, so record the names from the ions that made it into the table.
+        # The name is None for an ion whose dataset was set to a value directly.
+        for key in ('abundance', 'ionization_fraction'):
+            names = sorted({ion._dset_names[key] for ion in included}, key=str)
+            line_list.meta[key] = names[0] if len(names) == 1 else names
+        line_list.meta['chianti_version'] = str(included[0]._elvlc.version)
+        line_list.meta['fiasco_version'] = fiasco.__version__
+        return line_list
+
     @u.quantity_input
     def spectrum(self, density: u.cm**(-3), emission_measure: u.cm**(-5), wavelength_range=None,
                  bin_width=None, kernel=None, **kwargs):

fiasco/tests/test_collections.py

@@ -2,6 +2,7 @@
 Test collection functionality
 """
 import astropy.units as u
+import logging
 import numpy as np
 import pytest
 
@@ -198,6 +199,109 @@ def test_spectrum(hdf5_dbase_root):
     assert np.all(spec == spec2)
 
 
+line_list_density = 1e9 * u.cm**-3
+
+
+@pytest.fixture
+def line_list_collection(hdf5_dbase_root):
+    h1 = fiasco.Ion('H 1', temperature, hdf5_dbase_root=hdf5_dbase_root)
+    fe5 = fiasco.Ion('Fe 5', temperature, hdf5_dbase_root=hdf5_dbase_root)
+    # H 2 has no transition data and should be skipped
+    h2 = fiasco.Ion('H 2', temperature, hdf5_dbase_root=hdf5_dbase_root)
+    return h1 + fe5 + h2
+
+
+@pytest.fixture
+def line_list(line_list_collection):
+    return line_list_collection.build_line_list(line_list_density)
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_build_line_list_rows(line_list, line_list_collection):
+    h1, fe5, _ = line_list_collection
+    n_h1 = h1.transitions.is_bound_bound.sum()
+    n_fe5 = fe5.transitions.is_bound_bound.sum()
+    assert len(line_list) == n_h1 + n_fe5
+    assert set(line_list['ion name']) == {'H I', 'Fe V'}
+    assert 'intensity' not in line_list.colnames
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_build_line_list_levels(line_list, line_list_collection):
+    _, fe5, _ = line_list_collection
+    is_fe5 = line_list['ion name'] == 'Fe V'
+    assert np.array_equal(line_list['upper level'][is_fe5],
+                          fe5.transitions.upper_level[fe5.transitions.is_bound_bound])
+    assert np.all(line_list['upper level'] > line_list['lower level'])
+    assert np.all(line_list['energy'] > 0 * u.eV)
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_build_line_list_contribution_function(line_list, line_list_collection):
+    _, fe5, _ = line_list_collection
+    assert line_list['contribution function'].shape == (len(line_list), temperature.shape[0], 1)
+    g_fe5 = fe5.contribution_function(line_list_density)
+    assert u.allclose(line_list['contribution function'][line_list['ion name'] == 'Fe V'],
+                      np.moveaxis(g_fe5, -1, 0))
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_build_line_list_meta(line_list, line_list_collection):
+    h1, _, _ = line_list_collection
+    assert u.allclose(line_list.meta['temperature'], temperature)
+    assert u.allclose(line_list.meta['density'], line_list_density)
+    assert line_list.meta['abundance'] == h1._dset_names['abundance']
+    assert line_list.meta['ionization_fraction'] == h1._dset_names['ionization_fraction']
+    assert line_list.meta['chianti_version'] == str(h1._elvlc.version)
+    assert line_list.meta['fiasco_version'] == fiasco.__version__
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_build_line_list_minimum_abundance(line_list_collection):
+    # Selecting by abundance removes iron. H II passes the abundance cut but has no
+    # transitions, so it is never in the line list.
+    line_list = line_list_collection.build_line_list(line_list_density, minimum_abundance=1e-3)
+    assert set(line_list['ion name']) == {'H I'}
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_build_line_list_skipped_ions_logged(line_list_collection, caplog):
+    with caplog.at_level(logging.INFO, logger='fiasco'):
+        line_list_collection.build_line_list(line_list_density, minimum_abundance=1e-3)
+    assert 'Fe 5 not included in line list. Abundance' in caplog.text
+    caplog.clear()
+    with caplog.at_level(logging.INFO, logger='fiasco'):
+        line_list_collection.build_line_list(line_list_density, wavelength_range=[1215, 1216] * u.angstrom)
+    assert 'Fe 5 not included in line list. No transitions between' in caplog.text
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_build_line_list_wavelength_range(line_list_collection):
+    wavelength_range = [900, 1300] * u.angstrom
+    line_list = line_list_collection.build_line_list(line_list_density, wavelength_range=wavelength_range)
+    assert np.all(line_list['wavelength'] >= wavelength_range[0])
+    assert np.all(line_list['wavelength'] <= wavelength_range[1])
+    assert 'H I' in set(line_list['ion name'])
+
+
+@pytest.mark.requires_dbase_version('>= 8')
+def test_build_line_list_intensity(hdf5_dbase_root):
+    fe5 = fiasco.Ion('Fe 5', temperature, hdf5_dbase_root=hdf5_dbase_root)
+    c = fiasco.IonCollection(fe5)
+    density = [1e9, 1e10] * u.cm**-3
+    dem = fiasco.get_dem_model('quiet_sun', hdf5_dbase_root=hdf5_dbase_root)
+    line_list = c.build_line_list(density, emission_measure=dem)
+    assert line_list['intensity'].shape == (len(line_list), 2)
+    assert line_list['intensity'].unit.is_equivalent(u.erg / u.cm**2 / u.s / u.steradian)
+    assert np.all(line_list['intensity'] >= 0)
+
+
+def test_build_line_list_no_valid_ions(hdf5_dbase_root):
+    c = fiasco.IonCollection(fiasco.Ion('H 2', temperature, hdf5_dbase_root=hdf5_dbase_root))
+    with pytest.raises(ValueError, match='No lines found for any ion in collection.'):
+        c.build_line_list(1e9 * u.cm**-3)
+
+
 def test_spectrum_no_valid_ions(hdf5_dbase_root):
     # Consider the case of an collection with ions with no spectral information
     c2 = fiasco.IonCollection(fiasco.Ion('H 2', 1 * u.MK, hdf5_dbase_root=hdf5_dbase_root))