Skip to content

Commit

Permalink
tabulate pre-selection parameters
Browse files Browse the repository at this point in the history
  • Loading branch information
@anupama-reghunath
anupama-reghunath committed Nov 29, 2024
1 parent 7d891d6 commit e8254a8
Show file tree
Hide file tree
Showing 2 changed files with 204 additions and 61 deletions.
111 changes: 88 additions & 23 deletions macro/analysis_example.py
View file
Original file line number Diff line number Diff line change
@@ -1,41 +1,106 @@
#!/usr/bin/env python
"""Example script for usage of the analysis_toolkit for signal selection ."""

from argparse import ArgumentParser

import analysis_toolkit
import ROOT
import rootUtils as ut

parser = ArgumentParser(description=__doc__)
parser.add_argument("--path", help="Path to simulation file", default="./")
options = parser.parse_args()

geo_file = "geofile_full.conical.Pythia8-TGeant4.root"
input_file = "ship.conical.Pythia8-TGeant4_rec.root"

f = ROOT.TFile.Open(input_file)
geo_file = options.path + "/geofile_full.conical.Pythia8-TGeant4.root"
input_file = options.path + "/ship.conical.Pythia8-TGeant4_rec.root"

f = ROOT.TFile.Open(input_file, "read")
tree = f.cbmsim

selection = analysis_toolkit.selection_check(geo_file)

for event_nr, event in enumerate(tree):
if event_nr > 200:
break
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

for candidate_id_in_event, signal in enumerate(tree.Particles):
print(f"Event:{event_nr} Candidate_index: {candidate_id_in_event}")

print(f"\t vertex time: {selection.define_candidate_time(signal)} ns")
print(f"\t Impact Parameter: {selection.impact_parameter(signal)} cm")
print(f"\t is within Fiducial volume: {selection.is_in_fiducial(signal)}")
print(f"\t\tDist2InnerWall: {selection.dist_to_innerwall(signal)} cm")
print(f"\t\tDist2EntranceLid: {selection.dist_to_entrance_vessel(signal)} cm")
print(f"\t Daughter momentum>1 GeV: {selection.daughtermomentum(signal)} GeV")
print(f"\t Invariant mass: {selection.invariant_mass(signal)} GeV")
print(f"\t Degrees of Freedom ([d1,d2],Cut(>25)): {selection.nDOF(signal)}")
print(f"\t Reduced Chi^2 ([d1,d2],Cut(<5): {selection.chi2nDOF(signal)}")
print(f"\t DOCA: {selection.DOCA(signal)} cm")

print(
f"\t Preselection Cut passed:{selection.preselection_cut(signal,IP_cut=250)}"
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=True
)
print("------------------------------------------------------------")

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")
154 changes: 116 additions & 38 deletions python/analysis_toolkit.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import yaml
from rootpyPickler import Unpickler
from ShipGeoConfig import AttrDict
from tabulate import tabulate


class selection_check:
Expand Down Expand Up @@ -63,7 +64,7 @@ def define_candidate_time(self, candidate):
(candidate_pos.X() - hit.GetX()) ** 2
+ (candidate_pos.Y() - hit.GetY()) ** 2
+ (candidate_pos.Z() - hit.GetZ()) ** 2
) # distance to the vertex #in cm
) # distance to the vertex in cm

d_mom = self.tree.MCTrack[hit.GetTrackID()].GetP() / u.GeV
mass = self.tree.MCTrack[hit.GetTrackID()].GetMass()
Expand Down Expand Up @@ -138,39 +139,35 @@ def dist_to_innerwall(self, candidate):

return min_distance if min_distance < 100 * u.m else 0

def dist_to_entrance_vessel(self, candidate):
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."""
flag = True
nmeas = []
t1, t2 = candidate.GetDaughter(0), candidate.GetDaughter(1)
ndf_cut = 25 # cut on nDOF

for tr in [t1, t2]:
fit_status = self.tree.FitTracks[tr].getFitStatus()
nmeas.append(
int(round(fit_status.getNdf()))
) # nmeas.append(fit_status.getNdf())
if nmeas[-1] <= ndf_cut:
flag = False # dof<=25=False
return nmeas, flag

return np.array(nmeas)

def daughtermomentum(self, candidate):
"""Return the momentum(Mag) of the particle's daughter tracks."""
flag = True
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()))
if daughter_mom[-1] < 1 * u.GeV:
flag = False
return daughter_mom, flag

return np.array(daughter_mom)

def invariant_mass(self, candidate):
"""Invariant mass of the candidate."""
Expand Down Expand Up @@ -203,39 +200,120 @@ def is_in_fiducial(self, candidate):
def chi2nDOF(self, candidate):
"""Return the reduced chi^2 of the particle's daughter tracks."""
t1, t2 = candidate.GetDaughter(0), candidate.GetDaughter(1)
# cut on chi2nDOF
flag = True
chi2ndf_cut = 5

chi2ndf = []
for tr in [t1, t2]:
fit_status = self.tree.FitTracks[tr].getFitStatus()
chi2ndf.append(fit_status.getChi2() / fit_status.getNdf())
if chi2ndf[-1] >= chi2ndf_cut:
flag = False
return chi2ndf, flag

def preselection_cut(self, candidate, IP_cut=250):
return np.array(chi2ndf)

def preselection_cut(self, candidate, IP_cut=250, show_table=True):
"""Umbrella method to apply the pre-selection cuts on the candidate."""
flag = True

if len(self.tree.Particles) != 1:
return False
if self.dist_to_innerwall(candidate) <= 5 * u.cm:
return False
flag = False
if not (self.is_in_fiducial(candidate)):
return False
if self.dist_to_entrance_vessel(candidate) < 100 * u.cm:
return False
if self.impact_parameter(candidate) > IP_cut:
return False
nmeas, flag = self.nDOF(candidate)
if not (flag):
return False
chi2ndf, flag = self.chi2nDOF(candidate)
if not (flag):
return False
d_mom, flag = self.daughtermomentum(candidate)
if not (flag):
return False
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:
return False

return True
flag = False
if np.any(self.nDOF(candidate) <= 25 * u.cm):
flag = False
if np.any(self.chi2nDOF(candidate) >= 5 * u.cm):
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

0 comments on commit e8254a8

Please sign in to comment.