Source code for fiasco.io.sources.continuum_sources

"""
Source classes for continuum files
"""
import astropy.units as u
import h5py
import numpy as np
import pathlib
import plasmapy

from fiasco.io.generic import GenericParser

__all__ = [
    'GenericContinuumParser',
    'GffguParser',
    'GffintParser',
    'ItohIntegratedGauntParser',
    'ItohIntegratedGauntNonrelParser',
    'KlgfbParser',
    'VernerParser',
    'ItohParser',
    'HSeqParser',
    'HeSeqParser'
]




class GenericContinuumParser(GenericParser):

    @property
    def full_path(self):
        if hasattr(self, '_full_path'):
            return self._full_path
        elif self.standalone:
            return self.filename
        else:
            return pathlib.Path(self.ascii_dbase_root / 'continuum' / self.filename)





class GffguParser(GenericContinuumParser):
    """
    Free-free Gaunt factor as a function of scaled frequency and energy
    """
    filetype = 'gffgu'
    dtypes = [float, float, float]
    units = [u.dimensionless_unscaled, u.dimensionless_unscaled, u.dimensionless_unscaled]
    headings = ['u', 'gamma_squared', 'gaunt_factor']
    descriptions = ['scaled frequency', 'scaled temperature', 'free-free Gaunt factor']

    def __init__(self, filename, **kwargs):
        super().__init__(filename, **kwargs)
        self.body_index = 5



    def preprocessor(self, table, line, index):
        if index >= self.body_index and '--' not in line:
            super().preprocessor(table, line, index)




    def extract_footer(self, lines):
        comment = []
        for i, l in enumerate(lines):
            if i <= self.body_index - 3:
                comment.append(l)
            else:
                break

        footer = '\n'.join([l.strip() for l in comment])
        return footer




    def to_hdf5(self, hf, df, **kwargs):
        grp_name = '/'.join(['continuum', self.filetype])
        if grp_name not in hf:
            grp = hf.create_group(grp_name)
            grp.attrs['chianti_version'] = df.meta['chianti_version']
            grp.attrs['footer'] = df.meta['footer']
        else:
            grp = hf[grp_name]

        for name in df.colnames:
            col = df[name]
            if isinstance(col, u.Quantity):
                data = col.value
            else:
                data = col.data
            if '<U' in data.dtype.str:
                numchar = data.dtype.str[2:]
                data = data.astype(f'|S{numchar}')
            if name in grp:
                ds = grp[name]
            else:
                if data.dtype == np.dtype('O'):
                    ragged_dtype = h5py.special_dtype(vlen=np.dtype('float64'))
                    ds = grp.create_dataset(name, data=data, dtype=ragged_dtype)
                else:
                    ds = grp.create_dataset(name, data=data, dtype=data.dtype)
            if col.unit is None:
                ds.attrs['unit'] = 'SKIP'
            else:
                ds.attrs['unit'] = col.unit.to_string()
            ds.attrs['description'] = df.meta['descriptions'][name]






class GffintParser(GffguParser):
    """
    Total free-free Gaunt factor as a function of scaled temperature.
    """
    filetype = 'gffint'
    dtypes = [float, float, float, float, float]
    units = [u.dimensionless_unscaled, u.dimensionless_unscaled, u.dimensionless_unscaled,
             u.dimensionless_unscaled, u.dimensionless_unscaled]
    headings = ['log_gamma_squared', 'gaunt_factor', 's1', 's2', 's3']
    descriptions = ['log scaled temperature', 'total free-free Gaunt factor', 'spline coefficient',
                    'spline coefficient', 'spline coefficient']

    def __init__(self, filename, **kwargs):
        super().__init__(filename, **kwargs)
        self.body_index = 4




class ItohIntegratedGauntParser(GenericContinuumParser):
    """
    Total (frequency-integrated) relativistic free-free Gaunt factor as a
    function of a scaled temperature and scaled atomic number.
    """
    filetype='itoh_integrated_gaunt'
    dtypes=[float]
    units = [u.dimensionless_unscaled]
    headings = ['a_ik']
    descriptions = ['fitting coefficient']



    def preprocessor(self, table, line, index):
        line = line.strip().split()
        gf = np.array(line, dtype=float)
        table.append([gf])




    def extract_footer(self, *args):
        return """Fit parameters for relativistic free-free Gaunt factors integrated over frequency.
Itoh et al., 2002, A&A, 382, 722
comment: These are the coefficients a_ik tabulated in Table 1."""




    def to_hdf5(self, hf, df, **kwargs):
        grp_name = '/'.join(['continuum', self.filetype])
        if grp_name not in hf:
            grp = hf.create_group(grp_name)
            grp.attrs['chianti_version'] = df.meta['chianti_version']
            grp.attrs['footer'] = df.meta['footer']
        else:
            grp = hf[grp_name]

        for name in df.colnames:
            col = df[name]
            if isinstance(col, u.Quantity):
                data = col.value
            else:
                data = col.data
            if '<U' in data.dtype.str:
                numchar = data.dtype.str[2:]
                data = data.astype(f'|S{numchar}')
            if name in grp:
                ds = grp[name]
            else:
                if data.dtype == np.dtype('O'):
                    ragged_dtype = h5py.special_dtype(vlen=np.dtype('float64'))
                    ds = grp.create_dataset(name, data=data, dtype=ragged_dtype)
                else:
                    ds = grp.create_dataset(name, data=data, dtype=data.dtype)
            ds.attrs['unit'] = col.unit.to_string()
            ds.attrs['description'] = df.meta['descriptions'][name]






class ItohIntegratedGauntNonrelParser(GenericContinuumParser):
    """
    Total (frequency-integrated) non-relativistic free-free Gaunt factor as a
    function of a scaled temperature.
    """
    filetype='itoh_integrated_gaunt_nonrel'
    dtypes=[float]
    units = [u.dimensionless_unscaled]
    headings = ['b_i']
    descriptions = ['fitting coefficient']



    def extract_footer(self, *args):
        return """Fit coefficients for non-relativistic, frequency integrated free-free Gaunt factor
Itoh et al. (2002, A&A, 382, 722)
Comment: Data taken from Table 2 of this work."""




    def to_hdf5(self, hf, df, **kwargs):
        grp_name = '/'.join(['continuum', self.filetype])
        if grp_name not in hf:
            grp = hf.create_group(grp_name)
            grp.attrs['chianti_version'] = df.meta['chianti_version']
            grp.attrs['footer'] = df.meta['footer']
        else:
            grp = hf[grp_name]

        for name in df.colnames:
            col = df[name]
            if isinstance(col, u.Quantity):
                data = col.value
            else:
                data = col.data
            if '<U' in data.dtype.str:
                numchar = data.dtype.str[2:]
                data = data.astype(f'|S{numchar}')
            if name in grp:
                ds = grp[name]
            else:
                if data.dtype == np.dtype('O'):
                    ragged_dtype = h5py.special_dtype(vlen=np.dtype('float64'))
                    ds = grp.create_dataset(name, data=data, dtype=ragged_dtype)
                else:
                    ds = grp.create_dataset(name, data=data, dtype=data.dtype)
            ds.attrs['unit'] = col.unit.to_string()
            ds.attrs['description'] = df.meta['descriptions'][name]






class KlgfbParser(GenericContinuumParser):
    """
    Free-bound gaunt factor as a function of photon energy for several different energy levels.
    """
    filetype = 'klgfb'
    dtypes = [int, int, float, float]
    units = [None, None, u.dimensionless_unscaled, u.dimensionless_unscaled]
    headings = ['n', 'l', 'log_pe', 'log_gaunt_factor']
    descriptions = [
        'principal quantum number',
        'orbital angular momentum number',
        'log photon energy divided by ionization potential',
        'log free-bound Gaunt factor',
    ]



    def preprocessor(self, table, line, index):
        if index == 0:
            pass
        elif index == 1:
            self._photon_energy = np.array(line.strip().split(), dtype=float)
        else:
            line = line.strip().split()
            n, l = line[:2]
            gaunt_factor = np.array(line[2:], dtype=float)
            if gaunt_factor.shape != self._photon_energy.shape:
                raise ValueError('K&L gaunt factor and photon energy have unequal dimensions')
            for pe, gf in zip(self._photon_energy, gaunt_factor):
                table.append([n, l, pe, gf])




    def extract_footer(self, *args):
        return """Log of free-bound Gaunt factors as a function of log of photon energy divided by ionization potential
From Karzas, W. J. and Latter, R., 1961, ApJS, 6, 167"""




    def to_hdf5(self, hf, df, **kwargs):
        grp_name = '/'.join(['continuum', self.filetype])
        if grp_name not in hf:
            grp = hf.create_group(grp_name)
            grp.attrs['chianti_version'] = df.meta['chianti_version']
            grp.attrs['footer'] = df.meta['footer']
        else:
            grp = hf[grp_name]

        for name in df.colnames:
            col = df[name]
            if isinstance(col, u.Quantity):
                data = col.value
            else:
                data = col.data
            if '<U' in data.dtype.str:
                numchar = data.dtype.str[2:]
                data = data.astype(f'|S{numchar}')
            if name in grp:
                ds = grp[name]
            else:
                if data.dtype == np.dtype('O'):
                    ragged_dtype = h5py.special_dtype(vlen=np.dtype('float64'))
                    ds = grp.create_dataset(name, data=data, dtype=ragged_dtype)
                else:
                    ds = grp.create_dataset(name, data=data, dtype=data.dtype)
            if col.unit is None:
                ds.attrs['unit'] = 'SKIP'
            else:
                ds.attrs['unit'] = col.unit.to_string()
            ds.attrs['description'] = df.meta['descriptions'][name]




class KlgfbNParser(GenericContinuumParser):
    """
    Free-bound gaunt factor as a function of photon energy for a given principal quantum number.

    .. note::

        This parser returns the same information as `~fiasco.io.sources.continuum_sources.KlgfbParser`
        but exists because the file format for the Karzas and Latter Gaunt factors changed in v10
        of the database.
    """
    dtypes = [int, int, float, float]
    units = [None, None, u.Ry, u.dimensionless_unscaled]
    headings = ['n', 'l', 'photon_energy', 'gaunt_factor']
    descriptions = [
        'principal quantum number',
        'orbital angular momentum number',
        'photon energy',
        'free-bound Gaunt factor',
    ]

    @property
    def pqn(self):
        return int(self.filetype.split('_')[1])

    def extract_footer(self, *args):
        return f"""Free-bound Gaunt factors as a function of photon energy for principal quantum number n={self.pqn}
From Karzas, W. J. and Latter, R., 1961, ApJS, 6, 167"""

    def preprocessor(self, table, line, index):
        line = line.strip().split()
        pe = float(line[0])
        gaunt_factor = np.array(line[1:], dtype=float)
        for i,gf in enumerate(gaunt_factor):
            table.append([self.pqn, i, pe, gf])

    def to_hdf5(self, hf, df, **kwargs):
        # NOTE: The group name is the same for all klgfb_n files
        grp_name = 'continuum/klgfb'
        if grp_name not in hf:
            grp = hf.create_group(grp_name)
            grp.attrs['chianti_version'] = df.meta['chianti_version']
            grp.attrs['footer'] = df.meta['footer']
        else:
            grp = hf[grp_name]
        for name in df.colnames:
            col = df[name]
            data = col.value if isinstance(col, u.Quantity) else col.data
            # This conditional allows for concatenating multiple files into a single
            # dataset within the HDF5 database
            if name in grp:
                ds = grp[name]
                ds.resize(ds.shape[0]+data.shape[0], axis=0)
                ds[-data.shape[0]:] = data
            else:
                ds = grp.create_dataset(name, data=data, dtype=data.dtype, maxshape=(None,))
                if col.unit is None:
                    ds.attrs['unit'] = 'SKIP'
                else:
                    ds.attrs['unit'] = col.unit.to_string()
                ds.attrs['description'] = df.meta['descriptions'][name]


# These are implemented separately because of the way the parser factory works.
# This is more a quirk of this information being split amongst multiple files.
class Klgfb1Parser(KlgfbNParser):
    filetype = 'klgfb_1'


class Klgfb2Parser(KlgfbNParser):
    filetype = 'klgfb_2'


class Klgfb3Parser(KlgfbNParser):
    filetype = 'klgfb_3'


class Klgfb4Parser(KlgfbNParser):
    filetype = 'klgfb_4'


class Klgfb5Parser(KlgfbNParser):
    filetype = 'klgfb_5'


class Klgfb6Parser(KlgfbNParser):
    filetype = 'klgfb_6'




class VernerParser(GenericContinuumParser):
    """
    Fit parameters for calculating partial photoionization cross-sections using
    the method of :cite:t:`verner_analytic_1995`.
    """
    filetype = 'verner_short'
    dtypes = [int, int, int, int, float, float, float, float, float, float]
    units = [None, None, None, None, u.eV, u.eV, u.megabarn, u.dimensionless_unscaled,
             u.dimensionless_unscaled, u.dimensionless_unscaled]
    headings = ['Z', 'n_electrons', 'n', 'l', 'E_thresh', 'E_0_fit', 'sigma_0', 'y_a_fit', 'P_fit',
                'y_w_fit']
    descriptions = ['atomic number', 'number of electrons', 'principal quantum number',
                    'orbital angular momentum number',
                    'threshold energy below which cross-section is 0',
                    'E_0 fit parameter', 'nominal value of cross-section', 'y_a fit parameter',
                    'P fit parameter', 'y_w fit parameter']



    def extract_footer(self, *args):
        return """Fit parameters for calculating partial photoionization cross-sections for individual ions
From Verner, D. A . and Yakovlev, D. G., 1995, A&AS, 109, 125"""




    def to_hdf5(self, hf, df, **kwargs):
        for row in df:
            el = plasmapy.particles.atomic_symbol(int(row['Z'])).lower()
            stage = row['Z'] - row['n_electrons'] + 1
            grp_name = f'{el}/{el}_{stage}/continuum/{self.filetype}'
            if grp_name not in hf:
                grp = hf.create_group(grp_name)
                grp.attrs['chianti_version'] = df.meta['chianti_version']
                grp.attrs['footer'] = df.meta['footer']
            else:
                grp = hf[grp_name]
            for col in row.colnames:
                if col == 'Z' or col == 'n_electrons':
                    continue
                ds = grp.create_dataset(col, data=row[col])
                ds.attrs['description'] = df.meta['descriptions'][col]
                if not hasattr(row[col], 'unit'):
                    ds.attrs['unit'] = 'SKIP'
                else:
                    ds.attrs['unit'] = row[col].unit.to_string()






class ItohParser(GenericContinuumParser):
    """
    Fit parameters for calculating relativistic free-free Gaunt factor using the method of
    :cite:t:`itoh_relativistic_2000`.
    """
    filetype = 'itoh'
    dtypes = [int, float]
    units = [None, u.dimensionless_unscaled]
    headings = ['Z', 'a']
    descriptions = ['atomic number', 'fit coefficient']



    def preprocessor(self, table, line, index):
        a_matrix = np.array(line.strip().split()).reshape((11, 11))
        table.append([index+1] + [a_matrix])




    def extract_footer(self, *args):
        return """Analytic fit coefficients as a function of scaled temperature and energy for calculating the relativistic free-free Gaunt factor
From Itoh, N., et al., ApJS, 2000, 128, 125"""




    def to_hdf5(self, hf, df, **kwargs):
        grp_name = '/'.join(['continuum', self.filetype])
        if grp_name not in hf:
            grp = hf.create_group(grp_name)
            grp.attrs['chianti_version'] = df.meta['chianti_version']
            grp.attrs['footer'] = df.meta['footer']
        else:
            grp = hf[grp_name]

        for name in df.colnames:
            col = df[name]
            if isinstance(col, u.Quantity):
                data = col.value
            else:
                data = col.data
            if '<U' in data.dtype.str:
                numchar = data.dtype.str[2:]
                data = data.astype(f'|S{numchar}')
            if name in grp:
                ds = grp[name]
            else:
                if data.dtype == np.dtype('O'):
                    ragged_dtype = h5py.special_dtype(vlen=np.dtype('float64'))
                    ds = grp.create_dataset(name, data=data, dtype=ragged_dtype)
                else:
                    ds = grp.create_dataset(name, data=data, dtype=data.dtype)
            if col.unit is None:
                ds.attrs['unit'] = 'SKIP'
            else:
                ds.attrs['unit'] = col.unit.to_string()
            ds.attrs['description'] = df.meta['descriptions'][name]






class HSeqParser(GenericContinuumParser):
    r"""
    Parameters for calculating two-photon continuum for hydrogen-like ions

    Notes
    -----
    * The parameter :math: `\psi_{\text{norm}}` (called :math: `A_{\text{sum}}` in CHIANTI) is a
      normalization factor of the integral of the spectral distribution function :math:`\psi(y)` from 0 to 1,
      such that :math: `\frac{1}{\psi_{\text{norm}}} \int_{0}^{1} \psi(y) dy = 2`.
      This normalization is only used for hydrogenic ions.
    """
    filetype = 'hseq_2photon'
    dtypes = [int, float, int, float, float, float]
    units = [None, u.dimensionless_unscaled, None, 1/u.s, u.dimensionless_unscaled, u.dimensionless_unscaled]
    headings = ['Z', 'y', 'Z_0', 'A', 'psi_norm', 'psi']
    descriptions = ['atomic number', 'fraction of energy carried by one of the two photons',
                    'nominal atomic number', 'radiative decay rate', 'normalization of the integral of psi from 0 to 1',
                    'spectral distribution function']



    def preprocessor(self, table, line, index):
        if index == 0:
            self._y0 = np.array(line.strip().split(), dtype=float)
        elif index == 1:
            self._z0 = np.array(line.strip().split(), dtype=float)
        else:
            line = line.strip().split()
            table.append([line[0]] + [self._y0] + [self._z0] + line[1:3] + [np.array(line[3:])])




    def extract_footer(self, *args):
        return """Information needed for calculating two-photon continuum emission for hydrogen isoelectronic sequence
Radiative decay rates from Parpia, F. A., and Johnson, W. R., 1982, Phys. Rev. A, 26, 1142
Spectral distribution function from Goldman, S.P. and Drake, G.W.F., 1981, Phys Rev A, 24, 183"""




    def to_hdf5(self, hf, df, **kwargs):
        for row in df:
            el = plasmapy.particles.atomic_symbol(int(row['Z'])).lower()
            grp_name = f'{el}/continuum/{self.filetype}'
            if grp_name not in hf:
                grp = hf.create_group(grp_name)
                grp.attrs['chianti_version'] = df.meta['chianti_version']
                grp.attrs['footer'] = df.meta['footer']
            else:
                grp = hf[grp_name]
            for col in row.colnames:
                if col == 'Z':
                    continue
                ds = grp.create_dataset(col, data=row[col])
                ds.attrs['description'] = df.meta['descriptions'][col]
                if not hasattr(row[col], 'unit'):
                    ds.attrs['unit'] = 'SKIP'
                else:
                    ds.attrs['unit'] = row[col].unit.to_string()






class HeSeqParser(GenericContinuumParser):
    """
    Parameters for calculating two-photon continuum for helium-like ions.
    """
    filetype = 'heseq_2photon'
    dtypes = [int, float, float, float]
    units = [None, u.dimensionless_unscaled, 1/u.s, u.dimensionless_unscaled]
    headings = ['Z', 'y', 'A', 'psi']
    descriptions = ['atomic number', 'fraction of energy carried by one of the two photons',
                    'radiative decay rate', 'spectral distribution function']



    def preprocessor(self, table, line, index):
        if index == 0:
            self._y0 = np.array(line.strip().split(), dtype=float)
        else:
            line = line.strip().split()
            table.append([line[0]] + [self._y0] + line[1:2] + [np.array(line[2:])])




    def extract_footer(self, *args):
        return """Information needed for calculating two-photon continuum emission for helium isoelectronic sequence
Radiative decay rates from Drake, G.W.F., 1986, Phys. Rev. A, 34, 2871
Spectral distribution function from Drake, G.W.F., Victor, G.A., Dalgarno, A., 1969, Phys. Rev. A, 180, 25."""




    def to_hdf5(self, hf, df, **kwargs):
        for row in df:
            el = plasmapy.particles.atomic_symbol(int(row['Z'])).lower()
            grp_name = f'{el}/continuum/{self.filetype}'
            if grp_name not in hf:
                grp = hf.create_group(grp_name)
                grp.attrs['chianti_version'] = df.meta['chianti_version']
                grp.attrs['footer'] = df.meta['footer']
            else:
                grp = hf[grp_name]
            for col in row.colnames:
                if col == 'Z':
                    continue
                ds = grp.create_dataset(col, data=row[col])
                ds.attrs['description'] = df.meta['descriptions'][col]
                if not hasattr(row[col], 'unit'):
                    ds.attrs['unit'] = 'SKIP'
                else:
                    ds.attrs['unit'] = row[col].unit.to_string()