<reponame>henrikef/astromodels<gh_stars>1-10
import os
import sys
from functools import lru_cache, wraps

import astropy.units as astropy_units
import numpy as np
import six
from astropy.io import fits

from astromodels.functions.function import Function1D, FunctionMeta
from astromodels.utils import configuration
from astromodels.utils.data_files import _get_data_file_path


def cache_array_method(*args, **kwargs):
    """
    LRU cache implementation for methods whose FIRST parameter is a numpy array
    modified from: https://gist.github.com/Susensio/61f4fee01150caaac1e10fc5f005eb75
    """

    def decorator(function):
        @wraps(function)
        def wrapper(s, np_array, *args, **kwargs):
            hashable_array = tuple(np_array)
            return cached_wrapper(s, hashable_array, *args, **kwargs)

        @lru_cache(*args, **kwargs)
        def cached_wrapper(s, hashable_array, *args, **kwargs):
            array = np.array(hashable_array)
            return function(s, array, *args, **kwargs)

        # copy lru_cache attributes over too
        wrapper.cache_info = cached_wrapper.cache_info
        wrapper.cache_clear = cached_wrapper.cache_clear
        return wrapper

    return decorator


try:

    import pyatomdb

    has_atomdb = True

except:

    has_atomdb = False


if has_atomdb:
    # APEC class
    @six.add_metaclass(FunctionMeta)
    class APEC(Function1D):
        r"""
        description :
            The Astrophysical Plasma Emission Code (APEC, Smith et al. 2001)
            contributed by <NAME>
        parameters :
            K :
                desc : Normalization in units of 1e-14/(4*pi*(1+z)^2*dA*2)*EM
                initial value : 1.0
                is_normalization : True
                transformation : log10
                min : 1e-30
                max : 1e3
                delta : 0.1
            kT :
                desc : Plasma temperature
                initial value : 1.0
                min : 0.08
                max : 64
                delta : 0.1
            abund :
                desc : Metal abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            redshift :
                desc : Source redshift
                initial value : 0.1
                min : 0.0
                max : 10.0
                delta : 1e-3
                fix : yes

        """

        def _set_units(self, x_unit, y_unit):
            self.kT.unit = astropy_units.keV

            self.abund.unit = astropy_units.dimensionless_unscaled

            self.redshift.unit = astropy_units.dimensionless_unscaled

            self.K.unit = y_unit

        def init_session(self, abund_table="AG89"):
            # initialize PyAtomDB session
            self.session = pyatomdb.spectrum.CIESession(abundset=abund_table)

        def evaluate(self, x, K, kT, abund, redshift):
            assert self.session is not None, "please run init_session(abund)"

            sess = self.session

            nval = len(x)

            xz = x * (1.0 + redshift)

            ebplus = (np.roll(xz, -1) + xz)[: nval - 1] / 2.0

            ebounds = np.empty(nval + 1)

            ebounds[1:nval] = ebplus

            ebounds[0] = xz[0] - (ebplus[0] - xz[0])

            ebounds[nval] = xz[nval - 1] + (xz[nval - 1] - ebplus[nval - 2])

            binsize = (np.roll(ebounds, -1) - ebounds)[:nval]

            sess.set_response(ebounds, raw=True)

            sess.set_abund(
                [
                    6,
                    7,
                    8,
                    9,
                    10,
                    11,
                    12,
                    13,
                    14,
                    16,
                    17,
                    18,
                    19,
                    20,
                    21,
                    22,
                    23,
                    24,
                    25,
                    26,
                    27,
                    28,
                    29,
                    30,
                ],
                abund,
            )

            spec = sess.return_spectrum(kT) / binsize / 1e-14

            return K * spec

    # VAPEC class
    @six.add_metaclass(FunctionMeta)
    class VAPEC(Function1D):
        r"""
        description :
            The Astrophysical Plasma Emission Code (APEC, Smith et al. 2001), variable abundances for individual elements
            contributed by <NAME>
        parameters :
            K :
                desc : Normalization in units of 1e-14/(4*pi*(1+z)^2*dA*2)*EM
                initial value : 1.0
                is_normalization : True
                transformation : log10
                min : 1e-30
                max : 1e3
                delta : 0.1
            kT :
                desc : Plasma temperature
                initial value : 1.0
                min : 0.08
                max : 64
                delta : 0.1
            Fe :
                desc : Fe abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            C :
                desc : C abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            N :
                desc : N abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            O :
                desc : O abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ne :
                desc : Ne abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Mg :
                desc : Mg abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Al :
                desc : Al abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Si :
                desc : Si abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            S :
                desc : S abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ar :
                desc : Ar abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ca :
                desc : Ca abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ni :
                desc : Ni abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            redshift :
                desc : Source redshift
                initial value : 0.1
                min : 0.0
                max : 10.0
                delta : 1e-3
                fix : yes

        """

        def _set_units(self, x_unit, y_unit):
            self.kT.unit = astropy_units.keV

            self.Fe.unit = astropy_units.dimensionless_unscaled

            self.C.unit = astropy_units.dimensionless_unscaled

            self.N.unit = astropy_units.dimensionless_unscaled

            self.O.unit = astropy_units.dimensionless_unscaled

            self.Ne.unit = astropy_units.dimensionless_unscaled

            self.Mg.unit = astropy_units.dimensionless_unscaled

            self.Al.unit = astropy_units.dimensionless_unscaled

            self.Si.unit = astropy_units.dimensionless_unscaled

            self.Ar.unit = astropy_units.dimensionless_unscaled

            self.Ca.unit = astropy_units.dimensionless_unscaled

            self.Ni.unit = astropy_units.dimensionless_unscaled

            self.redshift.unit = astropy_units.dimensionless_unscaled

            self.K.unit = y_unit

        def init_session(self, abund_table="AG89"):
            # initialize PyAtomDB session
            self.session = pyatomdb.spectrum.CIESession(abundset=abund_table)

        def evaluate(
            self, x, K, kT, Fe, C, N, O, Ne, Mg, Al, Si, S, Ar, Ca, Ni, redshift
        ):
            assert self.session is not None, "please run init_session(abund)"

            sess = self.session

            nval = len(x)

            xz = x * (1.0 + redshift)

            ebplus = (np.roll(xz, -1) + xz)[: nval - 1] / 2.0

            ebounds = np.empty(nval + 1)

            ebounds[1:nval] = ebplus

            ebounds[0] = xz[0] - (ebplus[0] - xz[0])

            ebounds[nval] = xz[nval - 1] + (xz[nval - 1] - ebplus[nval - 2])

            binsize = (np.roll(ebounds, -1) - ebounds)[:nval]

            sess.set_response(ebounds, raw=True)

            sess.set_abund(
                [6, ], C,
            )

            sess.set_abund(
                [7, ], N,
            )

            sess.set_abund(
                [8, ], O,
            )

            sess.set_abund(
                [10, ], Ne,
            )

            sess.set_abund(
                [12, ], Mg,
            )

            sess.set_abund(
                [13, ], Al,
            )

            sess.set_abund(
                [14, ], Si,
            )

            sess.set_abund(
                [16, ], S,
            )

            sess.set_abund(
                [18, ], Ar,
            )

            sess.set_abund(
                [20, ], Ca,
            )

            sess.set_abund(
                [26, ], Fe,
            )

            sess.set_abund(
                [28, ], Ni,
            )

            sess.set_abund(
                [9, 11, 15, 17, 19, 21, 22, 23, 24, 25, 27, 29, 30], Fe
            )  # Remaining elements are set to Fe

            spec = sess.return_spectrum(kT) / binsize / 1e-14

            return K * spec


_abs_tables = {
    "phabs": {"AG89": "angr", "ASPL": "aspl"},
    "tbabs": {"AG89": "angr", "ASPL": "aspl", "WILM": "wilm"},
    "wabs": {"AG89": "angr"},
}
_abund_info = {}
_abund_info[
    "WILM"
] = "wilms\nfrom Wilms, <NAME> (2000), ApJ 542, 914 \n except for elements not listed which are given zero abundance)\n https://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html "
_abund_info[
    "AG89"
] = "angr\nfrom <NAME>. & <NAME>. (1989, Geochimica et Cosmochimica Acta 53, 197)\n https://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html"
_abund_info[
    "ASPL"
] = "aspl\nfrom <NAME>., <NAME>., <NAME>. & <NAME>. (2009, ARAA, 47, 481)\nhttps://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html"


def _get_xsect_table(model, abund_table):
    """
    contructs the abundance table from the values given
    """

    assert model in _abs_tables, "the model %s does not exist" % model
    assert abund_table in _abs_tables[model], (
        "the table %s does not exist" % abund_table
    )

    path_to_xsect = _get_data_file_path(
        os.path.join(
            "xsect", "xsect_%s_%s.fits" % (
                model, _abs_tables[model][abund_table])
        )
    )

    fxs = fits.open(path_to_xsect)
    dxs = fxs[1].data
    xsect_ene = dxs["ENERGY"]
    xsect_val = dxs["SIGMA"]

    return np.array(xsect_ene), np.array(xsect_val)

# PhAbs class


@six.add_metaclass(FunctionMeta)
class PhAbs(Function1D):
    r"""
    description :
        Photometric absorption (phabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by <NAME>
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : False
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1

        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True


    """

    def _setup(self):
        self._fixed_units = (
            astropy_units.keV, astropy_units.dimensionless_unscaled)
        self.init_xsect()

    def _set_units(self, x_unit, y_unit):
        self.NH.unit = astropy_units.cm ** (-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self, abund_table="AG89"):
        """
        Set the abundance table

        :param abund_table: "ASPL", "AG89" 
        :returns: 
        :rtype: 

        """

        # load cross section data

        try:
            self.xsect_ene, self.xsect_val = _get_xsect_table(
                "phabs", abund_table)
            self._abund_table = abund_table

        except:

            print("defaulting to AG89")
            self.xsect_ene, self.xsect_val = _get_xsect_table(
                "phabs", abund_table)

            self._abund_table = "AG89"

    @cache_array_method()
    def _cached_interp(self, x):

        return np.interp(x, self.xsect_ene, self.xsect_val)

    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm ** 2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value

        else:

            _unit = 1.0
            _y_unit = 1.0

            _x = x

        xsect_interp = self._cached_interp(_x * (1 + redshift))

        spec = np.exp(-NH * xsect_interp * _unit) * _y_unit

        return spec

# TbAbs class


@six.add_metaclass(FunctionMeta)
class TbAbs(Function1D):
    r"""
    description :
        Photometric absorption (Tbabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by <NAME>
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : True
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1

        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True


    """

    def _setup(self):

        self.init_xsect()

        self._fixed_units = (
            astropy_units.keV, astropy_units.dimensionless_unscaled)

    def _set_units(self, x_unit, y_unit):
        self.NH.unit = astropy_units.cm ** (-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self, abund_table="WILM"):
        """
        Set the abundance table

        :param abund_table: "WILM", "ASPL", "AG89" 
        :returns: 
        :rtype: 

        """

        try:
            self.xsect_ene, self.xsect_val = _get_xsect_table(
                "tbabs", abund_table)
            self._abund_table = abund_table

        except:

            print("defaulting to WILM")
            self.xsect_ene, self.xsect_val = _get_xsect_table(
                "tbabs", abund_table)

            self._abund_table = "WILM"

    @property
    def abundance_table(self):
        print(_abund_info[self._abund_table])

    @cache_array_method()
    def _cached_interp(self, x):

        return np.interp(x, self.xsect_ene, self.xsect_val)

    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm ** 2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value

        else:

            _unit = 1.0
            _y_unit = 1.0

            _x = x

        xsect_interp = self._cached_interp(_x * (1 + redshift))

        spec = np.exp(-NH * xsect_interp * _unit) * _y_unit

        return spec

# WAbs class


@six.add_metaclass(FunctionMeta)
class WAbs(Function1D):
    r"""
    description :
        Photometric absorption (Wabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by <NAME>
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : True
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1
        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True


    """

    def _setup(self):
        self._fixed_units = (
            astropy_units.keV, astropy_units.dimensionless_unscaled)
        self.init_xsect()

    def _set_units(self, x_unit, y_unit):
        self.NH.unit = astropy_units.cm ** (-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self):
        """
        Set the abundance table

        :returns:
        :rtype:

        """

        self.xsect_ene, self.xsect_val = _get_xsect_table("wabs", "AG89")

        self._abund_table = "AG89"

    @property
    def abundance_table(self):
        print(_abund_info[self._abund_table])

    @cache_array_method()
    def _cached_interp(self, x):

        return np.interp(x, self.xsect_ene, self.xsect_val)

    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm ** 2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value

        else:

            _unit = 1.0
            _y_unit = 1.0

            _x = x

        xsect_interp = self._cached_interp(_x * (1 + redshift))

        spec = np.exp(-NH * xsect_interp * _unit) * _y_unit

        return spec

import os
import sys
from functools import lru_cache, wraps

import astropy.units as astropy_units
import numpy as np
import six
from astropy.io import fits

from astromodels.functions.function import Function1D, FunctionMeta
from astromodels.utils import configuration
from astromodels.utils.data_files import _get_data_file_path


def cache_array_method(*args, **kwargs):
    """
    LRU cache implementation for methods whose FIRST parameter is a numpy array
    modified from: https://gist.github.com/Susensio/61f4fee01150caaac1e10fc5f005eb75
    """

    def decorator(function):
        @wraps(function)
        def wrapper(s, np_array, *args, **kwargs):
            hashable_array = tuple(np_array)
            return cached_wrapper(s, hashable_array, *args, **kwargs)

        @lru_cache(*args, **kwargs)
        def cached_wrapper(s, hashable_array, *args, **kwargs):
            array = np.array(hashable_array)
            return function(s, array, *args, **kwargs)

        # copy lru_cache attributes over too
        wrapper.cache_info = cached_wrapper.cache_info
        wrapper.cache_clear = cached_wrapper.cache_clear
        return wrapper

    return decorator


try:

    import pyatomdb

    has_atomdb = True

except:

    has_atomdb = False


if has_atomdb:
    # APEC class
    @six.add_metaclass(FunctionMeta)
    class APEC(Function1D):
        r"""
        description :
            The Astrophysical Plasma Emission Code (APEC, Smith et al. 2001)
            contributed by <NAME>
        parameters :
            K :
                desc : Normalization in units of 1e-14/(4*pi*(1+z)^2*dA*2)*EM
                initial value : 1.0
                is_normalization : True
                transformation : log10
                min : 1e-30
                max : 1e3
                delta : 0.1
            kT :
                desc : Plasma temperature
                initial value : 1.0
                min : 0.08
                max : 64
                delta : 0.1
            abund :
                desc : Metal abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            redshift :
                desc : Source redshift
                initial value : 0.1
                min : 0.0
                max : 10.0
                delta : 1e-3
                fix : yes

        """

        def _set_units(self, x_unit, y_unit):
            self.kT.unit = astropy_units.keV

            self.abund.unit = astropy_units.dimensionless_unscaled

            self.redshift.unit = astropy_units.dimensionless_unscaled

            self.K.unit = y_unit

        def init_session(self, abund_table="AG89"):
            # initialize PyAtomDB session
            self.session = pyatomdb.spectrum.CIESession(abundset=abund_table)

        def evaluate(self, x, K, kT, abund, redshift):
            assert self.session is not None, "please run init_session(abund)"

            sess = self.session

            nval = len(x)

            xz = x * (1.0 + redshift)

            ebplus = (np.roll(xz, -1) + xz)[: nval - 1] / 2.0

            ebounds = np.empty(nval + 1)

            ebounds[1:nval] = ebplus

            ebounds[0] = xz[0] - (ebplus[0] - xz[0])

            ebounds[nval] = xz[nval - 1] + (xz[nval - 1] - ebplus[nval - 2])

            binsize = (np.roll(ebounds, -1) - ebounds)[:nval]

            sess.set_response(ebounds, raw=True)

            sess.set_abund(
                [
                    6,
                    7,
                    8,
                    9,
                    10,
                    11,
                    12,
                    13,
                    14,
                    16,
                    17,
                    18,
                    19,
                    20,
                    21,
                    22,
                    23,
                    24,
                    25,
                    26,
                    27,
                    28,
                    29,
                    30,
                ],
                abund,
            )

            spec = sess.return_spectrum(kT) / binsize / 1e-14

            return K * spec

    # VAPEC class
    @six.add_metaclass(FunctionMeta)
    class VAPEC(Function1D):
        r"""
        description :
            The Astrophysical Plasma Emission Code (APEC, Smith et al. 2001), variable abundances for individual elements
            contributed by <NAME>
        parameters :
            K :
                desc : Normalization in units of 1e-14/(4*pi*(1+z)^2*dA*2)*EM
                initial value : 1.0
                is_normalization : True
                transformation : log10
                min : 1e-30
                max : 1e3
                delta : 0.1
            kT :
                desc : Plasma temperature
                initial value : 1.0
                min : 0.08
                max : 64
                delta : 0.1
            Fe :
                desc : Fe abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            C :
                desc : C abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            N :
                desc : N abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            O :
                desc : O abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ne :
                desc : Ne abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Mg :
                desc : Mg abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Al :
                desc : Al abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Si :
                desc : Si abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            S :
                desc : S abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ar :
                desc : Ar abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ca :
                desc : Ca abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            Ni :
                desc : Ni abundance
                initial value : 1
                min : 0.0
                max : 5.0
                delta : 0.01
                fix : yes
            redshift :
                desc : Source redshift
                initial value : 0.1
                min : 0.0
                max : 10.0
                delta : 1e-3
                fix : yes

        """

        def _set_units(self, x_unit, y_unit):
            self.kT.unit = astropy_units.keV

            self.Fe.unit = astropy_units.dimensionless_unscaled

            self.C.unit = astropy_units.dimensionless_unscaled

            self.N.unit = astropy_units.dimensionless_unscaled

            self.O.unit = astropy_units.dimensionless_unscaled

            self.Ne.unit = astropy_units.dimensionless_unscaled

            self.Mg.unit = astropy_units.dimensionless_unscaled

            self.Al.unit = astropy_units.dimensionless_unscaled

            self.Si.unit = astropy_units.dimensionless_unscaled

            self.Ar.unit = astropy_units.dimensionless_unscaled

            self.Ca.unit = astropy_units.dimensionless_unscaled

            self.Ni.unit = astropy_units.dimensionless_unscaled

            self.redshift.unit = astropy_units.dimensionless_unscaled

            self.K.unit = y_unit

        def init_session(self, abund_table="AG89"):
            # initialize PyAtomDB session
            self.session = pyatomdb.spectrum.CIESession(abundset=abund_table)

        def evaluate(
            self, x, K, kT, Fe, C, N, O, Ne, Mg, Al, Si, S, Ar, Ca, Ni, redshift
        ):
            assert self.session is not None, "please run init_session(abund)"

            sess = self.session

            nval = len(x)

            xz = x * (1.0 + redshift)

            ebplus = (np.roll(xz, -1) + xz)[: nval - 1] / 2.0

            ebounds = np.empty(nval + 1)

            ebounds[1:nval] = ebplus

            ebounds[0] = xz[0] - (ebplus[0] - xz[0])

            ebounds[nval] = xz[nval - 1] + (xz[nval - 1] - ebplus[nval - 2])

            binsize = (np.roll(ebounds, -1) - ebounds)[:nval]

            sess.set_response(ebounds, raw=True)

            sess.set_abund(
                [6, ], C,
            )

            sess.set_abund(
                [7, ], N,
            )

            sess.set_abund(
                [8, ], O,
            )

            sess.set_abund(
                [10, ], Ne,
            )

            sess.set_abund(
                [12, ], Mg,
            )

            sess.set_abund(
                [13, ], Al,
            )

            sess.set_abund(
                [14, ], Si,
            )

            sess.set_abund(
                [16, ], S,
            )

            sess.set_abund(
                [18, ], Ar,
            )

            sess.set_abund(
                [20, ], Ca,
            )

            sess.set_abund(
                [26, ], Fe,
            )

            sess.set_abund(
                [28, ], Ni,
            )

            sess.set_abund(
                [9, 11, 15, 17, 19, 21, 22, 23, 24, 25, 27, 29, 30], Fe
            )  # Remaining elements are set to Fe

            spec = sess.return_spectrum(kT) / binsize / 1e-14

            return K * spec


_abs_tables = {
    "phabs": {"AG89": "angr", "ASPL": "aspl"},
    "tbabs": {"AG89": "angr", "ASPL": "aspl", "WILM": "wilm"},
    "wabs": {"AG89": "angr"},
}
_abund_info = {}
_abund_info[
    "WILM"
] = "wilms\nfrom Wilms, <NAME> (2000), ApJ 542, 914 \n except for elements not listed which are given zero abundance)\n https://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html "
_abund_info[
    "AG89"
] = "angr\nfrom <NAME>. & <NAME>. (1989, Geochimica et Cosmochimica Acta 53, 197)\n https://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html"
_abund_info[
    "ASPL"
] = "aspl\nfrom <NAME>., <NAME>., <NAME>. & <NAME>. (2009, ARAA, 47, 481)\nhttps://heasarc.nasa.gov/xanadu/xspec/manual/XSabund.html"


def _get_xsect_table(model, abund_table):
    """
    contructs the abundance table from the values given
    """

    assert model in _abs_tables, "the model %s does not exist" % model
    assert abund_table in _abs_tables[model], (
        "the table %s does not exist" % abund_table
    )

    path_to_xsect = _get_data_file_path(
        os.path.join(
            "xsect", "xsect_%s_%s.fits" % (
                model, _abs_tables[model][abund_table])
        )
    )

    fxs = fits.open(path_to_xsect)
    dxs = fxs[1].data
    xsect_ene = dxs["ENERGY"]
    xsect_val = dxs["SIGMA"]

    return np.array(xsect_ene), np.array(xsect_val)

# PhAbs class


@six.add_metaclass(FunctionMeta)
class PhAbs(Function1D):
    r"""
    description :
        Photometric absorption (phabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by <NAME>
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : False
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1

        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True


    """

    def _setup(self):
        self._fixed_units = (
            astropy_units.keV, astropy_units.dimensionless_unscaled)
        self.init_xsect()

    def _set_units(self, x_unit, y_unit):
        self.NH.unit = astropy_units.cm ** (-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self, abund_table="AG89"):
        """
        Set the abundance table

        :param abund_table: "ASPL", "AG89" 
        :returns: 
        :rtype: 

        """

        # load cross section data

        try:
            self.xsect_ene, self.xsect_val = _get_xsect_table(
                "phabs", abund_table)
            self._abund_table = abund_table

        except:

            print("defaulting to AG89")
            self.xsect_ene, self.xsect_val = _get_xsect_table(
                "phabs", abund_table)

            self._abund_table = "AG89"

    @cache_array_method()
    def _cached_interp(self, x):

        return np.interp(x, self.xsect_ene, self.xsect_val)

    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm ** 2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value

        else:

            _unit = 1.0
            _y_unit = 1.0

            _x = x

        xsect_interp = self._cached_interp(_x * (1 + redshift))

        spec = np.exp(-NH * xsect_interp * _unit) * _y_unit

        return spec

# TbAbs class


@six.add_metaclass(FunctionMeta)
class TbAbs(Function1D):
    r"""
    description :
        Photometric absorption (Tbabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by <NAME>
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : True
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1

        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True


    """

    def _setup(self):

        self.init_xsect()

        self._fixed_units = (
            astropy_units.keV, astropy_units.dimensionless_unscaled)

    def _set_units(self, x_unit, y_unit):
        self.NH.unit = astropy_units.cm ** (-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self, abund_table="WILM"):
        """
        Set the abundance table

        :param abund_table: "WILM", "ASPL", "AG89" 
        :returns: 
        :rtype: 

        """

        try:
            self.xsect_ene, self.xsect_val = _get_xsect_table(
                "tbabs", abund_table)
            self._abund_table = abund_table

        except:

            print("defaulting to WILM")
            self.xsect_ene, self.xsect_val = _get_xsect_table(
                "tbabs", abund_table)

            self._abund_table = "WILM"

    @property
    def abundance_table(self):
        print(_abund_info[self._abund_table])

    @cache_array_method()
    def _cached_interp(self, x):

        return np.interp(x, self.xsect_ene, self.xsect_val)

    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm ** 2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value

        else:

            _unit = 1.0
            _y_unit = 1.0

            _x = x

        xsect_interp = self._cached_interp(_x * (1 + redshift))

        spec = np.exp(-NH * xsect_interp * _unit) * _y_unit

        return spec

# WAbs class


@six.add_metaclass(FunctionMeta)
class WAbs(Function1D):
    r"""
    description :
        Photometric absorption (Wabs implementation), f(E) = exp(- NH * sigma(E))
        contributed by <NAME>
    parameters :
        NH :
            desc : absorbing column density in units of 1e22 particles per cm^2
            initial value : 1.0
            is_normalization : True
            transformation : log10
            min : 1e-4
            max : 1e4
            delta : 0.1
        redshift :
            desc : the redshift of the source
            initial value : 0.
            is_normalization : False
            min : 0
            max : 15
            delta : 0.1
            fix: True


    """

    def _setup(self):
        self._fixed_units = (
            astropy_units.keV, astropy_units.dimensionless_unscaled)
        self.init_xsect()

    def _set_units(self, x_unit, y_unit):
        self.NH.unit = astropy_units.cm ** (-2)
        self.redshift.unit = astropy_units.dimensionless_unscaled

    def init_xsect(self):
        """
        Set the abundance table

        :returns:
        :rtype:

        """

        self.xsect_ene, self.xsect_val = _get_xsect_table("wabs", "AG89")

        self._abund_table = "AG89"

    @property
    def abundance_table(self):
        print(_abund_info[self._abund_table])

    @cache_array_method()
    def _cached_interp(self, x):

        return np.interp(x, self.xsect_ene, self.xsect_val)

    def evaluate(self, x, NH, redshift):

        if isinstance(x, astropy_units.Quantity):

            _unit = astropy_units.cm ** 2
            _y_unit = astropy_units.dimensionless_unscaled
            _x = x.value

        else:

            _unit = 1.0
            _y_unit = 1.0

            _x = x

        xsect_interp = self._cached_interp(_x * (1 + redshift))

        spec = np.exp(-NH * xsect_interp * _unit) * _y_unit

        return spec
LRU cache implementation for methods whose FIRST parameter is a numpy array     modified from: https://gist.github.com/Susensio/61f4fee01150caaac1e10fc5f005eb75 # copy lru_cache attributes over too # APEC class description :             The Astrophysical Plasma Emission Code (APEC, Smith et al. 2001)             contributed by <NAME>         parameters :             K :                 desc : Normalization in units of 1e-14/(4*pi*(1+z)^2*dA*2)*EM                 initial value : 1.0                 is_normalization : True                 transformation : log10                 min : 1e-30                 max : 1e3                 delta : 0.1             kT :                 desc : Plasma temperature                 initial value : 1.0                 min : 0.08                 max : 64                 delta : 0.1             abund :                 desc : Metal abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             redshift :                 desc : Source redshift                 initial value : 0.1                 min : 0.0                 max : 10.0                 delta : 1e-3                 fix : yes # initialize PyAtomDB session # VAPEC class description :             The Astrophysical Plasma Emission Code (APEC, Smith et al. 2001), variable abundances for individual elements             contributed by <NAME>         parameters :             K :                 desc : Normalization in units of 1e-14/(4*pi*(1+z)^2*dA*2)*EM                 initial value : 1.0                 is_normalization : True                 transformation : log10                 min : 1e-30                 max : 1e3                 delta : 0.1             kT :                 desc : Plasma temperature                 initial value : 1.0                 min : 0.08                 max : 64                 delta : 0.1             Fe :                 desc : Fe abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             C :                 desc : C abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             N :                 desc : N abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             O :                 desc : O abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             Ne :                 desc : Ne abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             Mg :                 desc : Mg abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             Al :                 desc : Al abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             Si :                 desc : Si abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             S :                 desc : S abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             Ar :                 desc : Ar abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             Ca :                 desc : Ca abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             Ni :                 desc : Ni abundance                 initial value : 1                 min : 0.0                 max : 5.0                 delta : 0.01                 fix : yes             redshift :                 desc : Source redshift                 initial value : 0.1                 min : 0.0                 max : 10.0                 delta : 1e-3                 fix : yes # initialize PyAtomDB session # Remaining elements are set to Fe contructs the abundance table from the values given # PhAbs class description :         Photometric absorption (phabs implementation), f(E) = exp(- NH * sigma(E))         contributed by <NAME>     parameters :         NH :             desc : absorbing column density in units of 1e22 particles per cm^2             initial value : 1.0             is_normalization : False             transformation : log10             min : 1e-4             max : 1e4             delta : 0.1          redshift :             desc : the redshift of the source             initial value : 0.             is_normalization : False             min : 0             max : 15             delta : 0.1             fix: True Set the abundance table          :param abund_table: "ASPL", "AG89"          :returns:          :rtype: # load cross section data # TbAbs class description :         Photometric absorption (Tbabs implementation), f(E) = exp(- NH * sigma(E))         contributed by <NAME>     parameters :         NH :             desc : absorbing column density in units of 1e22 particles per cm^2             initial value : 1.0             is_normalization : True             transformation : log10             min : 1e-4             max : 1e4             delta : 0.1          redshift :             desc : the redshift of the source             initial value : 0.             is_normalization : False             min : 0             max : 15             delta : 0.1             fix: True Set the abundance table          :param abund_table: "WILM", "ASPL", "AG89"          :returns:          :rtype: # WAbs class description :         Photometric absorption (Wabs implementation), f(E) = exp(- NH * sigma(E))         contributed by <NAME>     parameters :         NH :             desc : absorbing column density in units of 1e22 particles per cm^2             initial value : 1.0             is_normalization : True             transformation : log10             min : 1e-4             max : 1e4             delta : 0.1         redshift :             desc : the redshift of the source             initial value : 0.             is_normalization : False             min : 0             max : 15             delta : 0.1             fix: True Set the abundance table          :returns:         :rtype: