<reponame>jorana/wimprates
"""Dark matter - electron scattering
"""
import numericalunits as nu
import numpy as np
from scipy.interpolate import RegularGridInterpolator, interp1d
from scipy.integrate import quad, dblquad

import wimprates as wr
export, __all__ = wr.exporter()
__all__ += ['dme_shells', 'l_to_letter', 'l_to_number']

# Load form factor and construct interpolators
shell_data = wr.load_pickle('dme/dme_ionization_ff.pkl')
for _shell, _sd in shell_data.items():
    _sd['log10ffsquared_itp'] = RegularGridInterpolator(
        (_sd['lnks'], _sd['lnqs']),
        np.log10(_sd['ffsquared']),
        bounds_error=False, fill_value=-float('inf'),)


dme_shells = [(5, 1), (5, 0), (4, 2), (4, 1), (4, 0)]
l_to_number = dict(s=0, p=1, d=2, f=3)
l_to_letter = {v: k for k, v in l_to_number.items()}


@export
def shell_str(n, l):
    if isinstance(l, str):
        return str(n) + l
    return str(n) + l_to_letter[l]



@export
def dme_ionization_ff(shell, e_er, q):
    """Return dark matter electron scattering ionization form factor

    Outside the parametrized range, the form factor is assumed 0
    to give conservative results.

    :param shell: Name of atomic shell, e.g. '4p'
        Note not all shells are included in the data.
    :param e_er: Electronic recoil energy
    :param q: Momentun transfer
    """
    if isinstance(shell, tuple):
        shell = shell_str(*shell)

    lnq = np.log(q / (nu.me * nu.c0 * nu.alphaFS))
    # From Mathematica: (*ER*) (2 lnkvalues[[j]])/Log[10]
    # log10 (E/Ry) = 2 lnk / ln10
    # lnk = log10(E/Ry) * ln10 / 2
    #     = lng(E/Ry) / 2
    # Ry = rydberg = 13.6 eV
    ry = nu.me * nu.e ** 4 / (8 * nu.eps0 ** 2 * nu.hPlanck ** 2)
    lnk = np.log(e_er / ry) / 2
    return 10**(shell_data[shell]['log10ffsquared_itp'](
        np.vstack([lnk, lnq]).T))


@export
def binding_es_for_dme(n, l):
    """Return binding energy of Xenon's (n, l) orbital
    according to Essig et al. 2017 Table II

    Note these are different from e.g. Ibe et al. 2017!
    """

    return {'4s': 213.8,
            '4p': 163.5,
            '4d': 75.6,
            '5s': 25.7,
            '5p': 12.4}[shell_str(n, l)] * nu.eV


@export
def v_min_dme(eb, erec, q, mw):
    """Minimal DM velocity for DM-electron scattering
    :param eb: binding energy of shell
    :param erec: electronic recoil energy energy
    :param q: momentum transfer
    :param mw: DM mass
    """
    return (erec + eb) / q + q / (2 * mw)


# Precompute velocity integrals for t=None
@export 
def velocity_integral_without_time(halo_model=None):
    halo_model = wr.StandardHaloModel() if halo_model is None else halo_model
    _v_mins = np.linspace(0, 1, 1000) * wr.v_max(None, halo_model.v_esc)
    _ims = np.array([
        quad(lambda v: 1 / v * halo_model.velocity_dist(v,None),
            _v_min,
             wr.v_max(None, halo_model.v_esc ))[0]
        for _v_min in _v_mins])
    
    # Store interpolator in km/s rather than unit-dependent numbers
    # so we don't have to recalculate them when nu.reset_units() is called
    inverse_mean_speed_kms = interp1d(
        _v_mins / (nu.km/nu.s),
        _ims * (nu.km/nu.s),
        # If we don't have 0 < v_min < v_max, we want to return 0
        # so the integrand vanishes
        fill_value=0, bounds_error=False)
    return inverse_mean_speed_kms

inverse_mean_speed_kms = velocity_integral_without_time()


@export
@wr.vectorize_first
def rate_dme(erec, n, l, mw, sigma_dme,
             f_dm='1',
             t=None, halo_model = None, **kwargs):
    """Return differential rate of dark matter electron scattering vs energy
    (i.e. dr/dE, not dr/dlogE)
    :param erec: Electronic recoil energy
    :param n: Principal quantum numbers of the shell that is hit
    :param l: Angular momentum quantum number of the shell that is hit
    :param mw: DM mass
    :param sigma_dme: DM-free electron scattering cross-section at fixed
    momentum transfer q=0
    :param f_dm: One of the following:
     '1':     |F_DM|^2 = 1, contact interaction / heavy mediator (default)
     '1_q':   |F_DM|^2 = (\alpha m_e c / q), dipole moment
     '1_q2': |F_DM|^2 = (\alpha m_e c / q)^2, ultralight mediator
    :param t: A J2000.0 timestamp.
    If not given, a conservative velocity distribution is used.
    :param halo_model: class (default to standard halo model) containing velocity distribution
    """
    halo_model = wr.StandardHaloModel() if halo_model is None else halo_model
    shell = shell_str(n, l)
    eb = binding_es_for_dme(n, l)

    f_dm = {
        '1': lambda q: 1,
        '1_q': lambda q: nu.alphaFS * nu.me * nu.c0 / q,
        '1_q2': lambda q: (nu.alphaFS * nu.me * nu.c0 / q)**2
    }[f_dm]

    # No bounds are given for the q integral
    # but the form factors are only specified in a limited range of q
    qmax = (np.exp(shell_data[shell]['lnqs'].max())
            * (nu.me * nu.c0 * nu.alphaFS))

    if t is None:
        # Use precomputed inverse mean speed,
        # so we only have to do a single integral
        def diff_xsec(q):
            vmin = v_min_dme(eb, erec, q, mw)
            result = q * dme_ionization_ff(shell, erec, q) * f_dm(q)**2
            # Note the interpolator is in kms, not unit-carrying numbers
            # see above
            result *= inverse_mean_speed_kms(vmin / (nu.km/nu.s))
            result /= (nu.km/nu.s)
            return result

        r = quad(diff_xsec, 0, qmax)[0]

    else:
        # Have to do double integral
        # Note dblquad expects the function to be f(y, x), not f(x, y)...
        def diff_xsec(v, q):
            result = q * dme_ionization_ff(shell, erec, q) * f_dm(q)**2
            result *= 1 / v * halo_model.velocity_dist(v, t)
            return result

        r = dblquad(
            diff_xsec,
            0,
            qmax,
            lambda q: v_min_dme(eb, erec, q, mw),
            lambda _: wr.v_max(t, halo_model.v_esc),
            **kwargs)[0]

    mu_e = mw * nu.me / (mw + nu.me)

    return (
        # Convert cross-section to rate, as usual
        halo_model.rho_dm / mw * (1 / wr.mn())
        # d/lnE -> d/E
        * 1 / erec
        # Prefactors in cross-section
        * sigma_dme / (8 * mu_e ** 2)
        * r)

"""Dark matter - electron scattering
"""
import numericalunits as nu
import numpy as np
from scipy.interpolate import RegularGridInterpolator, interp1d
from scipy.integrate import quad, dblquad

import wimprates as wr
export, __all__ = wr.exporter()
__all__ += ['dme_shells', 'l_to_letter', 'l_to_number']

# Load form factor and construct interpolators
shell_data = wr.load_pickle('dme/dme_ionization_ff.pkl')
for _shell, _sd in shell_data.items():
    _sd['log10ffsquared_itp'] = RegularGridInterpolator(
        (_sd['lnks'], _sd['lnqs']),
        np.log10(_sd['ffsquared']),
        bounds_error=False, fill_value=-float('inf'),)


dme_shells = [(5, 1), (5, 0), (4, 2), (4, 1), (4, 0)]
l_to_number = dict(s=0, p=1, d=2, f=3)
l_to_letter = {v: k for k, v in l_to_number.items()}


@export
def shell_str(n, l):
    if isinstance(l, str):
        return str(n) + l
    return str(n) + l_to_letter[l]



@export
def dme_ionization_ff(shell, e_er, q):
    """Return dark matter electron scattering ionization form factor

    Outside the parametrized range, the form factor is assumed 0
    to give conservative results.

    :param shell: Name of atomic shell, e.g. '4p'
        Note not all shells are included in the data.
    :param e_er: Electronic recoil energy
    :param q: Momentun transfer
    """
    if isinstance(shell, tuple):
        shell = shell_str(*shell)

    lnq = np.log(q / (nu.me * nu.c0 * nu.alphaFS))
    # From Mathematica: (*ER*) (2 lnkvalues[[j]])/Log[10]
    # log10 (E/Ry) = 2 lnk / ln10
    # lnk = log10(E/Ry) * ln10 / 2
    #     = lng(E/Ry) / 2
    # Ry = rydberg = 13.6 eV
    ry = nu.me * nu.e ** 4 / (8 * nu.eps0 ** 2 * nu.hPlanck ** 2)
    lnk = np.log(e_er / ry) / 2
    return 10**(shell_data[shell]['log10ffsquared_itp'](
        np.vstack([lnk, lnq]).T))


@export
def binding_es_for_dme(n, l):
    """Return binding energy of Xenon's (n, l) orbital
    according to Essig et al. 2017 Table II

    Note these are different from e.g. Ibe et al. 2017!
    """

    return {'4s': 213.8,
            '4p': 163.5,
            '4d': 75.6,
            '5s': 25.7,
            '5p': 12.4}[shell_str(n, l)] * nu.eV


@export
def v_min_dme(eb, erec, q, mw):
    """Minimal DM velocity for DM-electron scattering
    :param eb: binding energy of shell
    :param erec: electronic recoil energy energy
    :param q: momentum transfer
    :param mw: DM mass
    """
    return (erec + eb) / q + q / (2 * mw)


# Precompute velocity integrals for t=None
@export 
def velocity_integral_without_time(halo_model=None):
    halo_model = wr.StandardHaloModel() if halo_model is None else halo_model
    _v_mins = np.linspace(0, 1, 1000) * wr.v_max(None, halo_model.v_esc)
    _ims = np.array([
        quad(lambda v: 1 / v * halo_model.velocity_dist(v,None),
            _v_min,
             wr.v_max(None, halo_model.v_esc ))[0]
        for _v_min in _v_mins])
    
    # Store interpolator in km/s rather than unit-dependent numbers
    # so we don't have to recalculate them when nu.reset_units() is called
    inverse_mean_speed_kms = interp1d(
        _v_mins / (nu.km/nu.s),
        _ims * (nu.km/nu.s),
        # If we don't have 0 < v_min < v_max, we want to return 0
        # so the integrand vanishes
        fill_value=0, bounds_error=False)
    return inverse_mean_speed_kms

inverse_mean_speed_kms = velocity_integral_without_time()


@export
@wr.vectorize_first
def rate_dme(erec, n, l, mw, sigma_dme,
             f_dm='1',
             t=None, halo_model = None, **kwargs):
    """Return differential rate of dark matter electron scattering vs energy
    (i.e. dr/dE, not dr/dlogE)
    :param erec: Electronic recoil energy
    :param n: Principal quantum numbers of the shell that is hit
    :param l: Angular momentum quantum number of the shell that is hit
    :param mw: DM mass
    :param sigma_dme: DM-free electron scattering cross-section at fixed
    momentum transfer q=0
    :param f_dm: One of the following:
     '1':     |F_DM|^2 = 1, contact interaction / heavy mediator (default)
     '1_q':   |F_DM|^2 = (\alpha m_e c / q), dipole moment
     '1_q2': |F_DM|^2 = (\alpha m_e c / q)^2, ultralight mediator
    :param t: A J2000.0 timestamp.
    If not given, a conservative velocity distribution is used.
    :param halo_model: class (default to standard halo model) containing velocity distribution
    """
    halo_model = wr.StandardHaloModel() if halo_model is None else halo_model
    shell = shell_str(n, l)
    eb = binding_es_for_dme(n, l)

    f_dm = {
        '1': lambda q: 1,
        '1_q': lambda q: nu.alphaFS * nu.me * nu.c0 / q,
        '1_q2': lambda q: (nu.alphaFS * nu.me * nu.c0 / q)**2
    }[f_dm]

    # No bounds are given for the q integral
    # but the form factors are only specified in a limited range of q
    qmax = (np.exp(shell_data[shell]['lnqs'].max())
            * (nu.me * nu.c0 * nu.alphaFS))

    if t is None:
        # Use precomputed inverse mean speed,
        # so we only have to do a single integral
        def diff_xsec(q):
            vmin = v_min_dme(eb, erec, q, mw)
            result = q * dme_ionization_ff(shell, erec, q) * f_dm(q)**2
            # Note the interpolator is in kms, not unit-carrying numbers
            # see above
            result *= inverse_mean_speed_kms(vmin / (nu.km/nu.s))
            result /= (nu.km/nu.s)
            return result

        r = quad(diff_xsec, 0, qmax)[0]

    else:
        # Have to do double integral
        # Note dblquad expects the function to be f(y, x), not f(x, y)...
        def diff_xsec(v, q):
            result = q * dme_ionization_ff(shell, erec, q) * f_dm(q)**2
            result *= 1 / v * halo_model.velocity_dist(v, t)
            return result

        r = dblquad(
            diff_xsec,
            0,
            qmax,
            lambda q: v_min_dme(eb, erec, q, mw),
            lambda _: wr.v_max(t, halo_model.v_esc),
            **kwargs)[0]

    mu_e = mw * nu.me / (mw + nu.me)

    return (
        # Convert cross-section to rate, as usual
        halo_model.rho_dm / mw * (1 / wr.mn())
        # d/lnE -> d/E
        * 1 / erec
        # Prefactors in cross-section
        * sigma_dme / (8 * mu_e ** 2)
        * r)
Dark matter - electron scattering # Load form factor and construct interpolators Return dark matter electron scattering ionization form factor      Outside the parametrized range, the form factor is assumed 0     to give conservative results.      :param shell: Name of atomic shell, e.g. '4p'         Note not all shells are included in the data.     :param e_er: Electronic recoil energy     :param q: Momentun transfer # From Mathematica: (*ER*) (2 lnkvalues[[j]])/Log[10] # log10 (E/Ry) = 2 lnk / ln10 # lnk = log10(E/Ry) * ln10 / 2 #     = lng(E/Ry) / 2 # Ry = rydberg = 13.6 eV Return binding energy of Xenon's (n, l) orbital     according to Essig et al. 2017 Table II      Note these are different from e.g. Ibe et al. 2017! Minimal DM velocity for DM-electron scattering     :param eb: binding energy of shell     :param erec: electronic recoil energy energy     :param q: momentum transfer     :param mw: DM mass # Precompute velocity integrals for t=None # Store interpolator in km/s rather than unit-dependent numbers # so we don't have to recalculate them when nu.reset_units() is called # If we don't have 0 < v_min < v_max, we want to return 0 # so the integrand vanishes Return differential rate of dark matter electron scattering vs energy     (i.e. dr/dE, not dr/dlogE)     :param erec: Electronic recoil energy     :param n: Principal quantum numbers of the shell that is hit     :param l: Angular momentum quantum number of the shell that is hit     :param mw: DM mass     :param sigma_dme: DM-free electron scattering cross-section at fixed     momentum transfer q=0     :param f_dm: One of the following:      '1':     |F_DM|^2 = 1, contact interaction / heavy mediator (default)      '1_q':   |F_DM|^2 = (\alpha m_e c / q), dipole moment      '1_q2': |F_DM|^2 = (\alpha m_e c / q)^2, ultralight mediator     :param t: A J2000.0 timestamp.     If not given, a conservative velocity distribution is used.     :param halo_model: class (default to standard halo model) containing velocity distribution # No bounds are given for the q integral # but the form factors are only specified in a limited range of q # Use precomputed inverse mean speed, # so we only have to do a single integral # Note the interpolator is in kms, not unit-carrying numbers # see above # Have to do double integral # Note dblquad expects the function to be f(y, x), not f(x, y)... # Convert cross-section to rate, as usual # d/lnE -> d/E # Prefactors in cross-section