from typing import List
import pandas as pd
import numpy as np
import math
import ROOT
from ROOT import Belle2
import basf2 as b2
import modularAnalysis as ma
import variables.collections as vc
from variables import variables as vm
import variables.utils as vu
import vertex as vtx
import argparse
import glob
import sys
from vibe.core.utils.misc import fancy_validation_mode_header
from vibe.core.validation_mode import ValidationModeBaseClass, add_output_formats
from vibe.core.helper.histogram_tools import HistVariable, Histogram, HistComponent
from vibe.core.helper.root_helper import makeROOTCompatible
[docs]
@fancy_validation_mode_header
class DimuonTrackingResolution(ValidationModeBaseClass):
name = "DimuonTrackingResolution"
[docs]
def create_basf2_path(self):
main = b2.Path()
# uncomment to print ERRORS only
#b2.set_log_level(b2.LogLevel.ERROR)
main.add_module('Progress')
track_cuts = 'abs(d0) < 2.0 and abs(z0) < 4.0 and pt > 2.0 and useCMSFrame(p) > 0.5'
line_mumu_tight_or_highm = "SoftwareTriggerResult(software_trigger_cut&skim&accept_mumu_tight_or_highm)"
line_mumutight = "SoftwareTriggerResult(software_trigger_cut&skim&accept_mumutight)"
line_mumu_2trk = "SoftwareTriggerResult(software_trigger_cut&skim&accept_mumu_2trk)"
eventCutString = f"[{line_mumu_tight_or_highm} == 1]"
eventCutString += f" or [[isNAN({line_mumu_tight_or_highm}) == 1] and [{line_mumutight} == 1]]"
eventCutString += f" or [[isNAN({line_mumutight}) == 1] and [{line_mumu_2trk} == 1]]"
ma.applyEventCuts(cut=eventCutString, path=main)
vm.addAlias('cdcecl_veto','L1Input(cdcecl_veto)')
vm.addAlias('passive_veto','L1Input(passive_veto)')
ma.fillParticleList('gamma:HLT', 'E>0.1', path=main)
ma.fillParticleList('mu+:HLT', track_cuts, path=main)
vm.addAlias('Etot', 'formula( [daughter(0,clusterE) + daughter(1,clusterE) + totalEnergyOfParticlesInList(gamma:HLT)] )')
ma.reconstructDecay('vpho:mumu -> mu+:HLT mu-:HLT', 'Etot < 2', path=main)
vm.addAlias('nCleanedTracks',f'nCleanedTracks({track_cuts})')
vtx.treeFit('vpho:mumu', updateAllDaughters=False, path=main)
vm.addAlias("micro_d0", "formula(1e4*d0)")
vm.addAlias("micro_z0", "formula(1e4*z0)")
vm.addAlias("micro_d0FromIP", "formula(1e4*d0FromIP)")
vm.addAlias("micro_z0FromIP", "formula(1e4*z0FromIP)")
cm_vars = vu.create_aliases(["E", "p", "pt", "pz", "cosTheta", "phi"], "useCMSFrame({variable})", "CM")
VARS = ['cdcecl_veto','passive_veto', 'IPX', 'IPY', 'IPZ', "nTracks", 'nCleanedTracks'] + vc.vertex
VARS += vu.create_aliases_for_selected(vc.track + vc.kinematics + vc.track_hits + vc.trackfit_parameters + cm_vars
+ ['cosTheta', 'phi', 'charge', 'trackTime', 'clusterE', 'isTrackFlippedAndRefitted'] + ["micro_d0", "micro_z0", "micro_d0FromIP", "micro_z0FromIP"],
'vpho -> ^mu+ ^mu-', ['mup', 'mum'])
vm.addAlias("ptRes", "formula((mum_CM_pt-mup_CM_pt)*1e3/1.414213562)")
vm.addAlias("d0FromIPRes", "formula((mum_micro_d0FromIP+mup_micro_d0FromIP)/1.414213562)")
vm.addAlias("z0FromIPRes", "formula((mum_micro_z0FromIP-mup_micro_z0FromIP)/1.414213562)")
vm.addAlias("trackTimeDiff", "formula(mum_trackTime - mup_trackTime)")
VARS+= ['d0FromIPRes', 'z0FromIPRes', 'ptRes', 'trackTimeDiff']
VARS = sorted(list(set(VARS)))
vm.addAlias("acosP", 'acos(mup_CM_cosTheta)')
vm.addAlias("acosM", 'acos(mum_CM_cosTheta)')
vm.addAlias("acosDiff", "formula(abs(acosP+acosM))")
goodEvent = 'mup_cosTheta > -0.573576436 and mup_cosTheta < 0.707106781 and mup_pt > 2'
goodEvent += ' and -0.573576436 < mum_cosTheta and mum_cosTheta < 0.707106781 and mum_pt > 2'
goodEvent += ' and abs(mup_CM_phi-mum_CM_phi) > 2.967059728' # // 170°"
goodEvent += ' and abs(acosDiff - 3.141592654) < 0.174532925' #// 10°
goodEvent += ' and nCleanedTracks==2'
ma.applyCuts("vpho:mumu", goodEvent, path=main)
ma.cutAndCopyList("vpho:0P", "vpho:mumu", "mup_nSVDHits <= 0.5", path=main)
ma.cutAndCopyList("vpho:0M", "vpho:mumu", "mum_nSVDHits <= 0.5", path=main)
ma.cutAndCopyList("vpho:1P", "vpho:mumu", "mup_nSVDHits > 0.5", path=main)
ma.cutAndCopyList("vpho:1M", "vpho:mumu", "mum_nSVDHits > 0.5", path=main)
ma.cutAndCopyList("vpho:1MP", "vpho:mumu", "mup_nSVDHits > 0.5 and mum_nSVDHits > 0.5", path=main)
self.variables_to_validation_ntuple(
decay_str="vpho:mumu",
variables=VARS,
path=main,
)
#1D histograms
hRes_1D_vpho_mumu = [
("ptRes", 1000, -500, 500),
("z0FromIPRes", 1000, -500, 500),
("d0FromIPRes", 1000, -200, 200),
]
hP_1D_vpho_mumu = [
("mup_micro_d0", 140, -700, 700),
("mup_micro_z0", 100, -2000, 2000),
("mup_micro_d0FromIP", 140, -700, 700),
("mup_micro_z0FromIP", 100, -2000, 2000),
("mup_CM_p", 140, 0, 7),
("mup_CM_pt", 140, 0, 7),
("mup_nCDCHits", 80, 0, 80),
("mup_nSVDHits", 14, 0, 14),
("mup_nPXDHits", 4, 0, 4),
]
hM_1D_vpho_mumu = [
("mum_micro_d0", 140, -700, 700),
("mum_micro_z0", 100, -2000, 2000),
("mum_micro_d0FromIP", 140, -700, 700),
("mum_micro_z0FromIP", 100, -2000, 2000),
("mum_CM_p", 140, 0, 7),
("mum_CM_pt", 140, 0, 7),
("mum_nCDCHits", 80, 0, 80),
("mum_nSVDHits", 14, 0, 14),
("mum_nPXDHits", 4, 0, 4),
]
h_1D_vpho_mumu_tracks = [
("nTracks", 20, 0, 20),
]
hM_1D_vpho_0 = [("mum_trackTime", 100, -50, 50),]
hP_1D_vpho_0 = [("mup_trackTime", 100, -50, 50),]
hM_1D_vpho_1 = [
("mum_trackTime", 100, -50, 50),
("mum_isTrackFlippedAndRefitted", 2, -0.5, 1.5),
]
hP_1D_vpho_1 = [
("mup_trackTime", 100, -50, 50),
("mup_isTrackFlippedAndRefitted", 2, -0.5, 1.5),
]
h_1D_vpho_1MP = [
("trackTimeDiff", 100, -50, 50),
]
# 2D histograms
h_2D_vpho_mumu = [
("mup_nSVDHits", 14, 0, 14, "mum_nSVDHits", 14, 0, 14),
]
h_2D_vpho_1MP = [
("mup_trackTime", 100, -50, 50, "mum_trackTime", 100, -50, 50),
]
self.variables_to_validation_histogram(
decay_str="vpho:mumu",
variables=hRes_1D_vpho_mumu+hP_1D_vpho_mumu+hM_1D_vpho_mumu+h_1D_vpho_mumu_tracks,
variables_2d=h_2D_vpho_mumu,
path=main,
)
self.variables_to_validation_histogram(
decay_str="vpho:0M",
variables=hM_1D_vpho_0,
variables_2d=None,
path=main,
)
self.variables_to_validation_histogram(
decay_str="vpho:0P",
variables=hP_1D_vpho_0,
variables_2d=None,
path=main,
)
self.variables_to_validation_histogram(
decay_str="vpho:1M",
variables=hM_1D_vpho_1,
variables_2d=None,
path=main,
)
self.variables_to_validation_histogram(
decay_str="vpho:1P",
variables=hP_1D_vpho_1,
variables_2d=None,
path=main,
)
self.variables_to_validation_histogram(
decay_str="vpho:1MP",
variables=h_1D_vpho_1MP,
variables_2d=h_2D_vpho_1MP,
path=main,
)
return main
[docs]
@add_output_formats('root')
def offline_df_manipulation(self, df: pd.DataFrame) -> pd.DataFrame:
'''
mup_helix = []
mup_moved_d0 = []
mup_moved_z0 = []
mum_helix = []
mum_moved_d0 = []
mum_moved_z0 = []
d0Res = []
z0Res = []
for i in range(len(df)):
mup_helix.append(Belle2.Helix(df.mup_d0[i], df.mup_phi0[i], df.mup_omega[i], df.mup_z0[i], df.mup_tanLambda[i]))
mum_helix.append(Belle2.Helix(df.mum_d0[i], df.mum_phi0[i], df.mum_omega[i], df.mum_z0[i], df.mum_tanLambda[i]))
mup_helix[i].passiveMoveBy(df.IPX[i], df.IPY[i], df.IPZ[i])
mum_helix[i].passiveMoveBy(df.IPX[i], df.IPY[i], df.IPZ[i])
mup_moved_d0.append((mup_helix[i].getD0())*1e4)
mum_moved_d0.append((mum_helix[i].getD0())*1e4)
mup_moved_z0.append((mup_helix[i].getZ0())*1e4)
mum_moved_z0.append((mum_helix[i].getZ0())*1e4)
d0Res.append((mum_moved_d0[i]+mup_moved_d0[i])/1.414213562)
z0Res.append((mum_moved_z0[i]-mup_moved_z0[i])/1.414213562)
df['mup_moved_d0'] = mup_moved_d0
df['mup_moved_z0'] = mup_moved_z0
df['mum_moved_d0'] = mum_moved_d0
df['mum_moved_z0'] = mum_moved_z0
df['d0Res'] = d0Res
df['z0Res'] = z0Res
'''
return df
@property
def analysis_validation_histograms(self) -> List[Histogram]:
return [
Histogram(
name="hD0",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"Original d0",
unit=r"$\mu m$",
bins=140,
scope=(-700, 700),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_micro_d0",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_micro_d0",
),
],
),
Histogram(
name="hMovedD0",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"Moved d0",
unit=r"$\mu m$",
bins=140,
scope=(-700, 700),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_micro_d0FromIP",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_micro_d0FromIP",
),
],
),
Histogram(
name="hZ0",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"Original z0",
unit=r"$\mu m$",
bins=100,
scope=(-2000, 2000),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_micro_z0",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_micro_z0",
),
],
),
Histogram(
name="hMovedZ0",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"Moved z0",
unit=r"$\mu m$",
bins=100,
scope=(-2000, 2000),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_micro_z0FromIP",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_micro_z0FromIP",
),
],
),
Histogram(
name="hP",
title=r"$\mu mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"$p_{CMS}$",
unit=r"GeV",
bins=140,
scope=(0, 7),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_CM_p",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_CM_p",
),
],
),
Histogram(
name="hPt",
title=r"$\mu mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"p$_{T, CM}$",
unit=r"GeV",
bins=140,
scope=(0, 7),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_CM_pt",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_CM_pt",
),
],
),
Histogram(
name="hNCDC",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"CDC Hits",
unit=r"",
bins=80,
scope=(0, 80),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_nCDCHits",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_nCDCHits",
),
],
),
Histogram(
name="hNSVD",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"SVD Hits",
unit=r"",
bins=14,
scope=(0, 14),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_nSVDHits",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_nSVDHits",
),
],
),
Histogram(
name="hNPXD",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"PXD Hits",
unit=r"$\mu m$",
bins=4,
scope=(0, 4),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_nPXDHits",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_nSVDHits",
),
],
),
Histogram(
name="hD0Res",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable="d0FromIPRes"),
label=r"$\Delta d0 (From IP)$",
unit=r"$\mu$m",
bins=1000,
scope=(-200, 200),
),
hist_components=[
HistComponent(
label="",
additional_cut_str=None,
color="black",
variable=None,
),
],
),
Histogram(
name="hZ0Res",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable="z0FromIPRes"),
label=r"$\Delta z0 (From IP)$",
unit=r"$\mu$m",
bins=1000,
scope=(-500, 500),
),
hist_components=[
HistComponent(
label="",
additional_cut_str=None,
color="black",
variable=None,
),
],
),
Histogram(
name="hTime",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"$Track time$",
unit=r"ns",
bins=100,
scope=(-50, 50),
),
hist_components=[
HistComponent(
label="mu+nSVDHits>0",
additional_cut_str="mup_nSVDHits > 0.5",
color="red",
variable='mup_trackTime',
),
HistComponent(
label="mu-nSVDHits>0",
additional_cut_str="mum_nSVDHits > 0.5",
color="blue",
variable='mum_trackTime',
),
HistComponent(
label="mu+nSVDHits=0",
additional_cut_str="mup_nSVDHits <= 0.5",
color="orange",
variable='mup_trackTime',
),
HistComponent(
label="mu-nSVDHits=0",
additional_cut_str="mum_nSVDHits <= 0.5",
color="cyan",
variable='mum_trackTime',
),
],
),
Histogram(
name="hFlipRefit",
title=r"$\mu\mu$",
hist_variable=HistVariable(
df_label=makeROOTCompatible(variable=""),
label=r"Track is Flipped&Refitted",
unit=r"",
bins=2,
scope=(-0.5, 1.5),
),
hist_components=[
HistComponent(
label="mup",
additional_cut_str=None,
color="red",
variable="mup_isTrackFlippedAndRefitted",
),
HistComponent(
label="mum",
additional_cut_str=None,
color="blue",
variable="mum_isTrackFlippedAndRefitted",
),
],
),
]
