aboutsummaryrefslogtreecommitdiff
path: root/macros/snogen
blob: d96a87f5dbb4b65d53d6d4a9d4167890b7409f9e (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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#!/usr/bin/env python
# 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/>.

"""
snogen is a script to generate SNOMAN command files.

Example:

    # simulate 100 1 GeV muons
    $ snogen -n 100 -e 1000 -p mu_minus | /path/to/snoman.exe

The output filename can be specified on the command line:

    $ snogen -o [output filename]

By default it will be [particle name]_[energy]_[number of events].zdab.
"""
from __future__ import print_function, division
import string

class MyTemplate(string.Template):
    delimiter = '@'

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser("generate a SNOMAN command file")
    parser.add_argument("-p", "--particle-id", default="e_minus",
                        help="particle id")
    parser.add_argument("-e", "--mc-energy", default=500.0, type=float,
                        help="total energy")
    parser.add_argument("-n", "--num-events", default=100, type=int,
                        help="number of events")
    parser.add_argument("-o", "--output", default=None,
                        help="output file name")
    args = parser.parse_args()

    if args.output is None:
        output = "%s_%.0f_%i.zdab" % (args.particle_id,args.mc_energy,args.num_events)
    else:
        output = args.output

    with open("mc_run_10000_macro.cmd") as f:
        s = f.read()

    template = MyTemplate(s)
    print(template.safe_substitute(particle_id=args.particle_id,mc_energy=args.mc_energy,num_events=args.num_events,output=output))
n class="vm">__name__ == '__main__': import argparse import numpy as np import h5py import pandas as pd import itertools from sddm import IDP_E_MINUS, IDP_MU_MINUS, SNOMAN_MASS from sddm.plot import plot_hist, plot_legend, get_stats, despine, iqr_std_err, iqr_std, quantile_error, q90_err, q90, median, median_err, std_err 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 plots") args = parser.parse_args() setup_matplotlib(args.save) import matplotlib.pyplot as plt # Read in all the data. # # Note: We have to add the filename as a column to each dataset since this # script is used to analyze MC data which all has the same run number. ev = pd.concat([pd.read_hdf(filename, "ev").assign(filename=filename) for filename in args.filenames],ignore_index=True) fits = pd.concat([pd.read_hdf(filename, "fits").assign(filename=filename) for filename in args.filenames],ignore_index=True) mcgn = pd.concat([pd.read_hdf(filename, "mcgn").assign(filename=filename) for filename in args.filenames],ignore_index=True) # get rid of 2nd events like Michel electrons ev = ev.sort_values(['run','gtid']).groupby(['filename','evn'],as_index=False).nth(0) # Now, we merge all three datasets together to produce a single # dataframe. To do so, we join the ev dataframe with the mcgn frame # on the evn column, and then join with the fits on the run and # gtid columns. # # At the end we will have a single dataframe with one row for each # fit, i.e. it will look like: # # >>> data # run gtid nhit, ... mcgn_x, mcgn_y, mcgn_z, ..., fit_id1, fit_x, fit_y, fit_z, ... # # Before merging, we prefix the primary seed track table with mcgn_ # and the fit table with fit_ just to make things easier. # Prefix track and fit frames mcgn = mcgn.add_prefix("mcgn_") fits = fits.add_prefix("fit_") # merge ev and mcgn on evn data = ev.merge(mcgn,left_on=['filename','evn'],right_on=['mcgn_filename','mcgn_evn']) # merge data and fits on run and gtid data = data.merge(fits,left_on=['filename','run','gtid'],right_on=['fit_filename','fit_run','fit_gtid']) # For this script, we only want the single particle fit results data = data[(data.fit_id2 == 0) & (data.fit_id3 == 0)] # Select only the best fit for a given run, gtid, and particle # combo data = data.sort_values('fit_fmin').groupby(['filename','run','gtid','fit_id1','fit_id2','fit_id3'],as_index=False).nth(0).reset_index(level=0,drop=True) # calculate true kinetic energy mass = [SNOMAN_MASS[id] for id in data['mcgn_id'].values] data['T'] = data['mcgn_energy'].values - mass data['dx'] = data['fit_x'].values - data['mcgn_x'].values data['dy'] = data['fit_y'].values - data['mcgn_y'].values data['dz'] = data['fit_z'].values - data['mcgn_z'].values data['dT'] = data['fit_energy1'].values - data['T'].values true_dir = np.dstack((data['mcgn_dirx'],data['mcgn_diry'],data['mcgn_dirz'])).squeeze() dir = np.dstack((np.sin(data['fit_theta1'])*np.cos(data['fit_phi1']), np.sin(data['fit_theta1'])*np.sin(data['fit_phi1']), np.cos(data['fit_theta1']))).squeeze() data['theta'] = np.degrees(np.arccos((true_dir*dir).sum(axis=-1))) # only select fits which have at least 2 fits data = data.groupby(['filename','run','gtid']).filter(lambda x: len(x) > 1) data_true = data[data['fit_id1'] == data['mcgn_id']] data_e = data[data['fit_id1'] == IDP_E_MINUS] data_mu = data[data['fit_id1'] == IDP_MU_MINUS] data_true = data_true.set_index(['filename','run','gtid']) data_e = data_e.set_index(['filename','run','gtid']) data_mu = data_mu.set_index(['filename','run','gtid']) data_true['ratio'] = data_mu['fit_fmin']-data_e['fit_fmin'] data_true['te'] = data_e['fit_time'] data_true['tm'] = data_mu['fit_time'] data_true['Te'] = data_e['fit_energy1'] # 100 bins between 50 MeV and 1 GeV bins = np.arange(50,1000,100) markers = itertools.cycle(('o', 'v')) fig3, ax3 = plt.subplots(3,1,num=3,sharex=True) fig4, ax4 = plt.subplots(3,1,num=4,sharex=True) for id in [IDP_E_MINUS, IDP_MU_MINUS]: events = data_true[data_true['mcgn_id'] == id] pd_bins = pd.cut(events['T'],bins) dT = events.groupby(pd_bins)['dT'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err]) dx = events.groupby(pd_bins)['dx'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err]) dy = events.groupby(pd_bins)['dy'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err]) dz = events.groupby(pd_bins)['dz'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err]) theta = events.groupby(pd_bins)['theta'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err,q90,q90_err]) label = 'Muon' if id == IDP_MU_MINUS else 'Electron' T = (bins[1:] + bins[:-1])/2 marker = markers.next() plt.figure(1) plt.errorbar(T,dT['median']*100/T,yerr=dT['median_err']*100/T,fmt=marker,label=label) plt.figure(2) plt.errorbar(T,dT['iqr_std']*100/T,yerr=dT['iqr_std_err']*100/T,fmt=marker,label=label) ax3[0].errorbar(T,dx['median'],yerr=dx['median_err'],fmt=marker,label=label) ax3[1].errorbar(T,dy['median'],yerr=dy['median_err'],fmt=marker,label=label) ax3[2].errorbar(T,dz['median'],yerr=dz['median_err'],fmt=marker,label=label) ax4[0].errorbar(T,dx['iqr_std'],yerr=dx['iqr_std_err'],fmt=marker,label=label) ax4[1].errorbar(T,dy['iqr_std'],yerr=dy['iqr_std_err'],fmt=marker,label=label) ax4[2].errorbar(T,dz['iqr_std'],yerr=dz['iqr_std_err'],fmt=marker,label=label) plt.figure(5) plt.errorbar(T,theta['std'],yerr=theta['std_err'],fmt=marker,label=label) plt.figure(6) plt.scatter(events['Te'],events['ratio'],marker=marker,label=label) fig = plt.figure(1) despine(fig,trim=True) plt.xlabel("Kinetic Energy (MeV)") plt.ylabel(r"Energy bias (\%)") plt.legend() plt.tight_layout() fig = plt.figure(2) despine(fig,trim=True) plt.xlabel("Kinetic Energy (MeV)") plt.ylabel(r"Energy resolution (\%)") plt.legend() plt.tight_layout() ax3[0].set_ylabel("X") ax3[0].set_ylim((-5,5)) ax3[1].set_ylabel("Y") ax3[1].set_ylim((-5,5)) ax3[2].set_xlabel("Kinetic Energy (MeV)") ax3[2].set_ylabel("Z") ax3[2].set_ylim((-5,5)) despine(ax=ax3[0],trim=True) despine(ax=ax3[1],trim=True) despine(ax=ax3[2],trim=True) h,l = ax3[0].get_legend_handles_labels() fig3.legend(h,l,loc='upper right') fig3.subplots_adjust(right=0.75) fig3.tight_layout() fig3.subplots_adjust(top=0.9) ax4[0].set_ylabel("X") ax4[0].set_ylim((0,ax4[0].get_ylim()[1])) ax4[1].set_ylabel("Y") ax4[1].set_ylim((0,ax4[1].get_ylim()[1])) ax4[2].set_xlabel("Kinetic Energy (MeV)") ax4[2].set_ylabel("Z") ax4[2].set_ylim((0,ax4[2].get_ylim()[1])) despine(ax=ax4[0],trim=True) despine(ax=ax4[1],trim=True) despine(ax=ax4[2],trim=True) h,l = ax4[0].get_legend_handles_labels() fig4.legend(h,l,loc='upper right') fig4.subplots_adjust(right=0.75) fig4.tight_layout() fig4.subplots_adjust(top=0.9) fig = plt.figure(5) despine(fig,trim=True) plt.xlabel("Kinetic Energy (MeV)") plt.ylabel("Angular resolution (deg)") plt.ylim((0,plt.gca().get_ylim()[1])) plt.legend() plt.tight_layout() fig = plt.figure(6) plt.xticks(range(0,1250,100)) plt.hlines(0,0,1200,color='k',linestyles='--',alpha=0.5) despine(fig,trim=True) plt.xlabel("Reconstructed Electron Energy (MeV)") plt.ylabel(r"Log Likelihood Ratio (e/$\mu$)") plt.legend() plt.tight_layout() if args.save: fig = plt.figure(1) plt.savefig("energy_bias.pdf") plt.savefig("energy_bias.eps") fig = plt.figure(2) plt.savefig("energy_resolution.pdf") plt.savefig("energy_resolution.eps") fig = plt.figure(3) plt.savefig("position_bias.pdf") plt.savefig("position_bias.eps") fig = plt.figure(4) plt.savefig("position_resolution.pdf") plt.savefig("position_resolution.eps") fig = plt.figure(5) plt.savefig("angular_resolution.pdf") plt.savefig("angular_resolution.eps") fig = plt.figure(6) plt.savefig("likelihood_ratio.pdf") plt.savefig("likelihood_ratio.eps") else: plt.figure(1) plt.title("Energy Bias") plt.figure(2) plt.title("Energy Resolution") plt.figure(3) fig3.suptitle("Position Bias (cm)") plt.figure(4) fig4.suptitle("Position Resolution (cm)") plt.figure(5) plt.title("Angular Resolution") plt.figure(6) plt.title("Log Likelihood Ratio vs Reconstructed Electron Energy") plt.show()