aboutsummaryrefslogtreecommitdiff
path: root/src/random.h
blob: e2cba7d0ce52d9c0d56bf7870318c2d1f56dcfca (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
/* Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
 *
 * This program is free software: you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.

 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.

 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#ifndef RANDOM_H
#define RANDOM_H

#include "mt19937ar.h"
#include <stdlib.h> /* for size_t */

double randn(void);
void rand_ball(double *pos, double radius);
void rand_sphere(double *dir);
void ran_choose_weighted(int *dest, double *w, size_t k, int *src, size_t n);

#endif
#fff0f0 } /* Literal.String.Char */ .highlight .dl { color: #dd2200; background-color: #fff0f0 } /* Literal.String.Delimiter */ .highlight .sd { color: #dd2200; background-color: #fff0f0 } /* Literal.String.Doc */ .highlight .s2 { color: #dd2200; background-color: #fff0f0 } /* Literal.String.Double */ .highlight .se { color: #0044dd; background-color: #fff0f0 } /* Literal.String.Escape */ .highlight .sh { color: #dd2200; background-color: #fff0f0 } /* Literal.String.Heredoc */ .highlight .si { color: #3333bb; background-color: #fff0f0 } /* Literal.String.Interpol */ .highlight .sx { color: #22bb22; background-color: #f0fff0 } /* Literal.String.Other */ .highlight .sr { color: #008800; background-color: #fff0ff } /* Literal.String.Regex */ .highlight .s1 { color: #dd2200; background-color: #fff0f0 } /* Literal.String.Single */ .highlight .ss { color: #aa6600; background-color: #fff0f0 } /* Literal.String.Symbol */ .highlight .bp { color: #003388 } /* Name.Builtin.Pseudo */ .highlight .fm { color: #0066bb; font-weight: bold } /* Name.Function.Magic */ .highlight .vc { color: #336699 } /* Name.Variable.Class */ .highlight .vg { color: #dd7700 } /* Name.Variable.Global */ .highlight .vi { color: #3333bb } /* Name.Variable.Instance */ .highlight .vm { color: #336699 } /* Name.Variable.Magic */ .highlight .il { color: #0000DD; font-weight: bold } /* Literal.Number.Integer.Long */
#!/usr/bin/env python3
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.
"""
Script to plot reconstructed quantities for instrumentals. To run it just run:

    $ ./plot-dc [list of fit results]

This will produce corner plots showing the distribution of the high level
variables used in the contamination analysis for all the different instrumental
backgrounds and external muons.
"""
from __future__ import print_function, division
import numpy as np
from scipy.stats import iqr, poisson
from matplotlib.lines import Line2D
from scipy.stats import iqr, norm, beta
from scipy.special import spence
from itertools import chain

particle_id = {20: 'e', 22: r'\mu'}

def plot_hist2(df, muons=False):
    for id, df_id in sorted(df.groupby('id')):
        if id == 20:
            plt.subplot(2,3,1)
        elif id == 22:
            plt.subplot(2,3,2)
        elif id == 2020:
            plt.subplot(2,3,4)
        elif id == 2022:
            plt.subplot(2,3,5)
        elif id == 2222:
            plt.subplot(2,3,6)

        if muons:
            plt.hist(np.log10(df_id.ke.values/1000), bins=np.linspace(0,4.5,100), histtype='step')
            plt.xlabel("log10(Energy (GeV))")
        else:
            bins = np.logspace(np.log10(20),np.log10(10e3),21)
            plt.hist(df_id.ke.values, bins=bins, histtype='step')
            plt.gca().set_xscale("log")
            major = np.array([10,100,1000,10000])
            minor = np.unique(list(chain(*list(range(i,i*10,i) for i in major[:-1]))))
            minor = np.setdiff1d(minor,major)
            plt.gca().set_xticks(major)
            plt.gca().set_xticks(minor,minor=True)
            plt.xlabel("Energy (MeV)")
        plt.title('$' + ''.join([particle_id[int(''.join(x))] for x in grouper(str(id),2)]) + '$')

    if len(df):
        plt.tight_layout()

def plot_hist(df, muons=False):
    for id, df_id in sorted(df.groupby('id')):
        if id == 20:
            plt.subplot(3,4,1)
        elif id == 22:
            plt.subplot(3,4,2)
        elif id == 2020:
            plt.subplot(3,4,5)
        elif id == 2022:
            plt.subplot(3,4,6)
        elif id == 2222:
            plt.subplot(3,4,7)
        elif id == 202020:
            plt.subplot(3,4,9)
        elif id == 202022:
            plt.subplot(3,4,10)
        elif id == 202222:
            plt.subplot(3,4,11)
        elif id == 222222:
            plt.subplot(3,4,12)

        if muons:
            plt.hist(np.log10(df_id.ke.values/1000), bins=np.linspace(0,4.5,100), histtype='step')
            plt.xlabel("log10(Energy (GeV))")
        else:
            plt.hist(df_id.ke.values, bins=np.linspace(20,10e3,100), histtype='step')
            plt.xlabel("Energy (MeV)")
        plt.title(str(id))

    if len(df):
        plt.tight_layout()

if __name__ == '__main__':
    import argparse
    import numpy as np
    import pandas as pd
    import sys
    import h5py
    from sddm.plot_energy import *
    from sddm.plot import *
    from sddm import setup_matplotlib

    parser = argparse.ArgumentParser("plot fit results")
    parser.add_argument("filenames", nargs='+', help="input files")
    parser.add_argument("--save", action='store_true', default=False, help="save corner plots for backgrounds")
    args = parser.parse_args()

    setup_matplotlib(args.save)

    import matplotlib.pyplot as plt

    # Loop over runs to prevent using too much memory
    evs = []
    data_filenames = [filename for filename in args.filenames if 'SNO' in filename]
    mc_filenames = [filename for filename in args.filenames if 'SNO' not in filename]
    if len(data_filenames):
        rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in data_filenames],ignore_index=True)
        for run, df in rhdr.groupby('run'):
            evs.append(get_events(df.filename.values, merge_fits=True))
    if len(mc_filenames):
        for filename in mc_filenames:
            evs.append(get_events([filename], merge_fits=True, mc=True))
    ev = pd.concat([ev for ev in evs if len(ev) > 0])

    # Hack to get rid of flasher and muon events in the breakdown sample.
    ev.breakdown &= ev.r_psup < 0.9

    ev = ev[ev.prompt & ~np.isnan(ev.fmin)]
    ev = ev[ev.ke > 20]

    #with pd.option_context('display.max_rows', None, 'display.max_columns', None):
    #    print("Noise events")
    #    print(ev[ev.noise][['psi','x','y','z','id1','id2']])
    #    print("Muons")
    #    print(ev[ev.muon][['psi','r','id1','id2','id3','energy1','energy2','energy3']])
    #    print("Neck")
    #    print(ev[ev.neck & ev.psi < 6][['psi','r','id1','cos_theta']])
    #    print("Flashers")
    #    print(ev[ev.flasher & ev.udotr > 0])
    #    print("Signal")
    #    print(ev[ev.signal])

    # save as PDF b/c EPS doesn't support alpha values
    if args.save:
        plot_corner_plot(ev[ev.breakdown],"Breakdowns",save="breakdown_corner_plot")
        plot_corner_plot(ev[ev.muon],"Muons",save="muon_corner_plot")
        plot_corner_plot(ev[ev.flasher],"Flashers",save="flashers_corner_plot")
        plot_corner_plot(ev[ev.neck],"Neck",save="neck_corner_plot")
        plot_corner_plot(ev[ev.noise],"Noise",save="noise_corner_plot")
        plot_corner_plot(ev[ev.signal],"Signal",save="signal_corner_plot")
    else:
        plot_corner_plot(ev[ev.breakdown],"Breakdowns")
        plot_corner_plot(ev[ev.muon],"Muons")
        plot_corner_plot(ev[ev.flasher],"Flashers")
        plot_corner_plot(ev[ev.neck],"Neck")
        plot_corner_plot(ev[ev.noise],"Noise")
        plot_corner_plot(ev[ev.signal],"Signal")

    fig = plt.figure()
    plot_hist2(ev[ev.flasher])
    despine(fig,trim=True)
    plt.suptitle("Flashers")
    fig = plt.figure()
    plot_hist2(ev[ev.muon],muons=True)
    despine(fig,trim=True)
    plt.suptitle("Muons")
    plt.show()