#!/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 o 1 (M-est) pos+time fit o 12 likelihood fits (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 The algorithm fits, in total, 7 parameters : x,y,z,t, theta, phi, logE The elements of the 7x7 error-covariance matrix correspond to these variables. At the fit algorithm fits the vertex (position and time) seperately from the direction and energy, the error-covariance matrix is block-diagonal with a 4x4 and 3x3 block. The final 3x3 is only filled for the best track. options: -f -input_file : Input file in any format aanet can read. default=../data/mc5.1.numuCC.ps10.aa.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=../data/KM3NeT_-00000001_20171212.detx -n : process this many events default=10000000000 -E : only process events w/ MC-Enu greater than this default=0 -a : run mc analysis on the first N events default=10 -R : recallibrate (dress) the hits with the detector default=False -X : check consistency of pos/dirs with detx for first N evts default=5 -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=1000 -r : rec type identifyer for written tracks default=definitions -m : do hit merging default=True -M : write copy and add meta data default=True -S : copy summary-slice data to output file default=True -T : require at least this many triggered hits, otherswise skip default=0 -P : dynamic calib. positions file default=none -Q : dynamic calib. rotations (orientations) file default=none -h : display this message default=False -i : start python prompt when done default=False -q : use numpy minimizer default=True -l : do dummy fit for debugging default=False -pdf : file with pdf histogram default=default -hit_type : use raw or mc hits default=raw -estart : starting value for energy fit default=5e5 -bgrate_hz : assumed background rate default=5000 -mest_hit_mint : M-estimator hit selection, min residual (ns) default=-800 -mest_hit_maxt : M-estimator hit selection, max residual (ns) default=800 -digi_hit_mint : Final fit hit selection, min resdidual (ns) default=-100 -digi_hit_maxt : Final fit hit selection, max resdidual (ns) default=900 -digi_hit_maxr : Final fit hit selection, max distance (m) default=800 -minimizer_gtol : scipy minimizer tollerance parameter default=0.0001 """ import ROOT,os,sys,collections, array, gc import ctypes from math import * from copy import copy import numpy import scipy from util import pyroot_util from scipy.optimize import minimize import numpy as np try : import aa except ImportError: print ("Failed to import aa; Make sure AADIR is in your PYHTONPATH!") sys.exit() import rec sys.path.append(aa.aadir+"/externals/km3net-dataformat/definitions/") from reconstruction import data as reco_ids recstages = [ reco_ids[x] for x in "AASHOWERFITPREFIT AASHOWERFITPOSITIONFIT AASHOWERFITDIRECTIONENERGYFIT".split() ] options = aa.Options( __doc__, sys.argv[1:] ) print( options ) from ROOT import vector, Hit, Det, Evt, Trk, Vec, EventFile, Timer from ROOT import MestShowerPdf, DigitalShowerPdf, ShowerFit #----------------------------------------------------- # Fancy-smancy setting of members #----------------------------------------------------- for T in Trk, DigitalShowerPdf, MestShowerPdf, ShowerFit : T.set = aa.__set #----------------------------------------------------- # Error matrix filling #--------------------------------------------------- def set_error_matrix( track, np_hessian, offset ): 'Insert the inverse of hessian as a block-matrix, starting at offset,offset in track.error_matrix.' N = np_hessian.shape[0] if N!= np_hessian.shape[1] : print ("error: hessian should be square meatrix") # allways make a 7x7 matrix track.error_matrix.resize( 7**2 ); try : errmat = numpy.linalg.pinv( np_hessian ) #pinv deals better with bad cases for i in range(N) : for j in range(N) : track.error_matrix[ ( i+offset ) * 7 + j+offset ] = errmat[i,j] except numpy.linalg.LinAlgError as e : print ("failed to invert error matrix" , e) print (errmat) track.error_matrix.clear() pdffile = aa.aadir+"/rec/pdf_data/genhen_JUNE_2___.hist.root" if options.pdf == "default" else options.pdf rec_type = reco_ids['AANET_RECONSTRUCTION_TYPE'] if options.r == "definitions" else int(options.r) def max_dist_func ( det, hits, pos ) : """ Return the radius of a sphere that contains almost all hits. """ dist_rmin = 300 # these have been determined on mc dist_rmax = 800 dist_fmax = 2.0 dist_fmin = 0.2 f = ROOT.hit_to_dom_ratio( det, hits, pos ) # count at 200 mn if f == 0 : # no doms in radius return dist_rmax f = log10(f) if f < dist_fmin : return dist_rmin return dist_rmin + (f-dist_fmin) * (dist_rmax-dist_rmin)/dist_fmax; #----------------------------------------- # input & output file #----------------------------------------- if options.o.startswith('auto') : # base output file on input file outfile = options.f+"."+"aashower"+".root" if options.o == "auto_pwd" : outfile = outfile.split("/")[-1] else : outfile = options.o print (" input file =", options.f ) print (" output file =", outfile ) f = ROOT.EventFile( options.f ) f.header["aashowerfit_detector"] = options.d # writen when we open output file f.set_output( outfile ) f.output_file.ls() #----------------------------------------------- # dynamic calibrations / detector file. #----------------------------------------------- dyncalib = None if options.P != 'none' or options.Q != 'none' : # we are going to do dynamic calibration det = Det( options.d ) if options.P == 'none' or options.Q == 'none' : print ("wrong options for dynamic calibration.") import sys sys.exit(1) if not 'JPP_DIR' in os.environ: print ("Dynamic Calibration requires JPP to be set up. This seems not to be the case (no $JPP_DIR)") sys.exit(1) # Use Jpp script to get mechanics file and do some checks os.system("getMechanics.sh " + options.d ) fn = f'mechanics_{det.id:08d}.txt' try : print( fn, "contains:", open(fn).read() ) except FileNotFoundError: print ("Failed to get file needed for dynamic calibration", fn ) exit(1) try : dyncalib = ROOT.DynamicCalibration( options.d, options.Q, options.P ) except : print ("failed to use DynamicCalibration; make sure aanet is compiled with JPP support") raise else : # static calibration. if options.d == "auto" : det = Det( f ) else : det = Det( options.d ) det.prnt() if det.doms.size() == 0 : print( "no detector -> exiting.") import sys sys.exit(1) #----------------------------------------------- # 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 ) digpdf = DigitalShowerPdf().set( infile = pdffile, k40prob = options.bgrate_hz*1e-6, eval_mode = DigitalShowerPdf().quick ) digpdf.init() digi_fit = ShowerFit ("digi", digpdf ).set ( verb = options.v, # not used if scipy minimizer minimizer_name ='Minuit2', minimizer_algo ='Migrad', ) digi_fit.fix_vars( 0, 1, 2, 5 ) fitterlist = [ mest_fit, digi_fit ] def mf( x ) : 'a function to call from the scipy minimizer' if True : # with Timer("mx"): df = np.array([1.0, 1.0, 1.0]) r = ctypes.c_double() digpdf.verb = 0 digpdf.eval_gradient_quick ( x[0],x[1],x[2], r, df ) return r.value, df ; #-------------------------------------------------- # Event Loop #-------------------------------------------------- iproc = 0 for evt in f: with ROOT.get_mem_counter("eventloop") : # -------------------------------------------------------------------------- # checks and preselections # -------------------------------------------------------------------------- if options.v > 0 : print (" === event ", evt.id , " ===") # update the detector alignment to the current time. if dyncalib : dyncalib.update( det, int(evt.t.AsDouble()) ); # detector geometry check if f.index < options.X : print("checking") ok = det.check( evt.hits ) if not ok : print ('Detector check failed!') print ("This means the hit position/directions in your input file are not consistent with the supplied") print ("detector. If you know what you're doing, you can apply the detector to the evets with the -R option.") if not options.R : sys.exit(1) else : print ("detector check passed!") # minimum mc neutrino energy requirement. if options.E > 0 and evt.neutrino() and evt.neutrino().E < options.E : continue # pre-selection on minimum number of triggered hits if options.T > 0 and sum( (h.trig!=0) for h in evt.hits ) < options.T : f.write() continue # count actually processed (not skipped) events iproc+=1 if iproc > options.n : print ('requested number of events reached') break with Timer("0) events") as event_timer : # from now on, there may be no 'continue' statements util write with Timer("1) calibration"): if options.hit_type == 'raw' : # Make the hit times more numerically palletable input_hits = copy(evt.hits) if options.R : det.apply( input_hits ) hit_t0 = min( h.t for h in input_hits ) - 1e-4 # prevent == 0 for h in input_hits : h.t -= hit_t0 if options.hit_type == 'mc' : input_hits = evt.mc_hits if options.R : det.apply( input_hits ) if options.m : hits = ROOT.merge_hits ( input_hits ) for h in hits: h.a = 1 else: hits = input_hits 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 last_id = max( (trk.id for trk in evt.trks), default = 0 ) T0 = Trk().set( id = last_id + 1, pos = pivot_hit.pos, t = pivot_hit.t, E = options.estart, rec_type = rec_type, rec_stages = aa.make_vector(recstages[:1]), comment = "start", mother_id = ROOT.TRK_MOTHER_NONE ) mest_hits = ROOT.filter_hits( coin_hits, T0, options.mest_hit_mint, options.mest_hit_maxt ); with Timer("4) vertex fit"): T1 = mest_fit.fit( T0, det, mest_hits, 1 ) T1.set( id = T0.id + 1, comment = "mest", rec_type = rec_type, rec_stages = aa.make_vector( recstages[:2] ), mother_id = T0.id ) # compute error matrix -- since *^*^$ ROOT wont do it for us TTT = copy(T1) def likfunc( x,y,z,t ) : TTT.pos.set(x,y,z) TTT.t = t return mpdf1.eval( TTT ) hess = numpy.array( pyroot_util.hessian( likfunc )( T1.pos.x, T1.pos.y, T1.pos.z, T1.t ) ) set_error_matrix( T1, hess, 0 ) if options.c : evt.trks.clear() #------------------------------------------------------------------ # likelihood fit, with starting-direction according to icosahedron #------------------------------------------------------------------ with Timer("5) final hit selection"): digi_hits = ROOT.filter_hits( input_hits, T1 , options.digi_hit_mint , options.digi_hit_maxt, options.digi_hit_maxr ) # the maximum distance to vertex for selecting doms/hits for final fit. digi_fit.max_dist = max_dist_func( det, digi_hits, T1.pos ) digpdf.force_precompute = True with Timer("6) direction fit"): if options.l : fits_ = [ T1 ] elif options.q : # use scipy miminizer digpdf.init_event( det, digi_hits, T1.pos, digi_fit.max_dist ) if digpdf.domsums.size() < 10 : f.write() continue x = np.array([1.0, 1.0, 1.0]) fits_, best = [], None start_dirs = ROOT.icosahedron(); for ii, T1.dir in enumerate( start_dirs ): 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': options.minimizer_gtol , 'eps': 1.4901161193847656e-08, 'maxiter': None, 'disp': False, 'return_all': False} 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] ) 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 # Set error matrix only for the the best one (it takes somewhat non-neglible cpu) # important to allocate on python side (not c++) for some reason (otherwise memory leak) hessian = np.empty([9], dtype=float) digpdf.fill_hessian( hessian, best.res.x[0], best.res.x[1], best.res.x[2] ) hessian.shape = 3,3 set_error_matrix( best, hessian, 4 ) else : # ROOT-fitter (buggy in root 6.22 hope for better in 6.24++) fits_ = [] for ii, T1.dir in enumerate( ROOT.icosahedron() ): do_init_pdf = ii==0 fit = digi_fit.fit( T1, det , digi_hits, do_init_pdf ) fit.comment = "fit"+str(ii) fits_.append( fit ) # invert the likelihoods to comply with the larger-value-is-better schema # that we have adopted in km3net-dataformats. And best-track first in the list. # shift back the time of the prefits T0.t += hit_t0 T1.t += hit_t0 evt.trks.push_back(T0) evt.trks.push_back(T1) for t in fits_ : t.lik *= -1 for i, t in enumerate( sorted( fits_, key = lambda t : t.lik ) ): t.set ( id = T1.id + 1 + i, mother_id = T1.id, rec_stages = aa.make_vector(recstages[:3]), fitinf = aa.make_vector([1.0 * t.E, digi_hits.size()]), status = ROOT.TRK_ST_FINALSTATE ) # from defitions/trkmembers correct_energy(t) # uncorrected energy saved in fitinf[0] t.t += hit_t0 # apply time shift f.evt.trks.push_back( t ) with Timer("7 post processing + writing "): # for the best track , fill the hit_ids besttrk = f.evt.trks[ f.evt.trks.size()-1] besttrk.hit_ids.clear() for h in digi_hits : besttrk.hit_ids.push_back( h.id ) evt.setusr("aashowerfit_time", event_timer.time_since_last_start() ) f.write(); evt.delusr("aashowerfit_time") evt.trks.clear() gc.collect() if iproc%options.t == 0 : print( evt ) Timer.report() ROOT.mem_report() f.report() if options.M : print ("writing meta data.") ourmeta = aa.make_meta( options ) f.meta_data.push_back( ourmeta ) f.write_meta_data() if options.S : print ("copying summary slice tree to output file.") ROOT.prepare_summary_frame( f.rootfile() ) s = f.rootfile().Get("KM3NET_SUMMARYSLICE") f.output_file.cd() if not s: print ("did not find KM3NET_SUMMARYSLICE") else : s.CloneTree().Write() print ("bye")