#!/usr/bin/env python3 """ usage : timeshowerfit.py [options] Timeshowerfit is a reconstruction program for cascade events. Inputs a detector file and event file, and outputs an eventfile with reconstructed tracks. 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 -al -algorithm : 'MIN', 'MINE', 'GRAD', or 'GRADE'. MIN: uses LL evaluation only GRAD: uses LL gradient evaluation GRADE: uses energy integrated PDF object. default=min -pdf : supply pdf object for E minimization default=/sps/km3net/users/jseneca/jpp/data/PDE_bl1.6_12Ebins_40cdbins_N50.dat -da : do det.apply default=False -cb : correct bias default=False -l -logs: path to log file default=./timeshowerfit_logs.txt -n : process this many events default=100000000 -s : start with this event default=0 -np: process every np events default=1 -nf: Number of files to read default=0 -g : generation mode default=False -id : select event by id default=0 -e : only process events w/ MC-Enu greater than this default=0 -me: only process events w/ MC-Enu smaller than this default=10000000 -cc: containment cut default=False -q : use numpy minimizer default=False -fv: fitted variables ['ExyzTPt']. default=ExyzTPt -mc: Use mc hits from file default=False -pr: Prefit mode. "mc" or "aa" or "jpp" or "fullx" where x is the number of prefit iterations default=mc -pro: Prefit offset "[transdev,longdev,angdev,tdev,Efrac]" default=[] -E : elongation steps default=0 -bg: take into account background default=False -bl: take blurring into account default=0.0 -fa: JSirene enhancement factor default=1.0 -nw: deactivate weight function default=False -t : report timing informtion every N events default=100 -v : verbose (0,1 or 2) default=0 -i : interactive mode default=False -h : display this message default=False """ print("timeshowerfit starts...") import ROOT, os, sys, collections, array from math import * from copy import copy, deepcopy import logging, ast import time try : import aa except ImportError: print("Failed to import aa; Make sure AADIR is in your PYTHONPATH!") sys.exit() import rec from rec import rec2 import resource import scipy import numpy #print( scipy.__version__ ) #print( scipy.__path__ ) #print( numpy.__version__ ) #print( numpy.__path__ ) def limit_memory(maxsize): soft, hard = resource.getrlimit(resource.RLIMIT_AS) resource.setrlimit(resource.RLIMIT_AS, (maxsize, maxsize)) log.debug( "Limited memory to {} MB".format( maxsize ) ) #limit_memory( 8 * 1024 * 1024 * 1024 ) n_threads = '1' os.environ['OPENBLAS_NUM_THREADS'] = n_threads os.environ['OMP_NUM_THREADS'] = n_threads os.environ['NUMEXPR_NUM_THREADS'] = n_threads os.environ['MKL_NUM_THREADS'] = n_threads os.environ['VECLIB_MAXIMUM_THREADS'] = n_threads print( "loading options" ) options = aa.Options( __doc__, sys.argv[1:] ) print( "loaded options" ) logging.basicConfig(level=logging.DEBUG, filename=options.l, filemode="a+", format="%(asctime)-15s %(levelname)-8s %(message)s") logger = logging.getLogger() logger.setLevel( { 0: 35, 1: 15, 2: 0 }[options.v] ) logger.addHandler( logging.StreamHandler( sys.stdout ) ) print( options ) logger.debug(options) from ROOT import vector, Hit, Det, Evt, Trk, Vec, EventFile, Timer, Mat, angle_between, v_light, c_light, getShowerMaximum, getShowerLength from numpy import random, log NFILES=options.nf #number of files to read SEED = 43 logger.debug( "seed: {}".format( SEED ) ) RNG = random.default_rng( seed = SEED ) Trk.set = aa.__set jshowerfile = [ os.environ["JPP_DIR"] + "/data/J14p.dat",\ os.environ["JPP_DIR"] + "/data/J13p.dat" ] estart = 5e5 background_rate = 7000 min_residual = -100 max_residual = 900 ln10 = log(10) correction = {} corr_file = open('/sps/km3net/users/jseneca/data/correction.txt', 'r') for line in corr_file: tmp = line.split(",") correction[ int(tmp[0]) ] = float(tmp[1][:-1]) corr_file.close() def correct_bias( hit_t, hit_tot ): return hit_t + correction[hit_tot] def correct_t( t, evt ): return t - ( evt.mc_t - max(0, (evt.frame_index-1) * 1e8) ) def check_hit_residuals( hits, trk, limit = 1000000): for h in hits: dt = h.t - trk.t - (trk.pos - h.pos).len() / v_light if abs(dt) < limit: print( dt ) 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; def deviate_trans_pos( mc_trk, trk_dev, dev ): """dev [m]""" if mc_trk.dir == trk_dev.dir: trk_dev.dir += Vec( 1e-6, 0, 0 ) #Otherwise we're going to run into issues trk_dev_pos_before = deepcopy( trk_dev.pos ) trk_dev.pos = trk_dev.pos + ( trk_dev.dir - mc_trk.dir ).normalize() * dev assert( abs( ( trk_dev_pos_before - trk_dev.pos ).len() - dev ) < 0.0001 ) return trk_dev def deviate_long_pos( mc_trk, trk_dev, dev ): """dev [m]""" trk_dev_pos_before = deepcopy( trk_dev.pos ) trk_dev.pos = trk_dev.pos + mc_trk.dir * dev assert( abs( ( trk_dev_pos_before - trk_dev.pos ).len() - dev ) < 0.0001 ) return trk_dev def deviate_dir( mc_trk, trk_dev, dev ): """dev [deg]""" if mc_trk.dir == trk_dev.dir: trk_dev.dir += Vec( 1e-6, 0, 0 ) #Otherwise we're going to run into issues rot = Mat() cross_vec = mc_trk.dir.cross( trk_dev.dir ) #Deviate in direction of trk_dev rot.set_rotation( cross_vec, cos( dev * pi / 180 ) ) trk_dev.dir = rot*mc_trk.dir smax = getShowerMaximum( trk_dev.E ) trk_dev.pos = trk_dev.pos + smax * mc_trk.dir + smax * ( -trk_dev.dir ) assert( angle_between( trk_dev.dir, mc_trk.dir ) - dev * pi / 180 < 0.0001 ) return trk_dev #----------------------------------------- # input & output file #----------------------------------------- if options.o.startswith('auto') : # base output file on input file outfile = options.f+".root" if options.o == "auto_pwd" : outfile = outfile.split("/")[-1] else : outfile = options.o logger.debug (" input file = {}".format( options.f ) ) logger.debug (" output file = {}".format( outfile ) ) if options.g: f = ROOT.EventFile() #if events generated else : from glob import glob files = glob(options.f) if not NFILES else glob(options.f)[:NFILES] f = ROOT.EventFile( files ) f.set_output( outfile ) mem1 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss logger.info( "Memory usage after loading files: {} MB".format( mem1 * 1 / 1024 ) ) #----------------------------------------------- # detector file. todo: get it from the evt file #----------------------------------------------- print( "chosen detector: ", options.d ) if options.d == "auto" : det = Det( f ), logger.info("Got detector from file") else : det = Det( options.d ) det.prnt() # For mc-analysis, we need the coord-origin. If it is missing, 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["timeshowerfit_detector"] = options.d f.output_file.WriteObject( f.header ,"Head" ) f.output_file.ls() #=================================================================================== # end of setup #=================================================================================== from ROOT import JppShower, JppGrad, JppGradE if "E" in options.al: pdf = JppShower[JppGradE]( [ options.pdf ], options.bl, background_rate, #buggy behavior when no background min_residual, max_residual, options.v, 1 ) else: pdf = JppShower[JppGrad]( [ "/pbs/throng/km3net/src/Jpp/master/data/J14p.dat", "/pbs/throng/km3net/src/Jpp/master/data/J13p.dat" ], options.bl, background_rate, #buggy behavior when no background min_residual, max_residual, options.v, 1, options.E ) #0 : do nothing, 1 = move to shower max, > 1 : sample #else: # from ROOT import JppShower # pdf = JppShower( jshowerfile[0], # jshowerfile[1], # options.bl, # min_residual, # max_residual, # options.bg, # options.E ) # if options.nw: # pdf.remove_weights() REC_ID = [ ( 105, 1005 ), ( 205, 2005 ) ][ "GRAD" in options.al ] logger.debug( "Made JppShower_simple" ) if options.q and "GRAD" in options.al: from scipy.optimize import minimize import numpy as np from math import pi import ctypes def fit5d( start_trk, hits, det, rec_id = 205, fitted_vars = "ExyzTPt", use_unhits = True, E_elongation = False, init_event = True, verbose = 0): #LL0 = gradpdf.eval_lik( start_trk, True ) t0 = time.time() print( " before init_event: {} MB".format( resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 ) ) if init_event: pdf.init_event( det, hits, start_trk.pos, max_dist = max_dist_func( det, hits, start_trk.pos ), incl_unhits = use_unhits) print( " after init_event: {} MB".format( resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 ) ) def mf( x ) : 'a function to call from the scipy minimizer' global sum_t global n_it with Timer("mx") as tt: t0 = tt.time_since_last_start() r, df = pdf.eval_lik_grad( 10**x[0],x[1],x[2],x[3],x[4],x[5],x[6] ) sum_t += tt.time_since_last_start() - t0 n_it += 1 #input( "This iteration: {} s".format( tt.time_since_last_start() - t0 ) ) df[0] = df[0] * (10**x[0]) * ln10 #ln(10), dE/dlogE logger.info( "LL-LL0: {}, r: {}, df: {}".format( r - LL0, r, df ) ) #Slow #logger.debug( "x: ", x ) #input() return r - LL0, df; def mf_lik( x ) : 'a function to call from the scipy minimizer' global sum_t_lik global n_it_lik with Timer("mx") as tt: t0 = tt.time_since_last_start() r = pdf.eval_lik( 10**x[0],x[1],x[2],x[3],x[4],x[5],x[6] ) sum_t_lik += tt.time_since_last_start() - t0 n_it_lik += 1 #input( "This iteration: {} s".format( tt.time_since_last_start() - t0 ) ) logger.info( "mf_lik LL-LL0: {}, r: {}".format( r - LL0, r ) ) #Slow #logger.debug( "x: ", x ) #input() return r - LL0; def mf_grad( x ) : 'a function to call from the scipy minimizer' global sum_t_grad global n_it_grad with Timer("mx") as tt: t0 = tt.time_since_last_start() r, df = pdf.eval_lik_grad( 10**x[0],x[1],x[2],x[3],x[4],x[5],x[6] ) sum_t_grad += tt.time_since_last_start() - t0 n_it_grad += 1 #input( "This iteration: {} s".format( tt.time_since_last_start() - t0 ) ) df[0] = df[0] * (10**x[0]) * ln10 #ln(10), dE/dlogE logger.debug( "mf_grad LL-LL0: {}, r: {}, df: {}".format( r - LL0, r, df ) ) #Slow #logger.debug( "x: ", x ) #input() return df; x = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) #fits_, best = [], None #start_dirs = ROOT.icosahedron(); #for ii, T1.dir in enumerate( start_dirs ): x[0],x[1],x[2],x[3],x[4],x[5],x[6] = log10(start_trk.E), start_trk.pos.x, start_trk.pos.y, start_trk.pos.z, start_trk.dir.theta(), start_trk.dir.phi(), start_trk.t LL0 = pdf.eval_lik ( 10**x[0],x[1],x[2],x[3],x[4],x[5],x[6] ) LL_mc = pdf.eval_lik ( evt.mc_trks[0] ) while x[5] < 0 : x[5] += 2*pi op={'disp': True, #bool(options.v>1), #'maxcor': 10, 'gtol': 0.0001, #'eps': 1.4901161193847656e-08, 'maxiter': 10000 } global sum_t global n_it sum_t = 0.0 n_it = 0.001 global sum_t_lik global n_it_lik sum_t_lik = 0.0 n_it_lik = 0 global sum_t_grad global n_it_grad sum_t_grad = 0.0 n_it_grad = 0 print( "before minimize: {} MB".format( resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 ) ) with Timer("scipy.minimize") : logger.debug( "Starting fit..." ) res = minimize(mf, x, method= "BFGS", jac = True, options = op) # bounds = [ ( 0, 14 ), # ( -1000, 1000 ), # ( -1000, 1000 ), # ( -1000, 1000 ), # ( 0.0, pi ), # ( 0.0, 2*pi ), # ( -1000, 1000 ) ] ) print( "after minimize: {} MB".format( resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 ) ) print( "n iterations: ", n_it ) print( "average time per iteration: {}".format( sum_t / float(n_it) ) ) #print( "average time per iteration lik: {}".format( sum_t_lik / float(n_it_lik) ) ) #print( "average time per iteration grad: {}".format( sum_t_grad / float(n_it_grad) ) ) fit_python_trk = Trk() fit_python_trk.set( lik = res.fun, pos = Vec( res.x[1], res.x[2], res.x[3] ), dir = Vec().set_angles( res.x[4],res.x[5] ), t = res.x[6], E = 10**res.x[0], comment = "grad_time_fit_5d{}".format( ["","_elongation"][E_elongation] ), rec_type = rec_id ) fit_python_trk.comment += ", gradfit_cputime:{}".format(time.time() - t0 ) LL_fit = pdf.eval_lik ( fit_python_trk ) logger.debug( "LL_mc - LL0: {}".format( LL_mc - LL0 ) ) #logger.debug( "LL_fit - LL0: {}".format( res.fun - LL0 ) ) logger.debug( "MC trk : {}".format( evt.mc_trks[0] ) ) logger.debug( "Fit python pdf: {}".format( fit_python_trk ) ) print( "end gradfit5d: {} MB".format( resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 ) ) return fit_python_trk """ def time_fit5d( start_trk, hits, det, E_elongation = False): rec_id = 2005 logger.debug( "Grad Time Prefit" ) LL0 = pdf.eval_lik ( start_trk ) pdf.init_event( det, hits, start_trk.pos, max_dist = 800, incl_unhits = False) def mf( x ) : 'a function to call from the scipy minimizer' with Timer("mx"): r, df = pdf.eval_lik_grad( start_trk.E, x[0],x[1],x[2],start_trk.dir.theta(), start_trk.dir.phi(), x[3] ) logger.debug( "LL-LL0: {}, r: {}, df: {}".format( r - LL0, r, df ) ) return r - LL0, np.array([ df[0], df[1], df[2], df[5] ]) ; x = np.array([1.0, 1.0, 1.0, 1.0]) #x, y, z, t x[0],x[1],x[2],x[3] = start_trk.pos.x, start_trk.pos.y, start_trk.pos.z, start_trk.t LL0 = pdf.eval_lik ( start_trk ) op={'disp': True, #'maxcor': 10, 'gtol': 0.0001, #'eps': 1.4901161193847656e-08, 'maxiter': 10000 } with Timer("scipy.minimize") : res = minimize(mf, x, method= "BFGS", jac = True, options = op, bounds = [ ( -1000, 1000 ), ( -1000, 1000 ), ( -1000, 1000 ), ( -1000, 1000 ) ] ) fit_trk = Trk() fit_trk.set( lik = res.fun, pos = Vec( res.x[0], res.x[1], res.x[2] ), dir = start_trk.dir, t = res.x[3], E = start_trk.E, comment = "grad_time_prefit_5d{}".format( ["","_elongation"][E_elongation] ), rec_type = rec_id ) return copy( fit_trk ) """ elif not options.q: from ROOT import JppShowerFit MIN_PACKAGE = [ "Minuit2", "Minuit2" ][ "GRAD" in options.al ] MIN_ALGO = [ "Simplex", "BFGS2" ][ "GRAD" in options.al ] fitter5d = JppShowerFit[ [JppGrad, JppGradE]["E" in options.al] ]( pdf, MIN_PACKAGE, MIN_ALGO, 2, #strategy options.v ) def fit5d( start_trk, hits, det, rec_id = 105, fitted_vars = "ExyzTPt", use_unhits = True, E_elongation = False, init_event = True, verbose = 0 ): logger.debug( "Fit 5d" ) t0 = time.time() fixed_vars = [] for i, c in enumerate("ExyzTPt"): if c in fitted_vars: continue else: fixed_vars.append( i ) assert( len(fixed_vars) == 7 - len( fitted_vars ) ) fitter5d.fix_vars( -1 ) #Set all variables free if fixed_vars: fitter5d.fix_vars( *fixed_vars ) logger.debug( "Fixed {}.".format( [ ["x","y","z","theta","phi","t","E"][p] for p in fixed_vars ] ) ) else: logger.debug( "All parameters left free" ) fitter5d.set_incl_unhits( use_unhits ) try: if( E_elongation ): fit_trk = fitter5d.fit( start_trk, det, hits, init_event, True ) else: fit_trk = fitter5d.fit( start_trk, det, hits, init_event ) fit_trk.comment = "fit_5d{}".format( ["","_elongation"][E_elongation] ) fit_trk.comment += ", minfit_cputime:{}".format(time.time() - t0 ) fit_trk.rec_type = rec_id return copy( fit_trk ) except Exception as err: #Bug in Minuit 2 causes assertion errors sometimes, if the ROOT build has been made with debug options: https://root-forum.cern.ch/t/change-in-minuit2-behaviour-from-6-20-00-onwards/40362 print( "Minimization failed! Error: {}. We continue anyway.".format( err ) ) fit_trk = copy( start_trk ) fit_trk.comment = "Failed minimization." fit_trk.rec_type = -9999 print( "failed fit, trk:", fit_trk ) return fit_trk def pos_t_prefit( start_trk, hits, det, evt, ite_count = 1, E_elongation = False, verbose = False): t0 = time.time() jumps = [ 1.0, 0.5, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1 ] #m ico_dirs = ROOT.icosahedron(); for d in ico_dirs: d.normalize() ico_dirs.push_back( Vec( 0, 0, 0 ) ) #Have one start at the previous fit fitter5d._init_pdf( start_trk, det, hits ) #We init the event here so all prefits use the same hits for ite in range(ite_count): prefit_trks = [] for ioffset, offset_dir in enumerate(ico_dirs): ite_trk = deepcopy( start_trk ) ite_trk.pos = start_trk.pos + offset_dir * jumps[ ite ] #Look m in every direction logger.info( "Prefit: offset {} of 13, trk: {}".format( ioffset, ite_trk ) ) prefit_trks.append( fit5d( ite_trk, hits, det, rec_id = REC_ID[1], fitted_vars = "xyzt", use_unhits = False, E_elongation = bool( E_elongation ), init_event = False ).set( comment = "time_prefit_5d, offset {}, iteration {}".format( ioffset, ite ) ) ) for trk in sorted( prefit_trks, key = lambda x: x.lik ): evt.trks.push_back( trk ) logger.debug( "{}: {} lik {}".format( trk.comment, trk, trk.lik) ) prefit_trk = sorted( prefit_trks, key = lambda x: x.lik )[0] prefit_trk.comment += ", pos_t_prefit_cputime:{}".format( time.time() - t0 ) return prefit_trk #-------------------------------------------------- # Event Loop #-------------------------------------------------- logger.info( "Event loop starting..." ) iproc = 0 if options.id: ievt = options.id + 1 elif f.size() == 1: ievt = 0 elif not f.size(): logger.error( "ERROR: no events in file" ) raise IndexError else: ievt = len(f) + 1 #start from the top while( ievt > len(f) - options.s ): ievt-=1 #Looping over events while iproc < options.n: if ievt < 0 or ievt > len(f) + 1: break if (options.s + ievt)%options.np != 0:#select every np events ievt -= 1 continue if not options.g: #Choose between generating events, or reading from file file_size = int(ievt/f.current_file_size()) if int(ievt/f.current_file_size()) < len(files) else len(files) - 1 assert( f.set_index( file_size, 2 ) ) #HAS to be above 1 (???) evt = f.get_event( ievt%f.current_file_size() ) #print("Opened file with global index {} index {} and event index {}".format( f.global_index(), f.file_index, f.index ) ) if len( evt.mc_trks ) == 0: #no event ievt -= 1 continue if options.id: if evt.id > int(options.id): ievt -= 1 continue elif evt.id < int(options.id): break if evt.mc_trks.size() > 0: #------ MC energy cut if ( options.e > 0 and evt.mc_trks[0].E < options.e ) or \ ( options.me > 0 and evt.mc_trks[0].E > options.me ) : logger.debug( "Energy {} out of selected range! Continuing...".format( evt.mc_trks[0].E ) ) ievt -= 1 continue if options.cc: #------- containment cut radius = 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 radius > (530 - 91): ievt -= 1 logger.debug( "Not contained! Continuing..." ) continue elif options.g: f.evt.clear() # important: the EventFile has only one Evt object : reset it f.evt = Evt() 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 ) f.evt.mc_trks.push_back( nu ) f.evt = generate_light( pdf, f.evt, det = det, first = False, min_e = 100000, background = options.bg, verbose = options.v ) evt = f.evt logger.debug("Found event: {}, energy {}\n{}".format(evt.id, evt.mc_trks[0].E, evt) ) with Timer("0) events") as t0 : t_start = t0.value() #No continues until write iproc+=1 # count actually processed (not skipped) events logger.info(" === event {} ===".format(evt.id) ) with Timer("1) calibration") : if options.mc or options.g: input_hits = vector() for hit in evt.mc_hits: if options.fa > 1 and random.random() > 1./factor: continue input_hits.push_back( hit ) else: input_hits = evt.hits for h in input_hits: h.t = correct_t( h.t, evt ) if options.da: det.apply( input_hits ) if options.cb: for h in input_hits: h.t = correct_bias( h.t, h.tot ) with Timer("4) prefit") : if "mc" in options.pr: T1 = Trk().set( dir = evt.mc_trks[0].dir, pos = evt.mc_trks[0].pos, t = evt.mc_trks[0].t, E = evt.mc_trks[0].E, comment = "use mc prefit" ) #T1.E = 10 #For testing elif "jpp" in options.pr: jpp_trk = Trk().set( lik = -1e12 ) for trk in evt.trks: #select jpp best fit if trk.rec_type == 4000 and ( 5 in trk.rec_stages ) and trk.lik > jpp_trk.lik: jpp_trk = deepcopy(trk) logger.debug( "raw jpp trk t: ", correct_t( jpp_trk.t, evt ) ) logger.debug( jpp_trk ) if jpp_trk.lik == -1e12: logger.error( "ERROR: no jpp trk found in this event. Continuing." ) continue #smax_d = getShowerLength( jpp_trk.E, 0.5 ) #jpp_trk.pos = jpp_trk.pos - jpp_trk.dir * smax_d #jpp_trk.t = correct_t( jpp_trk.t - smax_d / c_light, evt ) T1 = Trk().set( dir = jpp_trk.dir, pos = jpp_trk.pos, t = jpp_trk.t, E = jpp_trk.E, comment = "use jpp as prefit" ) elif "aa" in options.pr or "full" in options.pr: aashower_trk = Trk().set( lik = -1e12 ) for trk in evt.trks: #select aashower best fit if trk.rec_type == 101 and ( 304 in trk.rec_stages ) and trk.lik > aashower_trk.lik: aashower_trk = deepcopy(trk) logger.debug( "raw aashower trk t: {}".format( correct_t( aashower_trk.t, evt ) ) ) logger.debug( "{}".format( aashower_trk ) ) if aashower_trk.lik == -1e12: logger.error( "ERROR: no aashower trk found in this event. Continuing." ) continue smax_d = getShowerLength( aashower_trk.E, 0.5 ) aashower_trk.pos = aashower_trk.pos - aashower_trk.dir * smax_d aashower_trk.t = correct_t( aashower_trk.t - smax_d / c_light, evt ) T1 = Trk().set( dir = aashower_trk.dir, pos = aashower_trk.pos, t = aashower_trk.t, E = aashower_trk.E, comment = "use aashower as prefit" ) if "full" in options.pr: t0 = time.time() prefit_hits = ROOT.filter_hits( input_hits, aashower_trk, min_residual, max_residual, 300 ) il = [ s for s in options.pr if s in "23456789" ] ite_count = int(il[0]) if il else 1 logger.debug( "Will do {} prefit iterations".format( ite_count ) ) prefit_trk = pos_t_prefit( deepcopy(aashower_trk), prefit_hits, det, evt, ite_count, E_elongation = bool(options.E), verbose = True ) logger.debug( "prefit: {}".format( prefit_trk ) ) logger.debug( "aashower: {}".format( aashower_trk ) ) logger.debug( "mc_trk: {}".format( evt.mc_trks[0] ) ) logger.debug( "lik: {}, pos dif: {}, time dif: {}".format( prefit_trk.lik, ( prefit_trk.pos - evt.mc_trks[0].pos ).len(), abs( prefit_trk.t - evt.mc_trks[0].t ) ) ) T1 = Trk().set( dir = prefit_trk.dir, pos = prefit_trk.pos, t = prefit_trk.t, E = prefit_trk.E, rec_type = prefit_trk.rec_type, comment = "use aashower + time prefit" ) T1.comment += ", pos_t_prefit_cputime:{}".format(time.time() - t0 ) pr_list = ast.literal_eval( options.pro ) if pr_list: assert( len(pr_list) == 5 ) logger.debug( "Offsetting prefit: ", pr_list ) if pr_list[0]: deviate_trans_pos( evt.mc_trks[0], T1, pr_list[0] ) if pr_list[1]: deviate_long_pos ( evt.mc_trks[0], T1, pr_list[1] ) if pr_list[2]: deviate_dir ( evt.mc_trks[0], T1, pr_list[2] ) T1.t += pr_list[3] T1.E *= 1. + pr_list[4] T1.set( comment = ", prefit offset:" + options.pro ) """ else: for trk in evt.trks: if trk.rec_type == 101 and trk.rec_stages[-1] == 4: T1 = Trk().set( dir = trk.dir, pos = trk.pos, t = trk.t, E = trk.E / 2, comment = "use aashowerfit prefit" ) """ 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 ) if not options.q: fitter5d.set_max_dist( max_dist_func( det, digi_hits, T1.pos ) ) logger.info("max dist 5d: {}".format( fitter5d.get_max_dist() )) logger.debug( "Selected {} hits to be used in fit".format( len(digi_hits) ) ) """ for hit in digi_hits: res = hit.t - ( hit.pos - f.evt.mc_trks[0].pos ).len() / v_light print( "T res of hit {}: {} ns".format( hit.id, res ) ) if res < 0: input() """ mem0 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss logger.debug( " before fit Peak memory usage: {} MB".format( mem0 * 1 / 1024 ) ) with Timer("6) direction fit") : logger.debug( "Fitting 5d..." ) logger.debug( "Start track: {}".format( T1 ) ) logger.debug( "True track: {}".format( evt.mc_trks[0] ) ) fit_trk = Trk() fit_trk = fit5d( T1, digi_hits, det, rec_id = REC_ID[0], fitted_vars = options.fv, use_unhits = True, E_elongation = False, init_event = True, verbose = options.v ) print( "fit5d over here" ) mem1 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss logger.debug( " after fit Peak memory usage: {} MB".format( mem1 * 1 / 1024 ) ) #logger.info( "N ite: {}. mem before {} mem after {}, time {}".format( n_it, mem0 / 1024, mem1 / 1024, fit_trk.comment.split('cputime:')[-1] ) ) logger.info( "True track: {}".format( evt.mc_trks[0] ) ) logger.info( "Fit track: {}".format( fit_trk ) ) logger.debug( "deviation pos: {} m".format( ( fit_trk.pos - evt.mc_trks[0].pos ).len() ) ) logger.debug( "deviation dir: {} deg".format( ROOT.angle_between( fit_trk.dir, evt.mc_trks[0].dir ) * 180 / pi ) ) logger.debug( "deviation t: {} ns".format( fit_trk.t - evt.mc_trks[0].t ) ) logger.debug( "deviation E: {} GeV".format( fit_trk.E - evt.mc_trks[0].E ) ) logger.debug( "Reconstruction times, prefit: {} final fit: {}".format( T1.comment.split('cputime:')[-1], fit_trk.comment.split('cputime:')[-1] ) ) f.evt.trks.push_back( fit_trk ) print( "after push_back: {} MB".format( resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 ) ) with Timer("7 post processing - writing ") as t1: t_end = t1.value() logger.debug( "time: {} ".format( t_end - t_start ) ) f.write() if iproc%options.t == 0 : Timer.report() f.report() ievt -= 1 f.output_file.Write() # todo: is this really needed? mem = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss logger.debug( "Peak memory usage: {} MB".format( mem * 1 / 1024 ) ) del f logger.info("timeshowerfit done.")