aboutsummaryrefslogtreecommitdiff
path: root/utils/plot-root-results
AgeCommit message (Collapse)Author
2020-05-11add setup_matplotlib function and switch to logarithmic binstlatorre
This commit contains the following small updates: - create a setup_matplotlib() function to set up matplotlib correctly depending on if we are saving the plots or just displaying them - change default font size to 12 when displaying plots - switch to using logarithmic bins in plot-energy - fix despine() function when x axis is logarithmic
2020-05-11update utils/ folder to make a python package called sddmtlatorre
This commit adds an sddm python package to the utils/ folder. This allows me to consolidate code used across all the various scripts. This package is now installed by default to /home/tlatorre/local/lib/python2.7/site-packages so you should add the following to your .bashrc file: export PYTHONPATH=$HOME/local/lib/python2.7/site-packages/:$PYTHONPATH before using the scripts installed to ~/local/bin.
2020-04-27add a script to plot the results of the ROOT fitstlatorre
a id='n71' href='#n71'>71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322
#!/usr/bin/env python
"""
Script to combine the low energy Battistoni atmospheric flux with the high
energy flux from Giles Barr and apply oscillations.

Note: There are a couple of potential issues with this:

1. The Battistoni fluxes give a total flux integrated over the zenith angle
whereas we need to account for the zenith dependence. Therefore, we just use
the zenith dependence of the lowest energy in the Barr flux files.

2. I only have access to the Battistoni fluxes for electron neutrinos.
Therefore I will assume that the muon neutrino flux is twice as big which seems
to be roughly correct based on the paper "The Atmospheric Neutrino Flux Below
100 MeV: The FLUKA Results".

To run it:

    $ ./apply-atmospheric-oscillations --flux-dir ~/atm-production/fluxes/irc01_plus_lowe \
                                       --osc-dir ~/atm-production/nucraft_osc
"""
from __future__ import print_function, division
import numpy as np

def pdg_code_to_string(pdg):
    if pdg == 12:
        return "nue"
    elif pdg == -12:
        return "nbe"
    elif pdg == 14:
        return "num"
    elif pdg == -14:
        return "nbm"
    elif pdg == 16:
        return "nut"
    elif pdg == -16:
        return "nbt"
    else:
        raise ValueError("Unknown pdg code %i" % pdg)

if __name__ == '__main__':
    import argparse
    import matplotlib.pyplot as plt
    from os.path import split, splitext, join
    from scipy.interpolate import RectBivariateSpline
    from scipy.integrate import dblquad, nquad
    import sys

    parser = argparse.ArgumentParser("combine low energy and high energy atmospheric fluxes and apply oscillations")
    parser.add_argument("--flux-dir",required=True,help="atmospheric production directory")
    parser.add_argument("--osc-dir",required=True,help="directory with atmospheric oscillations")
    args = parser.parse_args()

    # Energy and cos(zenith) bins for the fluxes from Barr
    barr_cos_theta_bins = np.linspace(-1,1,21)
    barr_energy_bins = np.logspace(-1,1,41)

    # Final energy bins
    # We go here from 10 MeV -> 10 GeV with 20 points per decade like the Barr
    # fluxes
    energy_bins = np.logspace(-2,1,61)
    cos_theta_bins = barr_cos_theta_bins.copy()

    # Read in Battistoni fluxes
    nue_battistoni = np.genfromtxt(join(args.flux_dir,"lowe/nue.dat"))
    nbe_battistoni = np.genfromtxt(join(args.flux_dir,"lowe/nbe.dat"))

    # Convert Battistoni fluxes from MeV -> GeV
    nue_battistoni[:,0] /= 1000.0
    nue_battistoni[:,1] *= 1000.0
    nbe_battistoni[:,0] /= 1000.0
    nbe_battistoni[:,1] *= 1000.0

    # Convert Battistoni fluxes from cm^-2 -> m^-2
    nue_battistoni[:,1] *= 100.0**2
    nbe_battistoni[:,1] *= 100.0**2

    # Assume muon neutrino flux below 100 MeV is twice that of electron
    # neutrino flux
    num_battistoni = nue_battistoni.copy()
    num_battistoni[:,1] *= 2
    nbm_battistoni = nbe_battistoni.copy()
    nbm_battistoni[:,1] *= 2

    # Read in Barr fluxes
    nue_barr = np.genfromtxt(join(args.flux_dir,"irc01/fmax20_i0403z.sno_nue"))
    nbe_barr = np.genfromtxt(join(args.flux_dir,"irc01/fmax20_i0403z.sno_nbe"))
    num_barr = np.genfromtxt(join(args.flux_dir,"irc01/fmax20_i0403z.sno_num"))
    nbm_barr = np.genfromtxt(join(args.flux_dir,"irc01/fmax20_i0403z.sno_nbm"))

    # Read in oscillation probabilities
    nue_osc_prob = np.genfromtxt(join(args.osc_dir,"nue_osc_prob.txt"))
    nbe_osc_prob = np.genfromtxt(join(args.osc_dir,"nbe_osc_prob.txt"))
    num_osc_prob = np.genfromtxt(join(args.osc_dir,"num_osc_prob.txt"))
    nbm_osc_prob = np.genfromtxt(join(args.osc_dir,"nbm_osc_prob.txt"))

    # Reshape oscillation probability arrays. They should now look like:
    #
    #     [[[0.01,  -1,   P(nue), P(num), P(nut)],
    #       [0.01,  -0.8, P(nue), P(num), P(nut)],
    #                ...
    #       [0.01,  1.0,  P(nue), P(num), P(nut)]],
    #      [[0.02,  -1,   P(nue), P(num), P(nut)],
    #       [0.02,  -0.8, P(nue), P(num), P(nut)],
    #                ...
    #       [0.02,  1.0,  P(nue), P(num), P(nut)]],
    #                ...
    #      [[10.0,  -1,   P(nue), P(num), P(nut)],
    #       [10.0,  -0.8, P(nue), P(num), P(nut)],
    #                ...
    #       [10.0,  1.0,  P(nue), P(num), P(nut)]]]
    shape0 = len(np.unique(nue_osc_prob[:,0]))
    nue_osc_prob = nue_osc_prob.reshape((shape0,-1,5))
    nbe_osc_prob = nbe_osc_prob.reshape((shape0,-1,5))
    num_osc_prob = num_osc_prob.reshape((shape0,-1,5))
    nbm_osc_prob = nbm_osc_prob.reshape((shape0,-1,5))

    # convert dn/dlnE -> dn/dE
    nue_barr[:,2] /= nue_barr[:,0]
    nbe_barr[:,2] /= nbe_barr[:,0]
    num_barr[:,2] /= num_barr[:,0]
    nbm_barr[:,2] /= nbm_barr[:,0]

    # Reshape Barr flux arrays. They should now look like:
    #
    #     [[[0.1,  -1,   flux,...],
    #       [0.2,  -1,   flux,...],
    #               ...
    #       [10.0, -1,   flux,...]],
    # 
    #      [[0.1,  -0.9, flux,...],
    #       [0.2,  -0.9, flux,...],
    #              ...
    #       [10.0, -0.9, flux,...]],
    #              ...
    #              ...
    #      [[0.1,   1.0, flux,...],
    #       [0.2,   1.0, flux,...],
    #              ...
    #       [10.0,  1.0, flux,...]]]
    shape1 = len(np.unique(nue_barr[:,0]))
    nue_barr = nue_barr.reshape((-1,shape1,5))
    nbe_barr = nbe_barr.reshape((-1,shape1,5))
    num_barr = num_barr.reshape((-1,shape1,5))
    nbm_barr = nbm_barr.reshape((-1,shape1,5))

    nue_zenith_dist_x = nue_barr[:,0,1].copy()
    nue_zenith_dist_y = nue_barr[:,0,2].copy()
    nue_zenith_dist_y /= np.trapz(nue_zenith_dist_y,x=nue_zenith_dist_x)
    nbe_zenith_dist_x = nbe_barr[:,0,1].copy()
    nbe_zenith_dist_y = nbe_barr[:,0,2].copy()
    nbe_zenith_dist_y /= np.trapz(nbe_zenith_dist_y,x=nbe_zenith_dist_x)
    num_zenith_dist_x = num_barr[:,0,1].copy()
    num_zenith_dist_y = num_barr[:,0,2].copy()
    num_zenith_dist_y /= np.trapz(num_zenith_dist_y,x=num_zenith_dist_x)
    nbm_zenith_dist_x = nbm_barr[:,0,1].copy()
    nbm_zenith_dist_y = nbm_barr[:,0,2].copy()
    nbm_zenith_dist_y /= np.trapz(nbm_zenith_dist_y,x=nbm_zenith_dist_x)

    nue_battistoniz = np.empty((nue_barr.shape[0],nue_battistoni.shape[0],4))
    nue_battistoniz[:,:,0] = nue_battistoni[:,0]
    nue_battistoniz[:,:,1] = nue_zenith_dist_x[:,np.newaxis]
    nue_battistoniz[:,:,2] = nue_battistoni[:,1]*nue_zenith_dist_y[:,np.newaxis]/(2*np.pi)
    nue_battistoniz[:,:,3] = nue_battistoniz[:,:,2]*nue_battistoniz[:,:,0]

    nbe_battistoniz = np.empty((nbe_barr.shape[0],nbe_battistoni.shape[0],4))
    nbe_battistoniz[:,:,0] = nbe_battistoni[:,0]
    nbe_battistoniz[:,:,1] = nbe_zenith_dist_x[:,np.newaxis]
    nbe_battistoniz[:,:,2] = nbe_battistoni[:,1]*nbe_zenith_dist_y[:,np.newaxis]/(2*np.pi)
    nbe_battistoniz[:,:,3] = nbe_battistoniz[:,:,2]*nbe_battistoniz[:,:,0]

    num_battistoniz = np.empty((num_barr.shape[0],num_battistoni.shape[0],4))
    num_battistoniz[:,:,0] = num_battistoni[:,0]
    num_battistoniz[:,:,1] = num_zenith_dist_x[:,np.newaxis]
    num_battistoniz[:,:,2] = num_battistoni[:,1]*num_zenith_dist_y[:,np.newaxis]/(2*np.pi)
    num_battistoniz[:,:,3] = num_battistoniz[:,:,2]*num_battistoniz[:,:,0]

    nbm_battistoniz = np.empty((nbm_barr.shape[0],nbm_battistoni.shape[0],4))
    nbm_battistoniz[:,:,0] = nbm_battistoni[:,0]
    nbm_battistoniz[:,:,1] = nbm_zenith_dist_x[:,np.newaxis]
    nbm_battistoniz[:,:,2] = nbm_battistoni[:,1]*nbm_zenith_dist_y[:,np.newaxis]/(2*np.pi)
    nbm_battistoniz[:,:,3] = nbm_battistoniz[:,:,2]*nbm_battistoniz[:,:,0]

    nue_total = np.empty((nue_barr.shape[0],nue_barr.shape[1]+nue_battistoniz.shape[1],4))
    nue_total[:,:nue_battistoniz.shape[1],:] = nue_battistoniz
    nue_total[:,nue_battistoniz.shape[1]:,:] = nue_barr[:,:,:4]

    nbe_total = np.empty((nbe_barr.shape[0],nbe_barr.shape[1]+nbe_battistoniz.shape[1],4))
    nbe_total[:,:nbe_battistoniz.shape[1],:] = nbe_battistoniz
    nbe_total[:,nbe_battistoniz.shape[1]:,:] = nbe_barr[:,:,:4]

    num_total = np.empty((num_barr.shape[0],num_barr.shape[1]+num_battistoniz.shape[1],4))
    num_total[:,:num_battistoniz.shape[1],:] = num_battistoniz
    num_total[:,num_battistoniz.shape[1]:,:] = num_barr[:,:,:4]

    nbm_total = np.empty((nbm_barr.shape[0],nbm_barr.shape[1]+nbm_battistoniz.shape[1],4))
    nbm_total[:,:nbm_battistoniz.shape[1],:] = nbm_battistoniz
    nbm_total[:,nbm_battistoniz.shape[1]:,:] = nbm_barr[:,:,:4]

    flux_arrays = {12: nue_total, -12: nbe_total,
                   14: num_total, -14: nbm_total}

    # Save unoscillated fluxes
    for p in [12,-12,14,-14]:
        np.savetxt("fmax20_i0403z_plus_lowe.sno_%s" % pdg_code_to_string(p),flux_arrays[p].reshape((flux_arrays[p].shape[0]*flux_arrays[p].shape[1],-1)),
                   fmt='%10.4f %10.3f %20.8g %20.8g',
                   header="Energy (GeV), Cos(zenith), Flux (m^-2 s^-1 Gev^-1 steradian^-1), Flux (dn/dlnE)")

    osc_prob = {12: nue_osc_prob, -12: nbe_osc_prob,
                14: num_osc_prob, -14: nbm_osc_prob}

    def get_flux_interp(nu_type):
        flux_array = flux_arrays[nu_type]
        flux = flux_array[:,:,2]
        e = flux_array[0,:,0]
        cos_theta = flux_array[:,0,1]

        # Use first order splines here since we need to enforce the fact that
        # the flux is not negative
        return RectBivariateSpline(e,cos_theta,flux.T,kx=1,ky=1)

    def get_osc_interp(nu_type1,nu_type2):
        prob = osc_prob[nu_type1]
        e = prob[:,0,0]
        cos_theta = prob[0,:,1]

        if np.sign(nu_type1) != np.sign(nu_type2):
            raise ValueError("asking for oscillation probability from %s -> %s!" % tuple(map(pdg_code_to_string,(nu_type1,nu_type2))))

        # Use second order splines here since oscillation probabilities are
        # pretty smooth
        if abs(nu_type2) == 12:
            return RectBivariateSpline(e,cos_theta,prob[:,:,2],kx=2,ky=2)
        elif abs(nu_type2) == 14:
            return RectBivariateSpline(e,cos_theta,prob[:,:,3],kx=2,ky=2)
        elif abs(nu_type2) == 16:
            return RectBivariateSpline(e,cos_theta,prob[:,:,4],kx=2,ky=2)

    def get_oscillated_flux(nu_type1, nu_type2, elo, emid, ehi, zlo, zmid, zhi):
        """
        Returns the average flux of neutrinos coming from `nu_type1` neutrinos
        oscillating to `nu_type2` neutrinos for neutrinos in the bin between
        energies elo and ehi and with a cosine of the zenith angle between zlo
        and zhi.

        Note: We also pass the midpoint of the bin for both energy and
        cos(theta) as `emid` and `zmid`. The reason for this is that at least
        for the Barr fluxes, the midpoint of the bin was chosen as the midpoint
        of the log of the left and right edge instead of the actual middle of
        the bin. Therefore to be consistent when interpolating these fluxes we
        should use the same value.

        The returned flux has units of 1/(m^2 sec steradian GeV).
        """
        f_flux1 = get_flux_interp(nu_type1)
        f_osc1 = get_osc_interp(nu_type1,nu_type2)

        def f_flux(e,z):
            # Not sure exactly how scipy deals with boundary issues in
            # RectBivariateSpline so we just make sure everything is within the
            # range of the values given in the Barr and Battistoni fluxes
            if e < 0.0106:
                e = 0.0106
            if z < -0.95:
                z = -0.95
            if z > 0.95:
                z = 0.95
            return f_flux1(e,z)

        def f_osc(e,z):
            # FIXME: Should remove this
            if e < 0.02:
                e = 0.02
            return f_osc1(e,z)

        # Here we have two options for calculating the oscillated flux:
        #
        # 1. We interpolate the value of the flux and then multiply by an
        # average oscillation probability. This is more correct in the sense
        # that the original flux values are actually a histogram and so it
        # makes sense to just sample the flux at the middle of the bin.
        #
        # 2. We integrate the product of the flux and the oscillation
        # probability over the whole bin. This is technically more correct
        # since it will weight the oscillation probability by the flux, but it
        # may have issues near the edge of the first and last bin since we have
        # to extrapolate.
        #
        # For now, we do 1 since it is much faster. To do the second you can
        # just comment the next line.
        return f_flux(emid,zmid)*f_osc1.integral(elo,ehi,zlo, zhi)/(ehi-elo)/(zhi-zlo)
        return nquad(lambda e, z: f_flux(e,z)*f_osc(e,z), [(elo,ehi), (zlo, zhi)], opts={'limit':1000,'epsrel':1e-3})[0]/(ehi-elo)/(zhi-zlo)

    data = {nu_type: np.zeros((cos_theta_bins.shape[0]-1,energy_bins.shape[0]-1,4)) for nu_type in [12,-12,14,-14,16,-16]}
    for nu_type1 in [12,-12,14,-14]:
        if nu_type1 > 0:
            nu_type2s = [12,14,16]
        else:
            nu_type2s = [-12,-14,-16]
        for nu_type2 in nu_type2s:
            print("Calculating neutrino flux for %s -> %s" % tuple(map(pdg_code_to_string,(nu_type1,nu_type2))))
            for i, (elo, ehi) in enumerate(zip(energy_bins[:-1],energy_bins[1:])):
                for j, (zlo, zhi) in enumerate(zip(cos_theta_bins[:-1],cos_theta_bins[1:])):
                    print("\r%i/%i" % (i*(len(cos_theta_bins)-1)+j+1,(len(cos_theta_bins)-1)*(len(energy_bins)-1)),end="")
                    sys.stdout.flush()
                    # We'll follow the Barr convention here and take the midpoint in log space
                    emid = np.exp((np.log(elo)+np.log(ehi))/2)
                    zmid = (zlo + zhi)/2
                    data[nu_type2][j,i,0] = emid
                    data[nu_type2][j,i,1] = zmid
                    data[nu_type2][j,i,2] += get_oscillated_flux(nu_type1,nu_type2,elo,emid,ehi,zlo,zmid,zhi)
            print()

    for nu_type2 in data:
        for i, (elo, ehi) in enumerate(zip(energy_bins[:-1],energy_bins[1:])):
            for j, (zlo, zhi) in enumerate(zip(cos_theta_bins[:-1],cos_theta_bins[1:])):
                data[nu_type2][j,i,3] = data[nu_type2][j,i,2]*data[nu_type2][j,i,0]

    for p in [12,-12,14,-14,16,-16]:
        np.savetxt("fmax20_i0403z_plus_lowe.sno_%s.osc" % pdg_code_to_string(p),data[p].reshape((data[p].shape[0]*data[p].shape[1],-1)),
                    fmt='%10.4f %10.3f %20.8g %20.8g',
                   header="Energy (GeV), Cos(zenith), Flux (m^-2 s^-1 Gev^-1 steradian^-1), Flux (dn/dlnE), Junk")