Source code for irfpy.mima.gfactor0

''' Geometric factors for MEX/IMA.

.. note::

    Among a lot of definitions of geometric factor, the current best knowledge of MEX/IMA at IRF is
    to use :func:`gfactor_func`.  This is used for the official data archiving at PSA.

It is not straightforward, both theoretically and practically.

Here are some implementations, but carefully assess and use when you use them.

The current best proposal of the geometric factor for MEX IMA is only to use :func:`gfactor_func`.
'''
import numpy as np
from irfpy.mima import energy as _energy


class GL_CR4:
    r''' Geometric factor defined in MEX-ASP-CR-0004.

    P22 gives us impression on some values of g-factor.

    This is :math:`G_0/\Delta\Phi` [cm2 rad eV/eV].

    I am not very sure what the absolute value indicate here,
    but the relative value should at least be meaningful.

    >>> gl_h0 = GL_CR4.Proton.PAC0()
    >>> print(gl_h0([0, 2000, 5000, 10000, 30000]))    # doctest: +NORMALIZE_WHITESPACE
    [0.00000000e+00 1.55148151e-07 4.89672821e-06 4.48326188e-05
     4.48326188e-05]

    '''
    class Proton:

        class PAC0:

            def __init__(self):
                self.mid = lambda e: 5.6856e-20 * np.power(e, 3.7673)
                self.mid0 = 2000.
                self.mid1 = 9000.

            def __call__(self, energy):
                ''' Get approximated gfactor for PACC0.
                '''
                earray = np.array(energy)
                garray = np.zeros_like(earray)

                garray = np.where(earray < self.mid0, 0, garray)
                garray = np.where(np.logical_and(earray >= self.mid0, earray < self.mid1),
                                  self.mid(earray), garray)
                garray = np.where(earray >= self.mid1, self.mid(self.mid1) * np.ones_like(garray), garray)

                return garray

        class PAC1:
            def __init__(self):
                self.mid0 = 800.
                self.mid1 = 9000.
                self.mid = lambda e: 7.75129e-13 * np.power(e, 1.9619)

            def __call__(self, energy):
                earray = np.array(energy)
                garray = np.zeros_like(earray)

                garray = np.where(earray < self.mid0, 0, garray)
                garray = np.where(np.logical_and(earray >= self.mid0, earray < self.mid1),
                                  self.mid(earray), garray)
                garray = np.where(earray >= self.mid1, self.mid(self.mid1) * np.ones_like(garray), garray)

                return garray

        class PAC2:
            def __init__(self):
                self.mid0 = 100.
                self.mid1 = 5665.
                self.mid2 = 9000.

                self.mid01 = lambda e: 1.5975e-8 * np.power(e, 0.8747146)
                self.mid12 = lambda e: 3.3118e-17 * np.power(e, 3.18833)

            def __call__(self, energy):
                earray = np.array(energy)
                garray = np.zeros_like(earray)

                garray = np.where(earray < self.mid0, 0, garray)
                garray = np.where(np.logical_and(earray >= self.mid0, earray < self.mid1),
                                  self.mid01(earray), garray)
                garray = np.where(np.logical_and(earray >= self.mid1, earray < self.mid2),
                                  self.mid12(earray), garray)
                garray = np.where(earray >= self.mid2, self.mid12(self.mid2) * np.ones_like(garray), garray)

                return garray

    class Alpha:
        class PAC0:
            def __init__(self):
                e0 = 4343.47288876
                g0 = 7.11715324658e-06
                e1 = 9521.30113195
                g1 = 0.000100127259076

                # Follow  g = k0 * E ^ k1
                k1 = np.log(g1 / g0) / np.log(e1 / e0)
                k0 = g0 / np.power(e0, k1)

                self.func0 = lambda e: k0 * np.power(e, k1)
                self.e0 = 2000.
                self.e1 = 9000.

            def __call__(self, energy):
                earray = np.array(energy)
                garray = np.zeros_like(earray)

                garray = np.where(earray < self.e0, 0, garray)
                garray = np.where(np.logical_and(earray >= self.e0, earray < self.e1),
                                  self.func0(earray), garray)
                garray = np.where(earray >= self.e1, self.func0(self.e1) * np.ones_like(garray), garray)

                return garray

        class PAC1:
            def __init__(self):
                e0 = 591.676042159
                g0 = 1.87182187345e-06
                e1 = 8639.4700344
                g1 = 4.1855209508e-05

                # Follow  g = k0 * E ^ k1
                k1 = np.log(g1 / g0) / np.log(e1 / e0)
                k0 = g0 / np.power(e0, k1)

                self.func0 = lambda e: k0 * np.power(e, k1)
                self.e0 = 800.
                self.e1 = 9000.

            def __call__(self, energy):
                earray = np.array(energy)
                garray = np.zeros_like(earray)

                garray = np.where(earray < self.e0, 0, garray)
                garray = np.where(np.logical_and(earray >= self.e0, earray < self.e1),
                                  self.func0(earray), garray)
                garray = np.where(earray >= self.e1, self.func0(self.e1) * np.ones_like(garray), garray)

                return garray

        class PAC2:
            def __init__(self):
                e0, g0, e1, g1, e2, g2 = [float(val) for val in
                                          '''275.384500319  4.03092208896e-06
                                             612.651125795  1.30086655237e-05
                                             9288.46528267  5.44660332838e-05'''.split()]

                # Fisst func. Follow  g = k0 * E ^ k1
                k1 = np.log(g1 / g0) / np.log(e1 / e0)
                k0 = g0 / np.power(e0, k1)

                self.func0 = lambda e: k0 * np.power(e, k1)

                k3 = np.log(g2 / g1) / np.log(e2 / e1)
                k2 = g1 / np.power(e1, k3)

                self.func1 = lambda e: k2 * np.power(e, k3)

                self.e0 = 100
                self.e1 = np.exp(np.log(k2 / k0) / (k1 - k3))
                self.e2 = 9000

            def __call__(self, energy):
                earray = np.array(energy)
                garray = np.zeros_like(earray)

                garray = np.where(earray < self.e0, 0, garray)
                garray = np.where(np.logical_and(earray >= self.e0, earray < self.e1),
                                  self.func0(earray), garray)
                garray = np.where(np.logical_and(earray >= self.e1, earray < self.e2),
                                  self.func1(earray), garray)
                garray = np.where(earray >= self.e2, self.func1(self.e2) * np.ones_like(garray), garray)

                return garray


def _logE(energy):
    return np.log(energy)


def _indexbelow(energy, e0):
    return np.where(np.array(energy) < e0)


def _indexabove(energy, e0):
    return np.where(np.array(energy) > e0)


def _indexbetween(energy, e0, e1):
    earray = np.array(energy)
    return np.where(np.logical_and(earray >= e0, earray < e1))


class GL_extra:
    """ Geometric factor (GL) in IMA extra.

    Simplest way is to use the :func:`get` function.

    >>> gl = GL_extra.get('Proton', 0, 4)    # PROM section 0 and PAC4 for proton.
    >>> print(gl.shape)
    (96,)
    >>> print(gl[0], gl[10], gl[80], gl[95])
    5.218046635641179e-06 8.392430885762372e-06 0.0 0.0

    You can also get GL for arbitrary energy levels.

    >>> gl_h0 = GL_extra.Proton.PAC0()
    >>> print(gl_h0([0, 2000, 5000, 10000, 30000]))  # doctest: +NORMALIZE_WHITESPACE
    [             inf   6.59448303e-06   4.66195441e-06   3.58619490e-06
       2.36620604e-06]

    NOTE that the PACC0 is not well calibrated but not frequently used :)
    """
    _pacc = {0: 0, 4: 1, 7: 2}

    @classmethod
    def get(cls, species, eeprom, pacc):
        """ Return the GL factor for the given EEPROM channel and PACC.

        :param species: Either of 'Proton', 'Ghost', 'Alpha', or 'Heavy'
        :type species: ``str``
        :param eeprom: PROM/EEPROM section. Integer from 0 to 16.
        :param pacc: PACC index. Integer from 0 to 6. (Only 0, 4, 7 supported)
        :return: Table of gl factor, with a unit of cm2 rad eV/eV
        """
        etbl = _energy.get_default_table(eeprom, keep_negative=True)

        try:
            p = cls._pacc[pacc]
        except IndexError:
            raise ValueError('Given PACC (={}) not supported.'.format(pacc))

        if species.lower() in ['proton', 'hydrogen', 'h', 'p', 'hp']:
            if p == 0:
                return cls.Proton.PAC0()(etbl)
            elif p == 1:
                return cls.Proton.PAC1()(etbl)
            elif p == 2:
                return cls.Proton.PAC2()(etbl)
        elif species.lower() in ['ghost', 'protonghost', 'g', 'hg', 'pg']:
            if p == 0:
                return cls.ProtonGhost.PAC0()(etbl)
            elif p == 1:
                return cls.ProtonGhost.PAC1()(etbl)
            elif p == 2:
                return cls.ProtonGhost.PAC2()(etbl)
        elif species.lower() in ['alpha', 'helium', 'he', 'hepp', 'he2p', 'al']:
            if p == 0:
                return cls.Alpha.PAC0()(etbl)
            elif p == 1:
                return cls.Alpha.PAC1()(etbl)
            elif p == 2:
                return cls.Alpha.PAC2()(etbl)
        elif species.lower() in ['heavy', 'oxygen', 'o', 'op']:
            if p == 0:
                return cls.Heavy.PAC0()(etbl)
            elif p == 1:
                return cls.Heavy.PAC1()(etbl)
            elif p == 2:
                return cls.Heavy.PAC2()(etbl)

        raise ValueError("Some quantities are not supported. GIVEN={} {} {}".format(species, eeprom, pacc))

    class Heavy:

        class PAC0:
            def __init__(self):
                pass

            def __call__(self, energy):
                return np.exp(-0.378481 * _logE(energy) - 9.05248)

        class PAC1:
            def __init__(self):
                pass

            def __call__(self, energy):
                logE = _logE(energy)
                Gl = np.exp(-0.69 * logE - 5.63)
                e300index = _indexbelow(energy, 300)
                Gl[e300index] = 7e-5
                return Gl

        class PAC2:
            def __init__(self):
                pass

            def __call__(self, energy):
                logE = _logE(energy)
                e600index = _indexbelow(energy, 600)
                Gl = np.exp(-0.611844 * logE - 5.31592)
                Gl[e600index] = 1.0e-4
                return Gl

    class Alpha:

        class PAC0:
            def __init__(self):
                pass

            def __call__(self, energy):
                return np.exp(-0.378481 * _logE(energy) - 9.05248)

        class PAC1:
            def __init__(self):
                pass

            def __call__(self, energy):
                logE = _logE(energy)
                Gl = np.exp(-0.554935*logE - 6.40179)
                e300index = _indexbelow(energy, 300)

                Gl[e300index] = 7e-5
                return Gl

        class PAC2:
            def __init__(self):
                pass

            def __call__(self, energy):
                logE = _logE(energy)
                Gl = np.exp(-0.611844 * logE - 5.31592)
                e600index = _indexbelow(energy, 600)
                Gl[e600index] = 1.0e-4
                return Gl

    class Proton:

        class PAC0:
            def __init__(self):
                pass

            def __call__(self, energy):
                return np.exp(-0.378481 * _logE(energy) - 9.05248)

        class PAC1:
            def __init__(self):
                pass

            def __call__(self, energy):
                earray = np.array(energy)
                logE = _logE(earray)
                Gl = earray * 0.0 - 9876543210
                e700index = _indexbetween(earray, 700, 2000)
                e2000index = _indexabove(earray, 2000)
                Gl[e700index] = np.exp(3.0440526 * logE[e700index] - 33.7573)
                Gl[e2000index] = np.exp(-0.554935 * logE[e2000index] - 6.40179)
                Gl = np.where(Gl >= 0, Gl, 0)
                return Gl

        class PAC2:
            def __init__(self):
                pass

            def __call__(self, energy):
                earray = np.array(energy)
                Gl = earray * 0.0 - 9876543210
                e1500index = _indexabove(energy, 1500)
                logE = _logE(energy)
                Gl[e1500index] = np.exp(-0.611844 * logE[e1500index] - 5.31592)
                Gl = np.where(Gl >= 0, Gl, 0)
                return Gl

    class ProtonGhost:
        class PAC0:
            def __init__(self):
                pass

            def __call__(self, energy):
                return np.exp(-0.378481 * _logE(energy) - 9.05248)

        class PAC1:
            def __init__(self):
                pass

            def __call__(self, energy):
                earray = np.array(energy)
                Gl = earray * 0.0 - 9876543210
                e800index = _indexbetween(earray, 800, 2500)
                e100index = _indexbetween(earray, 100, 800)
                logE = _logE(earray)

                Gl[e800index] = np.exp(-1.2*logE[e800index] - 5.1)
                Gl[e100index] = np.exp(1.44064*logE[e100index] - 22.7525)

                Gl = np.where(Gl >= 0, Gl, 0)
                return Gl

        class PAC2:
            def __init__(self):
                pass

            def __call__(self, energy):
                earray = np.array(energy)
                Gl = earray * 0.0 - 9876543210
                e1500index = _indexbelow(earray, 1500)
                Gl[e1500index] = 2.0e-5
                Gl = np.where(Gl >= 0, Gl, 0)
                return Gl


class _LogLinearfunc:
    """log linear function from the data points

    log y = (log y1 - log y0) / (log x1 - log x0) (log x - log x0) + log y0
    """
    def __init__(self, x0, y0, x1, y1):
        lx0 = np.log10(x0)
        ly0 = np.log10(y0)
        lx1 = np.log10(x1)
        ly1 = np.log10(y1)

        self._inclination = (ly1 - ly0) / (lx1 - lx0)
        self._lx0 = lx0
        self._ly0 = ly0

    def __call__(self, val):
        log_y = self._inclination * (np.log10(val) - self._lx0) + self._ly0
        return np.power(10, log_y)

    def __str__(self):
        return 'log(y) = {:g} x (log(x) - {:g}) + ({:g})'.format(self._inclination, self._lx0, self._ly0)


class GF_SWM_draft:
    """ Geometric factor used in SWM report, Figure 6 middle. Read from the figure.

    The geometric factor used in SWM project.

    .. warning::

        Do not use. Use :func:`gfactor_func` instead.

    Only for H+ and He++. The unit is in ``cm-2 eV/eV`` which implies the ``omega`` is already multiplied.
    Only PAC1 (pacclevel4) supported.
    """

    _stolen_x = [327, 543, 622, 1390, 1560, 2800, 106, 297, 15200, 845, 2045, 15000, ]
    _stolen_y = [0.000006870000, 0.000017400000, 0.000018000000, 0.000018000000, 0.000016100000, 0.000003380000,
                 0.000636000000, 0.000636000000, 0.000071700000, 0.000016400000, 0.000218000000, 0.000071700000, ]

    _A = _LogLinearfunc(_stolen_x[0], _stolen_y[0], _stolen_x[1], _stolen_y[1])  # Proton ghost, roughly 300 - 500 eV
    _B = _LogLinearfunc(_stolen_x[2], _stolen_y[2], _stolen_x[3], _stolen_y[3])  # Proton ghost, roughly 500 - 1500 eV
    _C = _LogLinearfunc(_stolen_x[4], _stolen_y[4], _stolen_x[5], _stolen_y[5])  # Proton ghost, roughly 1500 - 3000 eV
    _D = _LogLinearfunc(_stolen_x[6], _stolen_y[6], _stolen_x[7], _stolen_y[7])  # Alpha, roughly below 300 eV
    _E = _LogLinearfunc(_stolen_x[7], _stolen_y[7], _stolen_x[8], _stolen_y[8])  # Alpha, roughly above 300 eV
    _F = _LogLinearfunc(_stolen_x[9], _stolen_y[9], _stolen_x[10], _stolen_y[10])  # Proton, roughly 800 - 2000 eV
    _G = _LogLinearfunc(_stolen_x[10], _stolen_y[10], _stolen_x[11], _stolen_y[11])  # Proton, roughly 2000 eV or higher

    class Proton:
        class PAC1:
            def __init__(self):
                pass

            def __call__(self, energy):
                g0 = GF_SWM_draft._F(energy)
                g1 = GF_SWM_draft._G(energy)

                g = np.min([g0, g1], axis=0)
                g = np.where(g < 1e-5, 0, g)
                return g

    class ProtonGhost:
        class PAC1:
            def __init__(self):
                pass

            def __call__(self, energy):
                g0 = GF_SWM_draft._A(energy)
                g0 = np.where(g0 < 6e-6, 0, g0)
                g1 = GF_SWM_draft._B(energy)
                g2 = GF_SWM_draft._C(energy)
                g = np.min([g0, g1, g2], axis=0)

                return g

    class Alpha:
        class PAC1:
            def __init__(self):
                pass

            def __call__(self, energy):
                g0 = GF_SWM_draft._D(energy)
                g1 = GF_SWM_draft._E(energy)

                g = np.min([g0, g1], axis=0)

                return g


# This shall be changed after the formulation is available.
GF_SWM = GF_SWM_draft
""" SWM gfactor, but deprecated. Use :func:`gfactor_func`.
"""


def gfactor_func(species, eeprom, pacc):
    """ Geometric factor used in MEX IMA PSA archiving (2021-01-27)

    :param species: Species. Eitherof ['heavy', 'proton', 'alpha', 'ghost']
    :param eeprom: Prom sectin number, 0-16.
    :param pacc: PACC, either (0, 1, 2)
    :returns: Geometric factor in cm2 sr eV/eV.
    :exception ValueError: Given the unsupported species name.

    .. note::

        This is the current (2021-01) best knowledge of geometric factor at IRF
        with the "normal" definition (in cm2 sr eV/eV) including the efficiency.

    Provides the G factor [cm2 sr eV/eV] for a species at a specific eeprom and pacc level.
    Use pacc 0, 1, or 2 (corresponding to 0, 4, 7).

    The G-factor is used as follows:
    diff_flux = counts / (tau[s] * energy[eV] * G)
    to calculate differential flux of unit [cm-2 s-1 sr-1 eV-1].

    This version of the G-factor was used for the MEx IMA PSA archiving project
    as of 2021-01-27.

    >>> from irfpy.mima import gfactor0
    >>> gf = gfactor0.gfactor_func('proton', 16, 1)
    >>> print(gf.shape)   # GEometric factor for each energy step is returned, so the shape should be (96,)
    (96,)
    >>> print('{:.2e}'.format(gf[20]))   # Geometric factor for energy step 20
    4.12e-06

    >>> gf = gfactor0.gfactor_func('oxygen', 16, 2)   # PACC level 2 (=7)
    >>> print(gf.shape)
    (96,)
    >>> print("{:.2e}".format(gf[20]))
    7.80e-06
    """

    etbl = _energy.get_default_table(eeprom, keep_negative=True)
    Ebreakpoint = [np.nan, 300, 600]  # Value determined after IMA Kiruna workshop 2007
    dphi = 22.5 * np.pi / 180  # Azimuth sector size to go from rad to sr

    # Heavy ions
    es = etbl.shape[0]
    Gl = np.ones((3, es)) * np.nan
    Glh = Gl.copy()
    Gl[0, :] = np.exp(-0.378481 * np.log(etbl[:es]) - 9.05248)
    Glh[0, :] = Gl[0, :]
    Gl[1, :] = 7.0e-5 * np.ones((1, es))
    Glh[1, :] = np.exp(-0.69 * np.log(etbl[:es]) - 5.63)
    Gl[2, :] = 1e-4
    Glh[2, :] = np.exp(-0.611844 * np.log(etbl[:es]) - 5.31592)  # from Andrei
    aa = etbl >= Ebreakpoint[1]
    Gl[1, aa] = Glh[1, aa]
    aa = etbl >= Ebreakpoint[2]
    Gl[2, aa] = Glh[2, aa]

    if species.lower() in ['heavy', 'oxygen', 'o', 'op', 'o2', 'o2p']:
        return Gl[pacc, :] * dphi / 2  # Multiply with azimuth sector size to get the correct unit

    elif species.lower() in ['proton', 'hydrogen', 'h', 'p', 'hp']:
        # H+
        GH = Gl.copy()
        aa = etbl[:es] > 2000
        GH[1, aa] = np.exp(-0.554935 * np.log(etbl[aa]) - 6.40179)
        aa = (etbl[:es] < 2000) & (etbl[:es] > 700)
        GH[1, aa] = np.exp(3.0440526 * np.log(etbl[aa]) - 33.7573)
        aa = etbl > 1500
        GH[2, aa] = np.exp(-0.611844 * np.log(etbl[aa]) - 5.31592)
        return GH[pacc, :] * dphi / 2  # Multiply with azimuth sector size to get the correct unit

    elif species.lower() in ['alpha', 'helium', 'he', 'hepp', 'he2p', 'al']:
        # He++
        GHe = Gl.copy()
        aa = etbl[:es] > 300
        GHe[1, aa] = np.exp(-0.554935 * np.log(etbl[aa]) - 6.40179)
        return GHe[pacc, :] * dphi / 2  # Multiply with azimuth sector size to get the correct unit

    elif species.lower() in ['ghost', 'protonghost', 'g', 'hg', 'pg']:
        # H+ ghosts
        Gghost = np.ones((3, es)) * np.nan
        aa = (etbl[:es] > 300) & (etbl[:es] < 800)
        Gghost[1, aa] = 2 * np.exp((14.0 - 13.1224) / (6.68461 - 6.21461) * (np.log(etbl[aa]) - 6.21461) - 14.0)
        aa = (etbl[:es] >= 800) & (etbl[:es] < 1500)
        Gghost[1, aa] = 2.0e-6
        aa = (etbl[:es] >= 1500) & (etbl[:es] < 3000)
        Gghost[1, aa] = np.exp((13.1224 - 14.4902) / (7.82405 - 7.31322) * (np.log(etbl[aa]) - 7.31322) - 13.1224)

        # From old code for low and high pacc
        Gghost[0, :] = 1.6e-6 * np.ones((1, es)) * 16 / np.pi  # Version from Markus, second IMA
        Gghost[2, :] = 2.8e-6 * np.ones((1, es)) * 16 / np.pi
        aa = etbl[:es] < 1500
        Gghost[2, aa] = np.exp(1.44064 * np.log(etbl[aa]) - 22.7525)
        Gghost[Gghost < 0] = np.nan

        return Gghost[pacc, :] * dphi / 2  # Multiply with azimuth sector size to get the correct unit

    else:
        raise ValueError('Species not recognised.')