diff --git a/CHANGELOG.md b/CHANGELOG.md index b5daa4285..3deab4014 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,9 @@ it in future. ### Added +* New analysis toolkit prototype added as part of the 'experimental' package. +* Simple analysis example script now available in 'examples/' + ### Fixed * Use ConstructedAt instead of remove pythonization for TClonesArray diff --git a/examples/analysis_example.py b/examples/analysis_example.py new file mode 100644 index 000000000..3929c9751 --- /dev/null +++ b/examples/analysis_example.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python +"""Example script for usage of the analysis_toolkit for signal selection .""" + +from argparse import ArgumentParser + +import ROOT +import rootUtils as ut +from experimental import analysis_toolkit + +parser = ArgumentParser(description=__doc__) +parser.add_argument("--path", help="Path to simulation file", default="./") +options = parser.parse_args() + + +f = ROOT.TFile.Open(options.path + "/ship.conical.Pythia8-TGeant4_rec.root", "read") +tree = f.cbmsim + +geo_file = ROOT.TFile.Open( + options.path + "/geofile_full.conical.Pythia8-TGeant4.root", "read" +) + +selection = analysis_toolkit.selection_check(geo_file) + + +def main(): + """Sample function to analyse the pre-selection parameters.""" + hist_dict = {} + + ut.bookHist(hist_dict, "event_weight", " Event weight; ; ", 100, 100, 100) + ut.bookHist( + hist_dict, + "candidate_time", + " candidate time @ decay vertex; ns ; ", + 200, + 100, + 100, + ) + ut.bookHist(hist_dict, "impact_parameter", "Impact Parameter;cm;", 200, 100, 100) + ut.bookHist( + hist_dict, "dist_to_innerwall", "Distance to Innerwall;cm;", 200, 100, 100 + ) + ut.bookHist( + hist_dict, + "dist_to_vesselentrance", + "Distance to Decay Vessel Entrance;cm;", + 200, + 100, + 100, + ) + ut.bookHist(hist_dict, "inv_mass", "Invariant mass ;GeV;", 200, 100, 100) + ut.bookHist(hist_dict, "DOCA", "Distance of Closest Approach ;cm;", 200, 100, 100) + ut.bookHist( + hist_dict, + "len_Particles", + " len(tree.Particles) ;Number of candidates per event;", + 200, + 100, + 100, + ) + ut.bookHist( + hist_dict, + "d_mom", + "momentum of Daughters ;d1 (GeV);d2 (GeV)", + 200, + 100, + 100, + 200, + 100, + 100, + ) + ut.bookHist(hist_dict, "nDOF", "nDOF ;d1;d2", 200, 100, 100, 200, 100, 100) + ut.bookHist(hist_dict, "chi2nDOF", "chi2nDOF ;d1;d2", 200, 100, 100, 200, 100, 100) + + for event_nr, event in enumerate(tree): + selection.access_event(event) + if len(event.Particles) == 0: + continue + + event_weight = event.MCTrack[2].GetWeight() + + hist_dict["event_weight"].Fill(event_weight) + + for candidate_id_in_event, signal in enumerate(event.Particles): + hist_dict["candidate_time"].Fill(selection.define_candidate_time(signal)) + hist_dict["impact_parameter"].Fill(selection.impact_parameter(signal)) + hist_dict["dist_to_innerwall"].Fill(selection.dist_to_innerwall(signal)) + hist_dict["dist_to_vesselentrance"].Fill( + selection.dist_to_vesselentrance(signal) + ) + hist_dict["DOCA"].Fill(selection.DOCA(signal)) + hist_dict["inv_mass"].Fill(selection.invariant_mass(signal)) + hist_dict["len_Particles"].Fill(len(tree.Particles)) + hist_dict["d_mom"].Fill(*selection.daughtermomentum(signal)) + hist_dict["nDOF"].Fill(*selection.nDOF(signal)) + hist_dict["chi2nDOF"].Fill(*selection.chi2nDOF(signal)) + + IP_cut = 250 + preselection_flag = selection.preselection_cut( + signal, IP_cut=IP_cut, show_table=False + ) + + if preselection_flag: + print( + f"Event:{event_nr} Candidate_index: {candidate_id_in_event} <--passes the pre-selection\n\n" + ) + else: + print(f"Event:{event_nr} Candidate_index: {candidate_id_in_event} \n\n") + + ut.writeHists(hist_dict, "preselectionparameters.root") + + +if __name__ == "__main__": + main() diff --git a/python/experimental/__init__.py b/python/experimental/__init__.py new file mode 100644 index 000000000..4821959d6 --- /dev/null +++ b/python/experimental/__init__.py @@ -0,0 +1,5 @@ +""" +This is the 'experimental' package, which contains scripts that are still being tested and may undergo significant changes. + +Suggestions & feedback are highly appreciated as this is a work-in-progress. +""" diff --git a/python/experimental/analysis_toolkit.py b/python/experimental/analysis_toolkit.py new file mode 100644 index 000000000..436d8ca37 --- /dev/null +++ b/python/experimental/analysis_toolkit.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python +"""Toolkit for Analysis.""" + +import numpy as np +import ROOT +import shipunit as u +import yaml +from rootpyPickler import Unpickler +from ShipGeoConfig import AttrDict +from tabulate import tabulate + + +class selection_check: + """Class to perform various selection checks on the candidate.""" + + def __init__(self, geo_file): + """Initialize the selection_check class with geometry and configuration.""" + self.geometry_manager = geo_file.FAIRGeom + unpickler = Unpickler(geo_file) + self.ship_geo = unpickler.load("ShipGeo") + + fairship = ROOT.gSystem.Getenv("FAIRSHIP") + + if self.ship_geo.DecayVolumeMedium == "helium": + with open(fairship + "/geometry/veto_config_helium.yaml", "r") as file: + config = yaml.safe_load(file) + self.veto_geo = AttrDict(config) + self.veto_geo.z0 + if self.ship_geo.DecayVolumeMedium == "vacuums": + with open(fairship + "/geometry/veto_config_vacuums.yaml", "r") as file: + config = yaml.safe_load(file) + self.veto_geo = AttrDict(config) + + def access_event(self, tree): + """Access event data.""" + self.tree = tree + + def define_candidate_time(self, candidate): + """Calculate time associated with the candidate decay vertex using strawtubes MCPoint info.""" + t0 = self.tree.ShipEventHeader.GetEventTime() + candidate_pos = ROOT.TVector3() + candidate.GetVertex(candidate_pos) + + d1, d2 = candidate.GetDaughter(0), candidate.GetDaughter(1) + d1_mc, d2_mc = self.tree.fitTrack2MC[d1], self.tree.fitTrack2MC[d2] + + time_vtx_from_strawhits = [] + + for hit in self.tree.strawtubesPoint: + if not ( + int(str(hit.GetDetectorID())[:1]) == 1 + or int(str(hit.GetDetectorID())[:1]) == 2 + ): + continue + + if not (hit.GetTrackID() == d1_mc or hit.GetTrackID() == d2_mc): + continue + + t_straw = hit.GetTime() + + dist = np.sqrt( + (candidate_pos.X() - hit.GetX()) ** 2 + + (candidate_pos.Y() - hit.GetY()) ** 2 + + (candidate_pos.Z() - hit.GetZ()) ** 2 + ) # distance to the vertex in cm + + d_mom = self.tree.MCTrack[hit.GetTrackID()].GetP() / u.GeV + mass = self.tree.MCTrack[hit.GetTrackID()].GetMass() + v = u.c_light * d_mom / np.sqrt(d_mom**2 + (mass) ** 2) + + t_vertex = t_straw - (dist / v) + + time_vtx_from_strawhits.append(t_vertex) + + t_vtx = np.average(time_vtx_from_strawhits) + t0 + + return t_vtx # units in ns + + def impact_parameter(self, candidate): + """Calculate the impact parameter of the candidate relative to (0,0,target.z0).""" + candidate_pos = ROOT.TVector3() + candidate.GetVertex(candidate_pos) + + candidate_mom = ROOT.TLorentzVector() + candidate.Momentum(candidate_mom) + target_point = ROOT.TVector3(0, 0, self.ship_geo.target.z0) + + projection_factor = 0 + if hasattr(candidate_mom, "P"): + P = candidate_mom.P() + else: + P = candidate_mom.Mag() + for i in range(3): + projection_factor += ( + candidate_mom(i) / P * (target_point(i) - candidate_pos(i)) + ) + + dist = 0 + for i in range(3): + dist += ( + target_point(i) + - candidate_pos(i) + - projection_factor * candidate_mom(i) / P + ) ** 2 + dist = ROOT.TMath.Sqrt(dist) + + return dist # in cm + + def dist_to_innerwall(self, candidate): + """Calculate the minimum distance(in XY plane) of the candidate decay vertex to the inner wall of the decay vessel. If outside the decay volume, or if distance > 100cm,Return 0.""" + candidate_pos = ROOT.TVector3() + candidate.GetVertex(candidate_pos) + position = (candidate_pos.X(), candidate_pos.Y(), candidate_pos.Z()) + + nsteps = 8 + dalpha = 2 * ROOT.TMath.Pi() / nsteps + min_distance = float("inf") + + node = self.geometry_manager.FindNode(*position) + if not node: + return 0 # is outside the decay volume + + # Loop over directions in the XY plane + for n in range(nsteps): + alpha = n * dalpha + direction = ( + ROOT.TMath.Sin(alpha), + ROOT.TMath.Cos(alpha), + 0.0, + ) # Direction vector in XY plane + self.geometry_manager.InitTrack(*position, *direction) + if not self.geometry_manager.FindNextBoundary(): + continue + # Get the distance to the boundary and update the minimum distance + distance = self.geometry_manager.GetStep() + min_distance = min(min_distance, distance) + + return min_distance if min_distance < 100 * u.m else 0 + + def dist_to_vesselentrance(self, candidate): + """Calculate the distance of the candidate decay vertex to the entrance of the decay vessel.""" + candidate_pos = ROOT.TVector3() + candidate.GetVertex(candidate_pos) + return candidate_pos.Z() - self.veto_geo.z0 + + def nDOF(self, candidate): + """Return the number of degrees of freedom (nDOF) for the particle's daughter tracks.""" + nmeas = [] + t1, t2 = candidate.GetDaughter(0), candidate.GetDaughter(1) + + for tr in [t1, t2]: + fit_status = self.tree.FitTracks[tr].getFitStatus() + nmeas.append( + int(round(fit_status.getNdf())) + ) # nmeas.append(fit_status.getNdf()) + + return np.array(nmeas) + + def daughtermomentum(self, candidate): + """Return the momentum(Mag) of the particle's daughter tracks.""" + daughter_mom = [] + t1, t2 = candidate.GetDaughter(0), candidate.GetDaughter(1) + for trD in [t1, t2]: + x = self.tree.FitTracks[trD] + xx = x.getFittedState() + daughter_mom.append((xx.getMom().Mag())) + + return np.array(daughter_mom) + + def invariant_mass(self, candidate): + """Invariant mass of the candidate.""" + return candidate.GetMass() + + def DOCA(self, candidate): + """Distance of Closest Approach.""" + return candidate.GetDoca() + + def is_in_fiducial(self, candidate): + """Check if the candidate decay vertex is within the Fiducial Volume.""" + candidate_pos = ROOT.TVector3() + candidate.GetVertex(candidate_pos) + + if candidate_pos.Z() > self.ship_geo.TrackStation1.z: + return False + if candidate_pos.Z() < self.veto_geo.z0: + return False + + # if self.dist2InnerWall(candidate)<=5*u.cm: return False + + vertex_node = ROOT.gGeoManager.FindNode( + candidate_pos.X(), candidate_pos.Y(), candidate_pos.Z() + ) + vertex_elem = vertex_node.GetVolume().GetName() + if not vertex_elem.startswith("DecayVacuum_"): + return False + return True + + def chi2nDOF(self, candidate): + """Return the reduced chi^2 of the particle's daughter tracks.""" + t1, t2 = candidate.GetDaughter(0), candidate.GetDaughter(1) + + chi2ndf = [] + for tr in [t1, t2]: + fit_status = self.tree.FitTracks[tr].getFitStatus() + chi2ndf.append(fit_status.getChi2() / fit_status.getNdf()) + + return np.array(chi2ndf) + + def preselection_cut(self, candidate, IP_cut=250, show_table=False): + """ + Umbrella method to apply the pre-selection cuts on the candidate. + + show_table=True tabulates the pre-selection parameters. + """ + flag = True + + if len(self.tree.Particles) != 1: + flag = False + if not (self.is_in_fiducial(candidate)): + flag = False + if self.dist_to_innerwall(candidate) <= 5 * u.cm: + flag = False + if self.dist_to_vesselentrance(candidate) <= 100 * u.cm: + flag = False + if self.impact_parameter(candidate) >= IP_cut * u.cm: + flag = False + if self.DOCA(candidate) >= 1 * u.cm: + flag = False + if np.any(self.nDOF(candidate) <= 25): + flag = False + if np.any(self.chi2nDOF(candidate) >= 5): + flag = False + if np.any(self.daughtermomentum(candidate) <= 1 * u.GeV): + flag = False + + if show_table: + table = [ + [ + "Number of candidates in event", + len(self.tree.Particles), + "==1", + len(self.tree.Particles) == 1, + ], + [ + "Time @ decay vertex (ns)", + self.define_candidate_time(candidate), + "", + "", + ], + [ + "Impact Parameter (cm)", + self.impact_parameter(candidate), + f"IP < {IP_cut*u.cm} cm", + self.impact_parameter(candidate) < IP_cut * u.cm, + ], + [ + "DOCA (cm)", + self.DOCA(candidate), + "DOCA < 1 cm", + self.DOCA(candidate) < 1 * u.cm, + ], + [ + "Is within Fiducial Volume?", + self.is_in_fiducial(candidate), + "True", + self.is_in_fiducial(candidate), + ], + [ + "Dist2InnerWall (cm)", + self.dist_to_innerwall(candidate), + "> 5 cm", + self.dist_to_innerwall(candidate) > 5 * u.cm, + ], + [ + "Dist2VesselEntrance (cm)", + self.dist_to_vesselentrance(candidate), + "> 100 cm", + self.dist_to_vesselentrance(candidate) > 100 * u.cm, + ], + ["Invariant Mass (GeV)", self.invariant_mass(candidate), "", ""], + [ + "Daughter Momentum [d1, d2] (GeV)", + self.daughtermomentum(candidate), + "> 1 GeV", + np.all(self.daughtermomentum(candidate) > 1 * u.GeV), + ], + [ + "Degrees of Freedom [d1, d2]", + self.nDOF(candidate), + "> 25", + np.all(self.nDOF(candidate) > 25), + ], + [ + "Reduced Chi^2 [d1, d2]", + self.chi2nDOF(candidate), + "< 5", + np.all(self.chi2nDOF(candidate) < 5), + ], + ["\033[1mPre-selection passed:\033[0m", "", "", flag], + ] + + for row in table: + row[3] = ( + f"\033[1;32m{row[3]}\033[0m" + if row[3] + else f"\033[1;31m{row[3]}\033[0m" + ) # Green for True, Red for False + + print( + tabulate( + table, + headers=[ + "Parameter", + "Value", + "Pre-selection cut", + "Pre-selection Check", + ], + tablefmt="grid", + ) + ) + return flag