<commit_before>""" Module for basic plots related to FRB host and foreground galaxies"""
import os
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline as IUS

from IPython import embed

from pkg_resources import resource_filename

from matplotlib import pyplot as plt

from astropy.cosmology import Planck15 as cosmo
from astropy.cosmology import z_at_value
from astropy import units

from frb.halos import ModifiedNFW
from frb import halos as frb_halos
from frb import igm as frb_igm

from ne2001 import density

def sub_cartoon(ax1, ax2, coord, zFRB, halos=False, host_DM=0., ymax=None,
                M31=False, fg_halos=None, dsmx=0.05, FRB_DM=None, yscl = 0.97):
    """
    Cartoon of DM cumulative

    Args:
        ax1:
        ax2:
        coord:
        zFRB:
        halos:
        host_DM:
        ymax:
        M31:
        fg_halos:
        dsmx (float): Padding on the x-axis;  Gpc
        FRB_DM (float): Observed value;  sets ymax = FRB_DM+50

    Returns:

    """

    if halos:
        embed()

    gcoord = coord.transform_to('galactic')
    l, b = gcoord.l.value, gcoord.b.value

    ds = []  # kpc
    DM_cumul = []

    # ISM
    ne = density.ElectronDensity()  # **PARAMS)
    for ss in np.linspace(2., 4, 5):  # log pc
        idd = 10. ** ss / 1e3  # kpc
        iDM = ne.DM(l, b, idd)
        # Append
        ds.append(idd)  # kpc
        DM_cumul.append(iDM.value)
        # print(idd, iDM)
    max_ISM = DM_cumul[-1]

    # MW
    Mhalo = np.log10(1.5e12)  # Boylan-Kolchin et al. 2013
    f_hot = 0.75  # Allows for disk + ISM
    c = 7.7
    mnfw_2 = ModifiedNFW(log_Mhalo=Mhalo, f_hot=f_hot, y0=2, alpha=2, c=c)
    # Zero out inner 10kpc
    mnfw_2.zero_inner_ne = 10.  # kpc
    params = dict(F=1., e_density=1.)
    model_ne = density.NEobject(mnfw_2.ne, **params)
    for ss in np.linspace(1., np.log10(mnfw_2.r200.value), 5):  # log kpc
        idd = 10. ** ss  # kpc
        iDM = model_ne.DM(l, b, idd).value
        # Add it in
        if idd == ds[-1]:
            DM_cumul[-1] = DM_cumul[-1] + iDM
        else:
            ds.append(idd)
            DM_cumul.append(max_ISM + iDM)

    DM_ISM_Halo = DM_cumul[-1]

    '''
    # M31
    m31 = M31()
    a,c =1,0
    x0, y0 = m31.distance.to('kpc').value, 0. # kpc (<NAME>., <NAME>., & Valls - <NAME>. 2012, ApJ, 745, 156)
    sep = m31.coord.separation(coord)
    atan = np.arctan(sep.radian)
    b = -1 * a / atan
    M31_Rperp = np.abs(a * x0 + b * y0 + c) / np.sqrt(a ** 2 + b ** 2)  # kpc
    zval, M31_DM_cumul = m31.Ne_Rperp(M31_Rperp * u.kpc, rmax=1., cumul=True)
    # Add em in
    ds += (zval+x0).tolist()
    DM_cumul += (M31_DM_cumul+DM_ISM_Halo).tolist()

    #DM_LG = 0.
    '''
    DM_LG = DM_cumul[-1]

    # IGM
    z0 = z_at_value(cosmo.comoving_distance, 1 * units.Mpc)
    zvals = np.linspace(z0, 0.5, 50)
    dz_vals = zvals[1] - zvals[0]
    #
    fhalos = frb_halos.frac_in_halos(zvals, 3e10, 1e16, rmax=1.)
    fIGM = 1. - fhalos
    #
    DM_cosmic_cumul, zeval = frb_igm.average_DM(zvals[-1], cumul=True)
    dzeval = zeval[1] - zeval[0]
    dDM_cosmic = DM_cosmic_cumul - np.roll(DM_cosmic_cumul, 1)
    dDM_cosmic[0] = dDM_cosmic[1]
    #
    DM_interp = IUS(zeval, dDM_cosmic)
    dDM_IGM = DM_interp(zvals) * fIGM * dz_vals / dzeval
    sub_DM_IGM = np.cumsum(dDM_IGM)
    f_IGM = IUS(zvals, sub_DM_IGM)
    zvals2 = np.linspace(z0, zFRB, 1000)
    DM_IGM = f_IGM(zvals2)

    DM_cosmic = DM_IGM.copy()

    # Halos at last
    if fg_halos is not None:
        for z, halo_DM, lbl in zip(fg_halos['z'], fg_halos['DM'], fg_halos['lbl']):
            iz = np.argmin(np.abs(zvals2 - z))
            DM_cosmic[iz:] += halo_DM
            # Label
            d = cosmo.comoving_distance(z)
            ax1.text(d.to('Gpc').value, DM_cosmic[iz], lbl, color='black',
                     fontsize=13, ha='left', va='top')

    Dc = cosmo.comoving_distance(zvals2).to('kpc')
    ds += Dc.value.tolist()
    DM_cumul += (DM_cosmic + DM_LG).tolist()

    # Host
    if host_DM > 0.:
        ds.append(ds[-1])
        DM_cumul.append(DM_cumul[-1] + host_DM)

    # Plot the DM curve
    ax1.plot(ds, DM_cumul, 'k')

    # max_y = np.max(DM_cumul)
    if FRB_DM is not None:
        ymax = FRB_DM + 50.
    if ymax is not None:
        max_y = ymax
    # Shade me
    lsz = 14.
    ax1.fill_between((0.1, 10.), 0, max_y, color='green', alpha=0.4)  # ISM
    ax1.text(0.15, max_y * yscl, r'\textbf{Galactic}'+'\n'+r'\textbf{ISM}', color='black', fontsize=lsz, ha='left', va='top')
    ax1.fill_between((10., mnfw_2.r200.value), 0, max_y, color='blue', alpha=0.4)  # Galactic Halo
    ax1.text(12., max_y * yscl, r'\textbf{Galactic}'+'\n'+r'\textbf{Halo}', color='black', fontsize=lsz, ha='left', va='top')
    if M31:
        ax1.fill_between((mnfw_2.r200.value, 2e3), 0, max_y, color='red', alpha=0.4)  # Galactic Halo
        ax1.text(300., max_y * yscl, r'\texgbf{M31}', color='black', fontsize=lsz, ha='left', va='top')

    ax1.set_xscale("log", nonposx='clip')
    # ax.set_yscale("log", nonposy='clip')
    if M31:
        ax1.set_xlim(0.1, 2e3)  # kpc
    else:
        ax1.set_xlim(0.1, mnfw_2.r200.value)  # kpc
    ax1.set_ylim(0., max_y)
    ax1.spines['right'].set_visible(False)
    ax1.set_xlabel(r'\textbf{Distance (kpc)}')
    ax1.set_ylabel(r'\textbf{Cumulative DM (pc cm$^{-3}$)}')

    # IGM
    Gds = np.array(ds) / 1e6
    ax2.plot(Gds, DM_cumul, 'k')
    ax2.spines['left'].set_visible(False)
    ax2.yaxis.tick_right()
    ax2.tick_params(labelright='off')
    ax2.minorticks_on()
    ax2.set_xlabel(r'\textbf{Distance (Gpc)}')

    smax = cosmo.comoving_distance(zFRB).to('Gpc').value
    ax2.set_xlim(0.1, smax+dsmx)  # Gpc
    ax2.fill_between((0.1, smax-dsmx), 0, max_y, color='gray', alpha=0.4)  # Galactic Halo
    ax2.fill_between((smax-dsmx, smax+dsmx), 0, max_y, color='red', alpha=0.4)  # Galactic Halo
    ilbl = r'\textbf{Cosmic}'
    ax2.text(0.2, max_y * yscl, ilbl, color='black', fontsize=lsz, ha='left', va='top')

    if FRB_DM is not None:
        ax1.axhline(y=FRB_DM, ls='--', color='black')
        ax2.axhline(y=FRB_DM, ls='--', color='black')

<commit_msg>dm cartoon continued<commit_after>""" Module for basic plots related to FRB host and foreground galaxies"""
import os
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline as IUS

from IPython import embed

from pkg_resources import resource_filename

from matplotlib import pyplot as plt

from astropy.cosmology import Planck15 as cosmo
from astropy.cosmology import z_at_value
from astropy import units

from frb.halos import ModifiedNFW
from frb import halos as frb_halos
from frb import igm as frb_igm

from ne2001 import density

def sub_cartoon(ax1, ax2, coord, zFRB, halos=False, host_DM=0., ymax=None,
                IGM_only=True,
                M31=False, fg_halos=None, dsmx=0.05, FRB_DM=None, yscl = 0.97):
    """
    Cartoon of DM cumulative

    Args:
        ax1:
        ax2:
        coord:
        zFRB:
        halos:
        host_DM:
        ymax:
        M31:
        fg_halos:
        dsmx (float): Padding on the x-axis;  Gpc
        FRB_DM (float): Observed value;  sets ymax = FRB_DM+50

    Returns:

    """

    if halos:
        embed()

    gcoord = coord.transform_to('galactic')
    l, b = gcoord.l.value, gcoord.b.value

    ds = []  # kpc
    DM_cumul = []

    # ISM
    ne = density.ElectronDensity()  # **PARAMS)
    for ss in np.linspace(2., 4, 5):  # log pc
        idd = 10. ** ss / 1e3  # kpc
        iDM = ne.DM(l, b, idd)
        # Append
        ds.append(idd)  # kpc
        DM_cumul.append(iDM.value)
        # print(idd, iDM)
    max_ISM = DM_cumul[-1]

    # MW
    Mhalo = np.log10(1.5e12)  # <NAME> al. 2013
    f_hot = 0.75  # Allows for disk + ISM
    c = 7.7
    mnfw_2 = ModifiedNFW(log_Mhalo=Mhalo, f_hot=f_hot, y0=2, alpha=2, c=c)
    # Zero out inner 10kpc
    mnfw_2.zero_inner_ne = 10.  # kpc
    params = dict(F=1., e_density=1.)
    model_ne = density.NEobject(mnfw_2.ne, **params)
    for ss in np.linspace(1., np.log10(mnfw_2.r200.value), 5):  # log kpc
        idd = 10. ** ss  # kpc
        iDM = model_ne.DM(l, b, idd).value
        # Add it in
        if idd == ds[-1]:
            DM_cumul[-1] = DM_cumul[-1] + iDM
        else:
            ds.append(idd)
            DM_cumul.append(max_ISM + iDM)

    DM_ISM_Halo = DM_cumul[-1]

    '''
    # M31
    m31 = M31()
    a,c =1,0
    x0, y0 = m31.distance.to('kpc').value, 0. # kpc (<NAME>., <NAME>., & Valls - <NAME>. 2012, ApJ, 745, 156)
    sep = m31.coord.separation(coord)
    atan = np.arctan(sep.radian)
    b = -1 * a / atan
    M31_Rperp = np.abs(a * x0 + b * y0 + c) / np.sqrt(a ** 2 + b ** 2)  # kpc
    zval, M31_DM_cumul = m31.Ne_Rperp(M31_Rperp * u.kpc, rmax=1., cumul=True)
    # Add em in
    ds += (zval+x0).tolist()
    DM_cumul += (M31_DM_cumul+DM_ISM_Halo).tolist()

    #DM_LG = 0.
    '''
    DM_LG = DM_cumul[-1]

    # IGM
    z0 = z_at_value(cosmo.comoving_distance, 1 * units.Mpc)
    zvals = np.linspace(z0, 0.5, 50)
    dz_vals = zvals[1] - zvals[0]
    #
    fhalos = frb_halos.frac_in_halos(zvals, 3e10, 1e16, rmax=1.)
    fIGM = 1. - fhalos
    #
    DM_cosmic_cumul, zeval = frb_igm.average_DM(zvals[-1], cumul=True)
    dzeval = zeval[1] - zeval[0]
    dDM_cosmic = DM_cosmic_cumul - np.roll(DM_cosmic_cumul, 1)
    dDM_cosmic[0] = dDM_cosmic[1]
    #
    DM_interp = IUS(zeval, dDM_cosmic)
    dDM_IGM = DM_interp(zvals) * fIGM * dz_vals / dzeval
    dDM_cosm = DM_interp(zvals) * dz_vals / dzeval
    sub_DM_IGM = np.cumsum(dDM_IGM)
    sub_DM_cosm = np.cumsum(dDM_cosm)
    f_IGM = IUS(zvals, sub_DM_IGM)
    f_cosm = IUS(zvals, sub_DM_cosm)
    zvals2 = np.linspace(z0, zFRB, 1000)
    DM_IGM = f_IGM(zvals2)
    DM_cosmic = f_cosm(zvals2)

    # Ignore halos?
    if IGM_only:
        DM_cosmic = DM_IGM.copy()

    # Halos at last
    if fg_halos is not None:
        for z, halo_DM, lbl in zip(fg_halos['z'], fg_halos['DM'], fg_halos['lbl']):
            iz = np.argmin(np.abs(zvals2 - z))
            DM_cosmic[iz:] += halo_DM
            # Label
            d = cosmo.comoving_distance(z)
            ax1.text(d.to('Gpc').value, DM_cosmic[iz], lbl, color='black',
                     fontsize=13, ha='left', va='top')

    Dc = cosmo.comoving_distance(zvals2).to('kpc')
    ds += Dc.value.tolist()
    DM_cumul += (DM_cosmic + DM_LG).tolist()

    # Host
    if host_DM > 0.:
        ds.append(ds[-1])
        DM_cumul.append(DM_cumul[-1] + host_DM)

    # Plot the DM curve
    ax1.plot(ds, DM_cumul, 'k')

    # max_y = np.max(DM_cumul)
    if FRB_DM is not None:
        ymax = FRB_DM + 50.
    if ymax is not None:
        max_y = ymax
    # Shade me
    lsz = 14.
    ax1.fill_between((0.1, 10.), 0, max_y, color='green', alpha=0.4)  # ISM
    ax1.text(0.15, max_y * yscl, r'\textbf{Galactic}'+'\n'+r'\textbf{ISM}', color='black', fontsize=lsz, ha='left', va='top')
    ax1.fill_between((10., mnfw_2.r200.value), 0, max_y, color='blue', alpha=0.4)  # Galactic Halo
    ax1.text(12., max_y * yscl, r'\textbf{Galactic}'+'\n'+r'\textbf{Halo}', color='black', fontsize=lsz, ha='left', va='top')
    if M31:
        ax1.fill_between((mnfw_2.r200.value, 2e3), 0, max_y, color='red', alpha=0.4)  # Galactic Halo
        ax1.text(300., max_y * yscl, r'\texgbf{M31}', color='black', fontsize=lsz, ha='left', va='top')

    ax1.set_xscale("log", nonposx='clip')
    # ax.set_yscale("log", nonposy='clip')
    if M31:
        ax1.set_xlim(0.1, 2e3)  # kpc
    else:
        ax1.set_xlim(0.1, mnfw_2.r200.value)  # kpc
    ax1.set_ylim(0., max_y)
    ax1.spines['right'].set_visible(False)
    ax1.set_xlabel(r'\textbf{Distance (kpc)}')
    ax1.set_ylabel(r'\textbf{Cumulative DM (pc cm$^{-3}$)}')

    # IGM
    Gds = np.array(ds) / 1e6
    ax2.plot(Gds, DM_cumul, 'k')
    ax2.spines['left'].set_visible(False)
    ax2.yaxis.tick_right()
    ax2.tick_params(labelright='off')
    ax2.minorticks_on()
    ax2.set_xlabel(r'\textbf{Distance (Gpc)}')

    smax = cosmo.comoving_distance(zFRB).to('Gpc').value
    ax2.set_xlim(0.1, smax+dsmx)  # Gpc
    ax2.fill_between((0.1, smax-dsmx), 0, max_y, color='gray', alpha=0.4)  # Galactic Halo
    ax2.fill_between((smax-dsmx, smax+dsmx), 0, max_y, color='red', alpha=0.4)  # Galactic Halo
    ilbl = r'\textbf{Cosmic}'
    ax2.text(0.2, max_y * yscl, ilbl, color='black', fontsize=lsz, ha='left', va='top')

    if FRB_DM is not None:
        ax1.axhline(y=FRB_DM, ls='--', color='black')
        ax2.axhline(y=FRB_DM, ls='--', color='black')


{'fraction_non_alphanumeric': 0.15346495084325792, 'fraction_numerical': 0.03988490128686756, 'mean_word_length': 2.54849688031764, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2081914', 'n_tokens_mistral': 5719, 'n_tokens_neox': 5313, 'n_words': 1073}