checkMemoryUsage.py

#!/usr/bin/env python3
 

 

 
import argparse
 
import basf2 as b2
import modularAnalysis as ma
import stdV0s
import variables.collections as vc
import variables.utils as vu
import vertex as vx
 
use_tflat = True
 
if use_tflat:
    import tflat.flavorTagger as ft
else:
    import flavorTagger as ft
 
b2.set_random_seed("aSeed")
 
# create path
main = b2.Path()
 
parser = argparse.ArgumentParser(
    description="Script to chech memory usage as part of light release validaiton"
)
 
parser.add_argument(
    "--release", type=str, help="The light release that is being tested."
)
args = parser.parse_args()
 
# load input ROOT file
ma.inputMdst(filename="", path=main)
 
# create muon particle list
ma.fillParticleList(decayString="mu+:all", cut="", path=main)
 
# reconstruct J/psi -> mu+ mu- decay
# keep only candidates with dM<0.11
ma.reconstructDecay(
    decayString="J/psi:mumu -> mu+:all mu-:all", cut="dM<0.11", path=main
)
 
# load KS from standard particle list
stdV0s.stdKshorts(fitter="TreeFit", path=main)
 
# reconstruct B0 -> J/psi KS decay
ma.reconstructDecay(
    decayString="B0:sig -> J/psi:mumu K_S0:merged",
    cut="Mbc > 5.2 and abs(deltaE)<0.15",
    path=main,
)
 
# match reconstructed with MC particles
ma.matchMCTruth(list_name="B0:sig", path=main)
 
# build the rest of the event associated to the B0
ma.buildRestOfEvent(target_list_name="B0:sig", path=main)
 
# set analysis global tag
b2.conditions.prepend_globaltag(ma.getAnalysisGlobaltag())
 
# Flavor Tagging
if use_tflat:
    # TFlaT
    ft.flavorTagger(
        particleLists=['B0:sig'],
        uniqueIdentifier='TFlaT_MC16rd_light_2601_hyperion',
        path=main)
    ft_var = 'qrTFLAT'
else:
    # GNN
    ft.flavorTagger(
        particleLists=["B0:sig"], weightFiles="B2nunubarBGx1", path=main, useGNN=True
    )
    ft_var = 'qrOutput(FBDT)'
 
# rank by highest r- factor
ma.rankByHighest(
    particleList="B0:sig",
    variable=f"abs({ft_var})",
    numBest=0,
    outputVariable="Dilution_rank",
    path=main,
)
 
# fit B vertex on the signal side
vx.treeFit(
    list_name="B0:sig",
    conf_level=-1,
    massConstraint=["K_S0"],
    ipConstraint=True,
    updateAllDaughters=True,
    path=main,
)
 
# fit B vertex on the tag side
vx.TagV(list_name="B0:sig", MCassociation="breco", path=main)
 
# select variables that will be stored to ntuple
fs_vars = vc.pid + vc.track + vc.track_hits + vc.mc_truth
jpsiandk0s_vars = vc.mc_truth
vertex_vars = vc.vertex + vc.mc_vertex + vc.kinematics + vc.mc_kinematics
bvars = (
    vc.reco_stats
    + vc.deltae_mbc
    + vc.mc_truth
    + vc.roe_multiplicities
    + vc.tag_vertex
    + vc.mc_tag_vertex
    + vertex_vars
)
 
if use_tflat:
    bvars += ['qrTFLAT']
else:
    # Attention: the collection of flavor tagging variables is defined in the flavorTagger
    bvars += ft.flavor_tagging
 
# Create aliases to save information for different particles
bvars = (
    bvars
    + vu.create_aliases_for_selected(
        list_of_variables=fs_vars,
        decay_string="B0 -> [J/psi -> ^mu+ ^mu-] [K_S0 -> ^pi+ ^pi-]",
    )
    + vu.create_aliases_for_selected(
        list_of_variables=jpsiandk0s_vars,
        decay_string="B0 -> [^J/psi -> mu+ mu-] [^K_S0 -> pi+ pi-]",
    )
    + vu.create_aliases_for_selected(
        list_of_variables=vertex_vars,
        decay_string="B0 -> [^J/psi -> ^mu+ ^mu-] [^K_S0 -> ^pi+ ^pi-]",
    )
)
 
# saving variables to ntuple
output_file = f"output/{args.release}.root"
ma.variablesToNtuple(
    decayString="B0:sig",
    variables=bvars,
    filename=output_file,
    treename="tree",
    path=main,
)
 
# summary of created Lists
ma.summaryOfLists(particleLists=["K_S0:merged", "J/psi:mumu", "B0:sig"], path=main)
 
# profile execution time and memory usage
main.add_module(
    "Profile",
    outputFileName=f"output/MemoryUsage_{args.release}.root",
    rssOutputFileName=f"output/RSSMemoryUsage_{args.release}.root",
)
 
# process the events
b2.process(main)
 
# print out the summary
print(b2.statistics)