Quantity speed test for gammafit.ProtonOZM

quantity_test.py

import numpy as np
## Constants and units
from astropy import units as u
# import constant values from astropy.constants
import astropy.constants as const
from astropy.constants import c, G, m_e, h, hbar, k_B, R_sun, sigma_sb, e, m_p
e = e.gauss
mec2 = (m_e*c**2).cgs

# variables for integrand
mp = (m_p*c**2)
m_pi = 1.349766e-4*u.TeV  # TeV/c2

erad = (e**2./mec2).cgs
sigt = (8*np.pi/3)*erad**2
ar = (4*sigma_sb/c).cgs

class ProtonOZM(object):
    def __init__(self):
        self.norm=1.0
        self.index=2.1
        self._norm_energy=1.0

    def Jp(self, Ep):
        """
        Particle distribution function [1/TeV]
        """
        Jp = self.norm*(Ep/self._norm_energy)**-self.index

        return Jp

    def Fgamma(self, x, Ep):
        """
        KAB06 Eq.58

        Parameters
        ----------
        x : Egamma/Eprot
        Ep : Eprot [TeV]
        """
        L = np.log(Ep)
        B = 1.30+0.14*L+0.011*L**2  # Eq59
        beta = (1.79+0.11*L+0.008*L**2)**-1  # Eq60
        k = (0.801+0.049*L+0.014*L**2)**-1  # Eq61
        xb = x**beta

        F1 = B*(np.log(x)/x)*((1-xb)/(1+k*xb*(1-xb)))**4
        F2 = 1./np.log(x)-(4*beta*xb)/(1-xb)-(4*k*beta*xb*(1-2*xb))/(1+k*xb*(1-xb))

        return F1*F2

    def sigma_inel(self, Ep):
        """
        Inelastic cross-section for p-p interaction. KAB06 Eq. 73, 79
        """
        L = np.log(Ep)
        sigma = 34.3 + 1.88*L + 0.25*L**2
        if Ep <= 0.1:
            Eth = 1.22e-3
            sigma *= (1-(Eth/Ep)**4)**2 * heaviside(Ep-Eth)
        return sigma*1e-27  # convert from mbarn to cm2

    def _photon_integrand(self, x, Egamma):
        """
        Integrand of Eq. 72
        """
        try:
            return self.sigma_inel(Egamma/x)*self.Jp(Egamma/x) \
                *self.Fgamma(x, Egamma/x)/x
        except ZeroDivisionError:
            return np.nan

    def calc_hiE_spec(self,Egamma):

        from scipy.integrate import quad
        Egamma = Egamma.to('TeV').value
        result = c.cgs.value*quad(self._photon_integrand, 0., 1., args=Egamma,
                        epsrel=1e-3, epsabs=0)[0]

        return result * u.Unit('1/(s TeV)')

    # variables for integrand
    _c=c.cgs.value
    _Kpi=0.17
    _mp = (m_p*c**2).to('TeV').value
    _m_pi = 1.349766e-4  # TeV/c2

    def _delta_integrand(self,Epi):
        Ep0 = self._mp+Epi/self._Kpi
        qpi = self._c/self._Kpi*self.sigma_inel(Ep0)*self.Jp(Ep0)
        return qpi/np.sqrt(Epi**2+self._m_pi**2)

    def calc_loE_spec(self,Egamma):
        """
        Delta-functional approximation for low energies Egamma < 0.1 TeV
        """
        from scipy.integrate import quad
        Egamma = Egamma.to('TeV').value
        Epimin = Egamma+self._m_pi**2/(4*Egamma)

        result = 2*quad(self._delta_integrand, Epimin, np.inf, epsrel=1e-3,
                epsabs=0)[0]

        return result * u.Unit('1/(s TeV)')


class QuantityProtonOZM(object):
    def __init__(self):
        self.norm=1.0
        self.index=2.1
        self.norm_energy=1.0*u.TeV

    def Jp(self, Ep):
        """
        Particle distribution function [1/TeV]
        """
        Jp = self.norm*(Ep/self.norm_energy).decompose().value**-self.index

        return Jp

    def Fgamma(self, x, Ep):
        """
        KAB06 Eq.58

        Parameters
        ----------
        x : Egamma/Eprot
        Ep : Eprot [TeV]
        """
        L = np.log(Ep.to('TeV').value)
        B = 1.30+0.14*L+0.011*L**2  # Eq59
        beta = (1.79+0.11*L+0.008*L**2)**-1  # Eq60
        k = (0.801+0.049*L+0.014*L**2)**-1  # Eq61
        xb = x**beta

        F1 = B*(np.log(x)/x)*((1-xb)/(1+k*xb*(1-xb)))**4
        F2 = 1./np.log(x)-(4*beta*xb)/(1-xb)-(4*k*beta*xb*(1-2*xb))/(1+k*xb*(1-xb))

        return F1*F2

    def sigma_inel(self, Ep):
        """
        Inelastic cross-section for p-p interaction. KAB06 Eq. 73, 79
        """
        L = np.log(Ep.to('TeV').value)
        sigma = 34.3 + 1.88*L + 0.25*L**2
        if Ep <= 0.1*u.TeV:
            Eth = 1.22e-3*u.TeV
            sigma *= (1-(Eth/Ep).decompose().value**4)**2 * heaviside(Ep-Eth)
        return sigma*1e-27  # convert from mbarn to cm2

    def _photon_integrand(self, x, Egamma):
        """
        Integrand of Eq. 72
        """
        try:
            return self.sigma_inel(Egamma/x)*self.Jp(Egamma/x) \
                *self.Fgamma(x, Egamma/x)/x
        except ZeroDivisionError:
            return np.nan

    def calc_hiE_spec(self,Egamma):

        from scipy.integrate import quad
        result = c.cgs.value*quad(self._photon_integrand, 0., 1., args=Egamma,
                        epsrel=1e-3, epsabs=0)[0]

        return result * u.Unit('1/(s TeV)')

    def _delta_integrand(self,Epi):
        Kpi=0.17
        Ep0 = mp+Epi*u.TeV/Kpi
        qpi = c.cgs.value/Kpi*self.sigma_inel(Ep0)*self.Jp(Ep0)
        return qpi/np.sqrt((Epi*u.TeV)**2+m_pi**2).value

    def calc_loE_spec(self,Egamma):
        """
        Delta-functional approximation for low energies Egamma < 0.1 TeV
        """
        from scipy.integrate import quad
        # scalar needed for quad
        Epimin = (Egamma+m_pi**2/(4*Egamma)).to('TeV').value

        result = 2*quad(self._delta_integrand, Epimin, np.inf, epsrel=1e-3,
                epsabs=0)[0]

        return result * u.Unit('1/(s TeV)')