# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# +
import matplotlib.pyplot as plt
import numpy as np

from hazma.scalar_mediator import ScalarMediator
# -

gsxx = 1.0
gsff = 1.0
gsGG = 1.0
gsFF = 1.0
ms = 1.0
mx = 144.

# + [markdown] heading_collapsed=true
# ### e-ASTROGAM's effective area

# + hidden=true
from hazma.gamma_ray_limits.gamma_ray_limit_parameters import A_eff_e_ASTROGAM

e_gams = np.logspace(0, np.log10(500), 500)

plt.loglog(e_gams, [A_eff_e_ASTROGAM(e) for e in e_gams], label="Actual")
plt.loglog(e_gams[[0, -1]], [500, 500], label="Estimate")

plt.xlabel(r"$E_\gamma$ (MeV)")
plt.ylabel(r"$A_{\text{eff}}$ (cm$^2$)")
plt.legend()

# + [markdown] heading_collapsed=true
# ### DM annihilation spectra

# + hidden=true
from hazma.gamma_ray_limits import gamma_ray_limit_parameters

sm = ScalarMediator(mx, ms, gsxx, gsff, gsGG, gsFF)

e_gams = np.logspace(0, np.log10(770), 250)

# DM spectra
for mx in [110., 130.]: # 144., 285., 550., 770.]:
    sm.mx = mx
    spectra = sm.spectra(e_gams, 2.01*mx)
    plt.loglog(e_gams, e_gams * spectra["total"], label=r"$m_X = %.0f$ MeV" % mx)

# Background spectrum
plt.loglog(gamma_ray_limit_parameters.e_Bs, gamma_ray_limit_parameters.e_Bs * gamma_ray_limit_parameters.dPhi_dEdOmega_Bs, label="Background")

plt.xlim(e_gams[[0, -1]])
plt.legend()
# -

# ### Compute limits

# +
sm = ScalarMediator(mx, ms, gsxx, gsff, gsGG, gsFF)

mx_min = 50.
mx_max = 1000.
n_mxs = 25

mxs = np.logspace(np.log10(mx_min), np.log10(mx_max), n_mxs)

# Compute limits
sv_lims = sm.compute_limits(mxs)

# +
plt.plot(mxs, sv_lims)

plt.grid()
plt.xlim(mxs[[0, -1]])
plt.ylim([1e-29, 1e-26])
plt.yscale("log")
plt.xlabel(r"$m_X$ (MeV)")
plt.ylabel(r"$\langle \sigma v \rangle$ (cm$^3$/s)")
plt.title("Projected e-ASTROGAM limits")

# + [markdown] heading_collapsed=true
# ### Scratch

# + hidden=true
from scipy.interpolate import interp1d
from hazma.gamma_ray_limits.gamma_ray_limit_parameters import eASTROGAM_params, dSph_params
from hazma.gamma_ray_limits.compute_limits import __f_lim, __jac_lim

sm = ScalarMediator(104., ms, gsxx, gsff, gsGG, gsFF)

e_gams = np.logspace(np.log10(sm.mx / 100.), np.log10(sm.mx), 200)                 
dN_dE_DM = interp1d(e_gams, sm.spectra(e_gams, 2.001*sm.mx)["total"])

def fn(e_a, e_b):
    e_a = min([e_a, e_b])
    e_b = max([e_a, e_b])
    
    if e_a == e_b:
        return 0.
    else:
        return __f_lim([e_a, e_b], dN_dE_DM, eASTROGAM_params, dSph_params)


# + hidden=true
def npmap2d(fun, x_spec, y_spec, doPrint=False):
    xs = np.linspace(*x_spec)
    ys = np.linspace(*y_spec)
    Z = np.empty(len(xs) * len(ys))
    i = 0
    for y in ys:
        for x in xs:
            Z[i] = fun(x, y)
            if doPrint: print([i, x, y, Z[i]])
            i += 1
    X, Y = np.meshgrid(xs, ys)
    Z.shape = X.shape
    return X, Y, Z


# + hidden=true
e_as, e_bs, fn_vals = npmap2d(fn, (sm.mx / 50., sm.mx, 10), (sm.mx / 50., sm.mx, 10))

# + hidden=true
CS = plt.contourf(e_as, e_bs, fn_vals)
plt.colorbar()  
plt.show()

# + hidden=true
import scipy
from scipy import optimize
from hazma.parameters import neutral_pion_mass as mpi0

def f(e_ab):
    e_a = min(e_ab)
    e_b = max(e_ab)
    
    if e_a == e_b:
        return 0.
    else:
        return __f_lim([e_a, e_b], dN_dE_DM, eASTROGAM_params, dSph_params)

e_a_0 = 0.5 * (sm.mx - sm.mx / 100.)
e_b_0 = 0.75 * (sm.mx - sm.mx / 100.)
e_bounds = [sm.mx / 100., sm.mx]

limit_obj = scipy.optimize.minimize(f, [e_a_0, e_b_0],
                                    bounds=[e_bounds, e_bounds],
#                                     args=(dN_dE_DM, eASTROGAM_params,
#                                           dSph_params),
                                    method="L-BFGS-B",
                                    options={"ftol": 1e-3})

# + hidden=true
limit_obj

# + hidden=true


StarcoderdataJupyterScriptsDedupFiltered