#!/usr/bin/env python """ usage : newshowerfit.py -f input_file Newshowerfit 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). Newshowerfit 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_v4/km3_v4_nueCC_1.JTE.aareco.root -o -output_file : The output file, aanet format default=auto_pwd -d -detector_file : det or detx. Set this to "auto" to get it from the header of the input file. default=/pbs/throng/km3net/detectors/km3net_jul13_90m.det -l : directory where to find newshowfit.so default=. -n : process this many events default=100000000 -np: process every np events default=1 -e : only process events w/ MC-Enu greater than this default=0 -me: only process events w/ MC-Enu smaller than this default=10000000 -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 -r3: reconstruct 3d method default=False -r4: reconstruct 4d method default=False -r5: reconstruct 5d method default=False -E : activate energy elongation default=False -mc: Use mc hits from file default=False -smc: Use mc position for fit default=False -cc: containment cut default=False -g : generation mode default=False -bg: take into account background default=False -bl: take blurring into account default=0.0 -v : verbose default=0 """ print("newshowerfit starts...") import ROOT, os, sys, collections, array from math import * from copy import copy print("here1") try : import aa except ImportError: print("Failed to import aa; Make sure AADIR is in your PYTHONPATH!") sys.exit() print("here2") import rec print("here3") import resource print("here4") #from toy_MC import * 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 JppShower, JppShowerFit from numpy import random SEED = 43 print( "seed: ", SEED ) RNG = random.default_rng( seed = SEED ) #----------------------------------------------------- # Fancy-smancy setting of members #----------------------------------------------------- for T in Trk, JppShower, 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", "mc" ][ options.mc ] muscale = 1.0 pdffile = "/sps/km3net/users/jseneca/aanet/rec/pdf_data/genhen_JUNE_2___.hist.root" jshowerfile = ["/sps/km3net/users/jseneca/jpp/data/J14p.dat",\ "/sps/km3net/users/jseneca/jpp/data/J13p.dat"] estart = 5e5 tolerance = 1 rec_type = options.r min_residual = -100; max_residual = 900; def max_dist_func ( det, hits, pos ) : rmax, fmax = 800, 2.0 rmin, fmin = 300, 0.2 h_d_ratio = ROOT.hit_to_dom_ratio( det, hits, pos, 200 ) #energy estimate if h_d_ratio == 0 : # no doms in radius return rmax log_h_d_ratio = log10(h_d_ratio) if log_h_d_ratio < fmin : return rmin return rmin + (log_h_d_ratio-fmin) * (rmax-rmin)/fmax; #----------------------------------------- # input & output file #----------------------------------------- if options.o.startswith('auto') : # base output file on input file outfile = options.f+"."+suffix+".root" if options.o == "auto_pwd" : outfile = outfile.split("/")[-1] else : outfile = options.o print (" input file =", options.f ) print (" output file =", outfile ) if not options.g: #if events not generated f = ROOT.EventFile( options.f ) f.set_output( outfile ) elif options.g: #if events generated f = ROOT.EventFile() f.set_output( outfile ) #----------------------------------------------- # 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. try: coord_origin = f.header.coord_origin() except: coord_origin = det.get_cg() #----------------------------------------------- # Add our info to the header and write it. #----------------------------------------------- if not options.g: #if events are not being generated f.header["newshowerfit_detector"] = options.d f.output_file.WriteObject( f.header ,"Head" ) f.output_file.ls() #=================================================================================== # 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) ff = 0.53972 + x * 0.16419 + x*x * -0.0168034; if ff > 1e-3 : trk.E = trk.E / ff; #-------------------------------------------------- # 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 ) digpdf = DigitalShowerPdf().set( infile = pdffile, k40prob = 5000*1e-6, eval_mode = DigitalShowerPdf().quick ) digpdf.init() fitter3d = ShowerFit ("digi", digpdf ).set ( verb = options.v, minimizer_name ='Minuit2', minimizer_algo ='Migrad' ) fitter3d.fix_vars( 0, 1, 2, 5 ) pdf5d = JppShower( jshowerfile[0], jshowerfile[1], True, options.bl, min_residual, max_residual, options.bg ) #pdf4d = JppShowerNpe( pdf5d ) #integrates time pdf fitter5d = JppShowerFit( "5d_pdf", pdf5d, "Minuit2", "MIGRAD", options.v ) fitter5d.fix_vars( 0, 1, 2, 5 ) """ fitter4d = JppShowerFitNpe( "4d_pdf", pdf4d, "Minuit2", "MIGRAD", options.v ) fitter4d.fix_vars( 0, 1, 2, 5 ) """ def fit3d( start_trk, trks, fitter, hits, det, rec_id = 103): fits_3d = [] for ii, start_trk.dir in enumerate( ROOT.icosahedron() ): do_init_pdf = ii==0 fit_3d = fitter.fit( start_trk, det, hits, do_init_pdf ) fit_3d.comment = "fit_3d"+str(ii) fits_3d.append( fit_3d ) for i, t in enumerate( sorted( fits_3d, key = lambda t : -t.lik ) ): t.set ( id = i + 2, rec_type = 1003, rec_stages = aa.make_vector([1,2,3,4]) if i+1==len(fits_3d) else aa.make_vector([1,2,3]), fitinf = aa.make_vector([1.0 * t.E, hits.size()] ) ) correct_energy(t) # uncorrected energy saved in fitinf[0] trks.push_back( t ) trks[-1].set( rec_type = rec_id ) return copy( trks[-1] ) def fit4d( start_trk, trks, fitter, hits, det, rec_id = 104, fast = False, E_elongation = False): fits_4d = [] if E_elongation: fit_4d = fitter4d.fit_elong( start_trk, det , hits, True ) fit_4d.comment = "fit_4d_elong_prefit" fit_4d.rec_type = 134 elif fast: fit_4d = fitter4d.fit_Efit( start_trk, det , hits, True ) fit_4d.comment = "fit_4d_fast_prefit" fit_4d.rec_type = 124 else: fit_4d = fitter4d.fit( start_trk, det , hits, True ) fit_4d.comment = "fit_4d_prefit" fit_4d.rec_type = 114 fits_4d.append( fit_4d ) for i, t in enumerate( sorted( fits_4d, key = lambda t : -t.lik ) ): t.set ( id = i + 2, rec_type = 1004, rec_stages = aa.make_vector([1,2,3,4]) if i+1==len(fits_4d) else aa.make_vector([1,2,3]) ) """ if i == len(fits4d_) - 1: #Could get different option to do this for ieval_t, eval_t in enumerate(pdf4d.get_eval_trks()): eval_t.rec_type = 204; trks.push_back(eval_t) """ trks.push_back( t ) trks[ trks.size( ) - 1 ].set( rec_type = rec_id ) #comment = "time:%.1f" %( t4d ) ) return copy( trks[-1] ) def fit5d( start_trk, trks, fitter, hits, det, rec_id = 105, E_elongation = False): fits_5d = [] if E_elongation: fit_5d = fitter5d.fit_elong( start_trk, det , hits, True ) fit_5d.comment = "fit_5d_elong_prefit" fit_5d.rec_type = 135 else: fit_5d = fitter5d.fit( start_trk, det , hits, True ) fit_5d.comment = "fit_5d_prefit" fit_5d.rec_type = 115 fits_5d.append( fit_5d ) for i, t in enumerate( sorted( fits_5d, key = lambda t : -t.lik ) ): t.set ( id = i + 2, rec_type = 1005, rec_stages = aa.make_vector([1,2,3,4]) if i+1==len(fits_5d) else aa.make_vector([1,2,3]) ) """ if i == len(fits5d_) - 1: #Could get different option to do this for ieval_t, eval_t in enumerate(pdf5d.get_eval_trks()): eval_t.rec_type = 204; trks.push_back(eval_t) """ trks.push_back( t ) trks[ trks.size( ) - 1 ].set( rec_type = rec_id ) #comment = "time:%.1f" %( t5d ) ) return copy( trks[-1] ) def fit5d_grad( start_trk, trks, fitter, hits, det, rec_id = 105, E_elongation = False): if options.q : # use scipy miminizer from scipy.optimize import minimize import numpy as np pdf.init_event( det, digi_hits, T1.pos, digi_fit.max_dist ) def mf( x ) : 'a function to call from the scipy minimizer' with Timer("mx"): df = np.array([ 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0 ]) r = ctypes.c_double() pdf.verb = 0 pdf.eval_lik_gradient( x[0],x[1],x[2],x[3],x[4],x[5],x[6], r, df ) return r.value, df ; x = np.array([ 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0 ]) fits_, best = [], None D = T1.dir.normalize() x[0],x[1],x[2] = acos(D.z), atan2( D.y, D.x ), log10(T1.E) while x[1] < 0 : x[1] += 2*pi op={'gtol': 0.0001, 'eps': 1.4901161193847656e-08, 'maxiter': None, 'disp': False, 'return_all': False} with Timer("scipy.minimize") : res = minimize( mf, x, method= "bfgs", jac = True, options = op ) if res.x[2] > 12 or res.x[2] != res.x[2] : res.x[2] = 12 fit = copy( T1 ).set ( lik = res.fun, dir = Vec().set_angles( res.x[0],res.x[1] ), E = 10**res.x[2], comment = "fit"+str(ii) ) fits_.append( fit ) if not best or fit.lik < best.lik : # inversion of lik will happen later; here lower is better best = fit best.res = res fits_5d = [] if E_elongation: fit_5d = fitter5d.fit_elong( start_trk, det , hits, True ) fit_5d.comment = "fit_5d_elong_prefit" fit_5d.rec_type = 135 else: fit_5d = fitter5d.fit( start_trk, det , hits, True ) fit_5d.comment = "fit_5d_prefit" fit_5d.rec_type = 115 fits_5d.append( fit_5d ) for i, t in enumerate( sorted( fits_5d, key = lambda t : -t.lik ) ): t.set ( id = i + 2, rec_type = 1005, rec_stages = aa.make_vector([1,2,3,4]) if i+1==len(fits_5d) else aa.make_vector([1,2,3]) ) """ if i == len(fits5d_) - 1: #Could get different option to do this for ieval_t, eval_t in enumerate(pdf5d.get_eval_trks()): eval_t.rec_type = 204; trks.push_back(eval_t) """ trks.push_back( t ) trks[ trks.size( ) - 1 ].set( rec_type = rec_id ) #comment = "time:%.1f" %( t5d ) ) return copy( trks[-1] ) #-------------------------------------------------- # Event Loop #-------------------------------------------------- iproc = 0 ievt = 2800 #start from the top while iproc < options.n: if ievt < 0 or ievt > 3000: break if ievt%options.np != 0:#select every np events ievt -= 1 continue #Choose between generating events, or reading from file if not options.g: evt = f.get_event( ievt ) if len( evt.mc_trks ) == 0: #no event ievt -= 1 continue elif options.g: f.evt.clear() nu = Trk() pos_z = RNG.uniform( 91, 677 - 91 ) r = 1e99 while r > 530 - 91: pos_x = RNG.uniform( 0, 530 - 91 ) pos_y = RNG.uniform( 0, 530 - 91 ) r = sqrt( pos_x*pos_x + pos_y*pos_y ) nu.pos = Vec( pos_x, pos_y, pos_z ) nu.dir = Vec( RNG.uniform(), RNG.uniform(), RNG.uniform() ).normalize() nu.E = RNG.uniform( options.e, options.me ) print( "nu E", nu.E ) f.evt = generate_evt( pdf5d, [nu], gen_mode = ["normal", "elongation"][options.E], evt_id = ievt, t0 = 0.0, background = options.bg, verbose = options.v ) evt = f.evt # minimum mc neutrino energy requirement. if evt.mc_trks.size() > 0: if ( options.e > 0 and evt.mc_trks[0].E < options.e ) or \ ( options.me > 0 and evt.mc_trks[0].E > options.me ) : ievt -= 1 continue print( "evt:", evt) # containment cut if options.cc: r = sqrt(evt.mc_trks[0].pos.x**2 + evt.mc_trks[0].pos.y**2) if evt.mc_trks[0].pos.z > (677 - 91) or evt.mc_trks[0].pos.z < (0 + 91) or r > (530 - 91): ievt -= 1 continue print("Found event: %i, energy %i" %(evt.id, evt.mc_trks[0].E) ) print( evt ) with Timer("0) events") as t0 : # from now on, there may be no 'continue' statements util write iproc+=1 # count actually processed (not skipped) events if options.v > 0 : print(" === event ", evt.id , " ===") with Timer("1) calibration") : if hit_type == 'raw' and not options.g: # 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( input_hits ) if hit_type == 'mc' or options.g: 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 coincidence 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") : if options.smc: T1 = Trk().set( id = 2, pos = evt.mc_trks[0].pos, t = evt.mc_trks[0].t, E = estart, rec_type = rec_type, rec_stages = aa.make_vector([1,2]), comment = "use mc vertex" ) else: 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, min_residual, max_residual, 800 ) # the maximum distance to vertex for selecting doms/hits for final fit. fitter3d.max_dist = max_dist_func( det, digi_hits, T1.pos ) digpdf.force_precompute = True fitter4d.max_dist = max_dist_func( det, digi_hits, T1.pos ) print("max dist: %f" %fitter4d.max_dist) fitter5d.max_dist = max_dist_func( det, digi_hits, T1.pos ) print("max dist 5d: %f" %fitter5d.max_dist) with Timer("6) direction fit") : besttrk3d = ROOT.Trk() if options.r3: with Timer( "6.1) direction fit 3d" ): print("Fitting 3d...") besttrk3d = fit3d( T1, f.evt.trks, fitter3d, digi_hits, det ) print( " mc E: {}, reco3d E: {}, reco3d angres: {}".format( evt.mc_trks[0].E, besttrk3d.E, ROOT.angle_between( besttrk3d.dir, evt.mc_trks[0].dir ) ) ) if options.r4: with Timer( "6.2) direction fit 4d" ): print("Fitting 4d...") besttrk4d = fit4d( besttrk3d, f.evt.trks, fitter4d, digi_hits, det, fast = False, E_elongation = options.E ) print( " mc E: {}, reco4d E: {}, reco4d angres: {}".format( evt.mc_trks[0].E, besttrk4d.E, ROOT.angle_between( besttrk4d.dir, evt.mc_trks[0].dir ) ) ) if options.r5: with Timer( "6.3) direction fit 5d" ): print("Fitting 5d...") besttrk5d = fit5d( besttrk3d, f.evt.trks, fitter5d, digi_hits, det, E_elongation = options.E) print( " mc E: {}, reco5d E: {}, reco5d angres: {}".format( evt.mc_trks[0].E, besttrk5d.E, ROOT.angle_between( besttrk5d.dir, evt.mc_trks[0].dir ) ) ) with Timer("7 post processing") : # for the best track , fill the hit_ids besttrk = f.evt.trks[ f.evt.trks.size()-1] besttrk.hit_ids.clear() if digi_hits.size(): for h in digi_hits : besttrk.hit_ids.push_back( h.id ) #------------------------------------------------------------- # optional, on-the-fly analysis of reconstruction performance #------------------------------------------------------------- #if evt.mc_trks.size() >0 and f.evt.mc_trks[0].E > 0 and iproc < options.a : #with Timer("analysis"): #from analysis import analyze #analyze( f.evt , 1 , pdf4d , coord_origin , rec_type = rec_type ) #analyze is broken... with Timer("7 post processing2 - undo time shift ") : # Undo the time shift if hit_type == 'raw' and not options.g: for h in evt.hits : h.t += hit_t0 for t in evt.trks : t.t += hit_t0 # tricky: as aanet does not clear the trks for some input # formats, we must make sure we delete our trks - otherwise # they just pile up. # (maybe aanet should clear the trks, always). with Timer("7 post processing - writing ") as t1: print() #evt.setusr("newshowerfit_time", t1.time_since_last_start() - t0.time_since_last_start() ) #print( "t0, t1, time: ", t0.time_since_last_start(), t1.time_since_last_start(), t1.time_since_last_start() - t0.time_since_last_start() ) print( "evt trks: ", len(evt.trks) ) print( "f evt trks: ", len(f.evt.trks) ) f.write() evt.trks.clear() if iproc%options.t == 0 : Timer.report() f.report() ievt -= 1 f.output_file.Write() # todo: is this really needed? print("Done with loop") mem = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss print( "Peak memory usage: ", mem * 1 / 1024 ) del f print("newshowerfit done.") """ prefitDirHit = Vec(0.0, 0.0, 0.0) prefitDirDom = Vec(0.0, 0.0, 0.0) n_forward = 0 n_backward = 0 if options.s: with Timer("5.5) principal direction pre-fit") : print( "number of pre_hits: %i " %( len( pre_hits ) ) ) for hit in pre_hits: prefitDirHit -= hit.dir prefitDirDom += hit.pos - T1.pos a = ROOT.angle_between( hit.pos - T1.pos, evt.mc_trks[0].dir ) print("Angle difference hit: %f" %( a ) ) if a < pi/2 : n_forward += 1 else : n_backward += 1 prefitDirHit.normalize() prefitDirDom.normalize() print( "prefit-hit dir theta %f phi %f" %( prefitDirHit.theta(), prefitDirHit.phi() ) ) print( "prefit-dom dir theta %f phi %f" %( prefitDirDom.theta(), prefitDirDom.phi() ) ) print( "true dir theta %f phi %f" %( evt.mc_trks[0].dir.theta(), evt.mc_trks[0].dir.phi() ) ) print("Angle difference prefit-hit dir: %f" %( ROOT.angle_between( prefitDirHit, evt.mc_trks[0].dir ) ) ) print("Angle difference prefit-dom dir: %f" %( ROOT.angle_between( prefitDirDom, evt.mc_trks[0].dir ) ) ) print( "n forward: %i, n backward: %i" %( n_forward, n_backward ) ) startDir = prefitDirHit """