#!/usr/bin/env python """ usage : vertex_fit.py -f input_file Vertex_fit is a reconstruction program for cascade events. options: -f -input_file : Input file in any format aanet can read. default=/sps/km3net/repo/mc/atm_neutrino/KM3NeT_-00000001_20171212/v5.1/reco/mcv5.1.genhen_nueCC.sirene.jte.jchain.aashower.1.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_-00000001_20171212.detx -n : process this many events default=100000000 -v : verbose (0,1 or 2) default=0 -h : display this message default=False -i : start python prompt when done default=False -E : minimum MC energy default=0.0 -e : elongation steps default=10 -m : use MC hits default=True -s : use MC start position default=True -b : blurring of the PDF default=0.0 -p : use first hit PDF default=True -F : energy fraction sample default=0.0 -t : starting time of the fit default=0.0 -X : distance shift wrt true default=0.0 """ print("vertex fit starts...") import ROOT, os, sys, collections, array from math import * from copy import copy from array import array import numpy as np import scipy.optimize try : import aa except ImportError: print("Failed to import aa; Make sure AADIR is in your PYHTONPATH!") sys.exit() import rec from rec import rec2 options = aa.Options( __doc__, sys.argv[1:] ) print(options) from ROOT import vector, Hit, Det, Evt, Trk, Vec, Pmt, EventFile, Timer, TFile, TGraph, TH1D, TH3D, angle_between, TTree, gROOT, best_reco_track # from ROOT import DigitalShowerPdf, MestShowerPdf, ShowerFit #----------------------------------------------------- # Fancy-smancy setting of members #----------------------------------------------------- # for T in Trk, DigitalShowerPdf, MestShowerPdf, ShowerFit : T.set = aa.__set #-------------------------------------------------- # Time slewing #-------------------------------------------------- aa.include("/pbs/throng/km3net/src/Jpp/master/software/JTrigger/JGetRiseTime.hh") slew = ROOT.JTRIGGER.JGetRiseTime() #-------------------------------------------------- # Constants #-------------------------------------------------- water_index = 1.3800851282 dndl = 0.0298 c_light = 299792458*1E-9 v_light = c_light/(water_index) jshowerfile = ["/sps/km3net/users/jseneca/jpp/data/J14p.dat",\ "/sps/km3net/users/jseneca/jpp/data/J13p.dat"] min_residual = -100 max_residual = 900 if options.m: R_bg = 0 else: R_bg = 7e3 #----------------------------------------- # input & output file #----------------------------------------- if options.o.startswith('auto') : # base output file on input file outfile = options.f+"."+"vertex"+".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.set_output( outfile ) f.autosave = False # manually decide which events to write #----------------------------------------------- # detector file #----------------------------------------------- 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) # 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["vertex_fit_detector"] = options.d f.output_file.WriteObject( f.header ,"Head" ) f.output_file.ls() #----------------------------------------------- # Helper functions #----------------------------------------------- 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; def find_lepton(PDG_number, mc_trks): i = 0 for track in mc_trks: if track.type == PDG_number: return i i+=1 return None def is_contained(track): r = sqrt(track.pos.x**2 + track.pos.y**2) if track.pos.z <= (677 - 91) and track.pos.z >= (0 + 91) and r <= (530 - 91): return True else: return False def max_dist_func ( det, hits, pos ) : """ Return the radius of a sphere that contains almost all hits. """ 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; #-------------------------------------------------- # setup the ShowerPdf #-------------------------------------------------- L = ROOT.CascadeLikelihood( options.p, # first hit probability options.b, # sigma blurring of pdf options.e, # energy elongation steps R_bg, min_residual, max_residual, options.F, # extra sampling point at the vertex with energy fraction F 1 ) #verbosity # #-------------------------------------------------- # # Let's go # #-------------------------------------------------- iproc=0 for evt in f: # minimum mc neutrino energy requirement. if options.E > 0 and evt.neutrino() and evt.neutrino().E < options.E : continue if is_contained(evt.mc_trks[0]) != True: continue # count actually processed (not skipped) events iproc+=1 if iproc > options.n : print ('requested number of events reached') break # if options.v > 0 : print (" ============= event ", evt.id , " =============") nu = evt.mc_trks[0] with Timer("calibration"): det.apply(evt) start_trk = Trk() if options.s: start_trk = nu else: start_trk = best_reco_track(evt, 101) with Timer("hit_selection"): if options.m: hits = evt.mc_hits max_dist = 1e10 else: hits = evt.hits for hit in hits: hit.t -= slew( hit.tot ) max_dist = max_dist_func( det, hits, start_trk.pos ) max_dist = 800 if not options.m and options.s: offset = evt.mc_t - evt.t.AsDouble()*1e9 for hit in evt.hits : hit.t -= offset vertex_hits = ROOT.select_hits( hits, start_trk, min_residual, max_residual, max_dist, options.p ); with Timer("vertex fit"): pmts = ROOT.vector(Pmt)() for i,pmt in det.pmts : pmts.push_back( pmt ) print("vertex hits", len(vertex_hits)) L.set_event( pmts, vertex_hits ) # starting point x = np.array([start_trk.pos.x + options.X, start_trk.pos.y + options.X, start_trk.pos.z + options.X, start_trk.t + options.t]) print("starting point: ", x) print("nu.pos", nu.pos) print("nu.dir", nu.dir) # track to be fitted test_trk = Trk() test_trk.dir = nu.dir test_trk.E = nu.E # function to be minimized def mf( x ) : with Timer("mx"): # function found in rec/JppShowerPdf test_trk.pos = ROOT.Vec( x[0], x[1], x[2] ) test_trk.t = x[3] # return pdf4d.eval_L_vertex( test_trk, vertex_hits, options.p, False, options.e ) return -L.eval( test_trk ) res = scipy.optimize.minimize(mf, x, method = 'Nelder-Mead') print(res) T = Trk() T.dir = nu.dir T.pos = ROOT.Vec(res.x[0], res.x[1], res.x[2]) T.t = res.x[3] T.E = evt.mc_trks[0].E T.rec_type = 555 T.comment = "tstart=" + str(options.t) T.lik = L.eval( T ) T.fitinf.push_back(options.t) T.fitinf.push_back(len(vertex_hits)) # T = Trk().set( id = 100, # pos = Vec(res.x[0], res.x[1], res.x[2]), # dir = start_trk.dir, # t = res.x[3], # E = start_trk.E, # rec_type = 1010, # comment = "vertex_fit" ) evt.trks.push_back(T) d = start_trk.pos - evt.mc_trks[0].pos d_along = d.dot( evt.mc_trks[0].dir ) d_perp = (d - d_along * evt.mc_trks[0].dir).len() print("d_along_start", d_along) print("d_perp_start", d_perp) print("lik start", -L.eval(start_trk) ) d = T.pos - evt.mc_trks[0].pos d_along = d.dot( evt.mc_trks[0].dir ) d_perp = (d - d_along * evt.mc_trks[0].dir).len() print("d_along", d_along) print("d_perp", d_perp) print("lik end", -L.eval( T ) ) print("fitted time", T.t) f.write(); evt.trks.clear() Timer.report() f.output_file.Write() # todo: is this really needed? del f print ("bye")