#!/usr/bin/env python """ usage : aashowerfit.py -f input_file Aashowerfit is a reconstruction program for cascade events. The options can be set via the command line (see below), or via the following environment variables AA_INFILE, AA_OUTFILE, AA_DETFILE, AA_LIBDIR, AA_NEVE, and AA_VERB. Environment variables are intended for use in batch jobs and take precident over command line arugments. The reconstructed tracks are appended to the Evt.trks container in the aanet event. This means that previous reconstructions are not overwritten by default (see option -c). Aashowerfit will write the following tracks (in order): o 1 starting point of the M-estimator fit , rec_type = 101, rec_stages = {1} o 1 (M-est) pos+time fit , rec-type = 101, rec_stages = {1,2} o 11 sub-optimal likelihood fits , rec-type = 101, rec_stages = {1,2,3} o 1 optimal likelihood fit (best fit) , rec-type = 101, rec_stages = {1,2,3,4} (In fact the 12 likelihood fits are simply sorted by likelihood, but the best (i.e. last) one receives a diffent value of rec_stages for easy identificaion in analysis.) The likelihood fit tracks (rec_stage>=3) have: fitinf[0] = uncorrected energy, fitinf[1] = number of hits used in the fit options: -f -input_file : Input file in any format aanet can read. default=../data/km3_sim/km3_v4_nueCC_1.JTE.aareco.root -o -out_file: output default=./rec_files/ -d -detector_file : det or detx. Set this to "auto" to get it from the header of the input file. default=/afs/in2p3.fr/home/throng/km3net/detectors/km3net_jul13_90m.det -l : directory where to find aashowfit.so default=. -n : process this many events default=100000000 -e : only process events w/ MC-Enu greater than this default=0 -a : run mc analysis on the first N events default=10 -c : clear the trks container before writing result default=False -v : verbose (0,1 or 2) default=0 -t : report timing informtion every N events default=100 -r : rec type identifyer for written tracks default=101 -m : do hit merging default=True -h : display this message default=False -i : start python prompt when done default=False -q : experimental numpy minimizer default=True """ print "aashowerfit" import ROOT,os,sys,collections, array from math import * from copy import copy import numpy as np import pickle try : import aa except ImportError: print "Failed to import aa; Make sure AADIR is in your PYHTONPATH!" sys.exit() import rec options = aa.Options( __doc__, sys.argv[1:] ) #options.read_env() print options from ROOT import vector, Hit, Det, Evt, Trk, Vec, EventFile, Timer from ROOT import MestShowerPdf, DigitalShowerPdf, ShowerFit from ROOT import JppShowerPdf, JppShowerFit def todeg(rad): return rad*180./pi def torad(deg): return deg*pi/180. #----------------------------------------------------- # Fancy-smancy setting of members #----------------------------------------------------- for T in Trk, JppShowerPdf, JppShowerFit, DigitalShowerPdf, MestShowerPdf, ShowerFit : T.set = aa.__set #----------------------------------------------------- # Only R0 tune is included in this production version #----------------------------------------------------- tune = "R0" if tune == 'R0' : suffix = "R0" hit_type = "raw" muscale = 1.0 jshowerfile = ["J13p.dat", "J14p.dat"] estart = 5e5 tollerance = 1 rec_type = options.r def max_dist_func ( det, hits, pos ) : rmax, fmax = 800, 2.0 rmin, fmin = 300, 0.2 f = ROOT.hit_to_dom_ratio( det, hits, pos, 200 ) if f == 0 : # no doms in radius return rmax f = log10(f) if f < fmin : return rmin return rmin + (f-fmin) * (rmax-rmin)/fmax; #----------------------------------------- # input & output file #----------------------------------------- basefilename = os.path.basename(options.f) f = ROOT.EventFile( options.f ) f5d = ROOT.EventFile( "../data/km3_reco/" + basefilename + ".5dR0.root" ) f3d = ROOT.EventFile( "../data/km3_reco/" + basefilename + ".3dR0.root" ) #----------------------------------------------- # detector file. todo: get it from the evt file #----------------------------------------------- if options.d == "auto" : det = Det( f ) else : det = Det( options.d ) det.prnt() # For mc-analysis, we need the coord-origin. If it is mising, we # use the cg of the detector and hope for the best. if f.header.have_line("coord_origin") : coord_origin = f.header.coord_origin() else : coord_origin = det.get_cg() #----------------------------------------------- # Add our info to the header and write it. #----------------------------------------------- f.header["aashowerfit_detector"] = options.d #=================================================================================== # end of setup #=================================================================================== #----------------------------------------------- # Energy correction function #----------------------------------------------- def correct_energy( trk ): "2nd deg polynomial fitted to median of E/Etrue distribution" if ( trk.E > 10 and trk.E < 1e10 ) : x = log10(trk.E) f = 0.53972 + x * 0.16419 + x*x * -0.0168034; if f > 1e-3 : trk.E = trk.E / f; #-------------------------------------------------- # setup the ShowerFit'ers #-------------------------------------------------- mpdf1 = MestShowerPdf().set ( verb = options.v ) mest_fit = ShowerFit ( "mest1", mpdf1 ).set( verb = options.v ) mest_fit.fix_vars( 3,4,6 ) pdf5d = JppShowerPdf() fit5d = JppShowerFit("5d_pdf", pdf5d ).set ( verb = options.v, minimizer_name ='Minuit2', minimizer_algo ='MIGRAD', ) fit5d.fix_vars( 0,1,2,5, 6) #-------------------------------------------------- # Event Loop #-------------------------------------------------- iproc = 0 logfile = open("/sps/km3net/users/jseneca/aanet/rec/log_" + options.f.split("/")[-1], "w") logfile.write("Opened logfile for writing...") logfile.write("MC file: %s" %(options.f)) for evt in f: f5d.next() f3d.next() evt5d = f5d.evt evt3d = f3d.evt # minimum mc neutrino energy requirement. if options.e > 0 and \ evt.mc_trks.size()>0 and \ evt.mc_trks[0].E < options.e : continue logfile = open("/sps/km3net/users/jseneca/aanet/rec/log_" + options.f.split("/")[-1], "a") iproc+=1 # count actually processed (not skipped) events evttimer = Timer("0) events") if iproc > options.n : print 'requested number of events reached' break if options.v > 0 : print " === event ", evt.id , " ===" print " === event 3d", evt3d.id , " ===" print " === event 5d", evt5d.id , " ===" with Timer("1) calibration"): if hit_type == 'raw' : # Make the hit times more numerically palletable hit_t0 = min( h.t for h in evt.hits ) for h in evt.hits : h.t -= hit_t0 input_hits = evt.hits det.apply_fast( input_hits ) if hit_type == 'mc' : input_hits = evt.mc_hits det.apply( input_hits ) with Timer("merge"): if options.m : hits = ROOT.merge_hits ( input_hits ) for h in hits: h.a = 1 print "merged ", input_hits.size() - hits.size(), " of " , hits.size() else: hits = input_hits with Timer("coincidences"): coin_hits = ROOT.get_coincidences( hits, 2, 20 ) print " coin hits : ", coin_hits.size() with Timer("3) hit selection"): # hit with the largest-multiplicity cooincidence pivot_hit = Hit() for h in coin_hits : if h.a > pivot_hit.a : pivot_hit = h T0 = Trk().set( id = 1, pos = pivot_hit.pos, t = pivot_hit.t, E = estart, rec_type = rec_type, rec_stages = aa.make_vector([1]), comment = "start" ) mest_hits = ROOT.filter_hits( coin_hits, T0, -800, 800 ); if options.v > 1 : print(" --- *** mest hits *** --- ") print "mest hits size/purity = " , len(mest_hits) #, purity( mest_hits ) with Timer("4) vertex fit"): T1 = mest_fit.fit( T0, det, mest_hits, 1 ) T1.set( id = 2, comment = "mest", rec_type = rec_type, rec_stages = aa.make_vector( [1,2] ) ) if options.c : evt.trks.clear() evt.trks.push_back(T0) evt.trks.push_back(T1) #------------------------------------------------------------------ # likelihood fit, with starting-direction according to icosahedron #------------------------------------------------------------------ with Timer("5) final hit selection"): digi_hits = ROOT.filter_hits( input_hits, T1 , -100, 900, 800 ) fit5d.max_dist = max_dist_func( det, digi_hits, T1.pos ) print("max dist: %f, total n hits: %i" %(fit5d.max_dist, len(digi_hits) ) ) #not implemented in JppShowerPdf yet #pdf5d.force_precompute = True with Timer("6) likelihood space calulation"): if (fit5d.max_dist == 0): pdf5d.init_event( det, digi_hits ) else: pdf5d.init_event( det, digi_hits, T1.pos, fit5d.max_dist ) true_LL = pdf5d.eval_L( Trk().set( id = 1, pos = T1.pos, t = T1.t, E = evt.mc_trks[0].E, dir = evt.mc_trks[0].dir.normalize(), rec_type = 0 ) ) zzzscale = 0.005 zzscale = 1 zscale = 10 pscale = 50 neut_D = evt.mc_trks[0].dir.normalize() neut_E = evt.mc_trks[0].E lep_D = evt.mc_trks[1].dir.normalize() reco3d = evt3d.trks[-1] reco3d_D = reco3d.dir.normalize() reco3d_E = reco3d.E reco5d = evt5d.trks[-1] reco5d_D = reco5d.dir.normalize() reco5d_E = reco5d.E evals5d = [] for itrk, trk in enumerate(evt5d.trks): #print("rec type: %i" %trk.rec_type) if trk.rec_type == 205: evals5d.append(trk) minLL_D = Vec(0.0, 0.0, 0.0) #use reco Energy LL_5d = np.zeros((pscale, 2*pscale)) ThetaArray = np.linspace(0, pi, pscale) PhiArray = np.linspace(-pi, pi, 2*pscale) try: pickled_file = open(options.o + basefilename + "%iLL_TESTF" %(evt.id), 'r') LL_5d = pickle.load(pickled_file) test_dir = Vec().set_angles(ThetaArray[pscale/2], PhiArray[pscale/2]) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) if LL_5d[pscale/2][pscale/2] != pdf5d.eval_L(Ttest): print( "LL From file: ", LL_5d[pscale/2][pscale/2], "real: ", pdf5d.eval_L(Ttest) ) raise IOError except IOError: print("Re-making LL") for itheta, theta in enumerate(ThetaArray): if itheta%(pscale/10) == 0 : print("%i out of 10" %(itheta/(pscale/10))) for iphi, phi in enumerate(PhiArray): test_dir = Vec().set_angles(theta, phi) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) this_LL = 0 this_LL = pdf5d.eval_L(Ttest) LL_5d[itheta, iphi] = this_LL dump_file = open( options.o + basefilename + "%iLL_TESTF" %(evt.id), 'w') pickle.dump(LL_5d, dump_file) dump_file.close() ZLL_5d = np.zeros((pscale, pscale)) ZThetaArray = np.linspace(neut_D.theta() - torad(zscale), neut_D.theta() + torad(zscale), pscale) ZPhiArray = np.linspace(neut_D.phi() - torad(zscale), neut_D.phi() + torad(zscale), pscale) try: pickled_file = open( options.o + basefilename + "%iZLL_TESTF" %(evt.id), 'r') ZLL_5d = pickle.load(pickled_file) test_dir = Vec().set_angles(ZThetaArray[pscale/2], ZPhiArray[pscale/2]) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) if ZLL_5d[pscale/2][pscale/2] != pdf5d.eval_L(Ttest): print( "ZLL From file: ", ZLL_5d[pscale/2][pscale/2], "real: ", pdf5d.eval_L(Ttest) ) raise IOError except IOError: print("Re-making ZLL") for itheta, theta in enumerate(ZThetaArray): if itheta%(pscale/10) == 0 : print("%i out of 10" %(itheta/(pscale/10))) for iphi, phi in enumerate(ZPhiArray): test_dir = Vec().set_angles(theta, phi) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) this_LL = 0 this_LL = pdf5d.eval_L(Ttest) ZLL_5d[itheta, iphi] = this_LL dump_file = open( options.o + basefilename + "%iZLL_TESTF" %(evt.id), 'w') pickle.dump(ZLL_5d, dump_file) dump_file.close() minLL_indices = np.unravel_index(ZLL_5d.argmin(), ZLL_5d.shape) minLL_D.set_angles(ZThetaArray[minLL_indices[0]], ZPhiArray[minLL_indices[1]]) #extra zoomed version ZZLL_5d = np.zeros((pscale, pscale)) ZZThetaArray = np.linspace(reco5d_D.theta() - torad(zzscale), reco5d_D.theta() + torad(zzscale), pscale) ZZPhiArray = np.linspace(reco5d_D.phi() - torad(zzscale), reco5d_D.phi() + torad(zzscale), pscale) try: pickled_file = open( options.o + basefilename + "%iZZLL_TESTF" %(evt.id), 'r') ZZLL_5d = pickle.load(pickled_file) test_dir = Vec().set_angles(ZZThetaArray[pscale/2], ZZPhiArray[pscale/2]) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) if abs(ZZLL_5d[pscale/2][pscale/2] - pdf5d.eval_L(Ttest)) >= 1: #approximately equal print( "ZZL From file: ", ZZLL_5d[pscale/2][pscale/2], "real: ", pdf5d.eval_L(Ttest) ) raise IOError except IOError: print("Re-making reco ZZLL") for itheta, theta in enumerate(ZZThetaArray): if itheta%(pscale/10) == 0 : print("%i out of 10" %(itheta/(pscale/10))) for iphi, phi in enumerate(ZZPhiArray): test_dir = Vec().set_angles(theta, phi) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) this_LL = 0 this_LL = pdf5d.eval_L(Ttest) ZZLL_5d[itheta, iphi] = this_LL dump_file = open( options.o + basefilename + "%iZZLL_TESTF" %(evt.id), 'w') pickle.dump(ZZLL_5d, dump_file) dump_file.close() #extra super zoomed version ZZZLL_5d = np.zeros((pscale, pscale)) ZZZThetaArray = np.linspace(reco5d_D.theta() - torad(zzzscale), reco5d_D.theta() + torad(zzzscale), pscale) ZZZPhiArray = np.linspace(reco5d_D.phi() - torad(zzzscale), reco5d_D.phi() + torad(zzzscale), pscale) try: pickled_file = open( options.o + basefilename + "%iZZZLL_TESTF" %(evt.id), 'r') ZZZLL_5d = pickle.load(pickled_file) test_dir = Vec().set_angles(ZZZThetaArray[pscale/2], ZZZPhiArray[pscale/2]) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) if abs(ZZZLL_5d[pscale/2][pscale/2] - pdf5d.eval_L(Ttest)) >= 1: #approximately equal print( "ZZZL From file: ", ZZZLL_5d[pscale/2][pscale/2], "real: ", pdf5d.eval_L(Ttest) ) raise IOError except IOError: print("Re-making reco ZZZLL") for itheta, theta in enumerate(ZZZThetaArray): if itheta%(pscale/10) == 0 : print("%i out of 10" %(itheta/(pscale/10))) for iphi, phi in enumerate(ZZZPhiArray): test_dir = Vec().set_angles(theta, phi) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) this_LL = 0 this_LL = pdf5d.eval_L(Ttest) ZZZLL_5d[itheta, iphi] = this_LL dump_file = open( options.o + basefilename + "%iZZZLL_TESTF" %(evt.id), 'w') pickle.dump(ZZZLL_5d, dump_file) dump_file.close() #1d version one_LL_5d = np.zeros( pscale * pscale ) one_Theta = pi/2. one_PhiArray = np.linspace(-pi, pi, pscale*pscale) try: pickled_file = open( options.o + basefilename + "%ione_LL_TESTF" %(evt.id), 'r') one_LL_5d = pickle.load(pickled_file) test_dir = Vec().set_angles(one_Theta, one_PhiArray[pscale*pscale/2]) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) if one_LL_5d[pscale*pscale/2] != pdf5d.eval_L(Ttest): print( "LL From file: ", one_LL_5d[pscale*pscale/2], "real: ", pdf5d.eval_L(Ttest) ) raise IOError except IOError: print("Re-making one_LL") for iphi, phi in enumerate(one_PhiArray): test_dir = Vec().set_angles(one_Theta, phi) Ttest = Trk().set( id = 1, pos = T1.pos, t = T1.t, E = neut_E, dir = test_dir) this_LL = 0 this_LL = pdf5d.eval_L(Ttest) one_LL_5d[iphi] = this_LL dump_file = open( options.o + basefilename + "%ione_LL_TESTF" %(evt.id), 'w') pickle.dump(one_LL_5d, dump_file) dump_file.close() print("MC Nu 5d LL: %f, Zoomed Min 5d LL: %f, Reco 5d LL: %f" %(true_LL, ZLL_5d.min(), reco5d.lik)) if(true_LL < reco5d.lik or ZLL_5d.min() < reco5d.lik): print("WARNING: Reco 5d LL not minimal, reconstruction for this event CAN BE IMPROVED") res3dAng = (180/pi)*acos(neut_D.dot( reco3d_D ) ) res5dAng = (180/pi)*acos(neut_D.dot( reco5d_D ) ) res3dE = (reco3d_E - neut_E) / neut_E res5dE = (reco5d_E - neut_E ) / neut_E print("3dAng reco diff: %f" %res3dAng ) print("5dAng reco diff: %f" %res5dAng ) print("3dE reco diff: %f" %res3dE ) print("5dE reco diff: %f" %res5dE ) logfile.write( "Event n. %i\n" %(evt.id)) logfile.write( "MC Nu 5d LL: %f, Min 5d LL: %f, Reco 5d LL: %f\n" %(true_LL, ZLL_5d.min(), reco5d.lik)) logfile.write( "MC Nu dir theta: %f, phi: %f" %(neut_D.theta(), neut_D.phi())) logfile.write( "MC Lep dir theta: %f, phi: %f" %(lep_D.theta(), lep_D.phi())) if(len(evals5d)): logfile.write( "Start 5d reco dir theta: %f, phi: %f" %(evals5d[0].dir.theta(), evals5d[0].dir.phi())) logfile.write( " 5d reco dir theta: %f, phi: %f" %(reco5d_D.theta(), reco5d_D.phi())) logfile.write( " 3d reco dir theta: %f, phi: %f" %(reco3d_D.theta(), reco3d_D.phi())) logfile.write( " 5d minLL dir theta: %f, phi: %f" %(minLL_D.theta(), minLL_D.phi())) logfile.write( "Neut E: %f" %(neut_E)) logfile.write( "3d reco E: %f" %(reco3d_E)) logfile.write( "5d reco E: %f" %(reco5d_E)) logfile.write( "Dir 1, Dir 2: diff between [deg]\n") logfile.write( "Fit 5d start, Fit 5d end: %f \n" %( (180/pi) * ROOT.angle_between(T1.dir, reco5d_D) ) ) logfile.write( "Min 5d LL, Fit 5d: %f\n" %((180/pi)*ROOT.angle_between(minLL_D, reco5d_D)) ) logfile.write( "Min 5d LL, MC Nu: %f\n" %((180/pi)*ROOT.angle_between(minLL_D, neut_D)) ) logfile.write( "MC Nu, Fit 3d: %f deg\n" %(res3dAng)) logfile.write( "MC Nu, Fit 5d: %f deg\n" %(res5dAng)) logfile.write( "MC Nu E : %f\n" %(neut_E) ) logfile.write( "3dE reco : %f\n" %(res3dE)) logfile.write( "5dE reco : %f\n" %(res5dE)) logfile.write( " --- \n" ) logfile.close() print "bye"