ShipReco.py

Go to the documentation of this file.

#!/usr/bin/env python
# SPDX-License-Identifier: LGPL-3.0-or-later
# SPDX-FileCopyrightText: Copyright CERN for the benefit of the SHiP Collaboration
 
from argparse import ArgumentParser
 
withHists = True
pidProton = False  # if true, take truth, if False fake with pion mass
 
import resource
 
 
def mem_monitor() -> None:
    # Getting virtual memory size
    pid = os.getpid()
    with open(os.path.join("/proc", str(pid), "status")) as f:
        lines = f.readlines()
    _vmsize = [line for line in lines if line.startswith("VmSize")][0]
    vmsize = int(_vmsize.split()[1])
    # Getting physical memory size
    pmsize = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    print(f"memory: virtual = {vmsize / 1.0e3:5.2F} MB  physical = {pmsize / 1.0e3:5.2F} MB")
 
 
import os
 
import global_variables
import ROOT
import rootUtils as ut
import shipRoot_conf
 
shipRoot_conf.configure()
 
parser = ArgumentParser()
 
parser.add_argument("-f", "--inputFile", dest="inputFile", help="Input file", required=True)
parser.add_argument(
    "-n", "--nEvents", dest="nEvents", help="Number of events to reconstruct", required=False, default=999999, type=int
)
parser.add_argument("-g", "--geoFile", dest="geoFile", help="ROOT geofile", required=True)
parser.add_argument(
    "--noVertexing", dest="noVertexing", help="switch off vertexing", required=False, action="store_true"
)
parser.add_argument(
    "--noStrawSmearing",
    dest="withNoStrawSmearing",
    help="no smearing of distance to wire, default on",
    required=False,
    action="store_true",
)
parser.add_argument(
    "--withT0", dest="withT0", help="simulate arbitrary T0 and correct for it", required=False, action="store_true"
)
parser.add_argument(
    "-i",
    "--firstEvent",
    dest="firstEvent",
    help="First event of input file to use",
    required=False,
    default=0,
    type=int,
)
parser.add_argument(
    "--realPR",
    dest="realPR",
    help="Option for pattern recognition without MC truth. \n\
           FH                        : Hough transform.\n\
           AR                        : Artificial retina.\n\
      TemplateMatching               : Tracks are searched for based on the template: track seed + hits within a window around the seed.",
    required=False,
    choices=["FH", "AR", "TemplateMatching"],
    default="",
)
parser.add_argument("-dy", dest="dy", help="Max height of tank", required=False, default=None, type=int)
parser.add_argument("--Debug", dest="Debug", help="Switch on debugging", required=False, action="store_true")
 
options = parser.parse_args()
vertexing = not options.noVertexing
 
 
# need to figure out which geometry was used, only needed if no geo file
if not options.dy:
    # try to extract from input file name
    tmp = options.inputFile.split(".")
    try:
        dy = float(tmp[1] + "." + tmp[2])
    except (ValueError, IndexError):
        dy = None
 
print(
    "configured to process ",
    options.nEvents,
    " events from ",
    options.inputFile,
    " starting with event ",
    options.firstEvent,
    " with option Yheight = ",
    dy,
    " with vertexing",
    vertexing,
    " and real pattern reco ",
    options.realPR,
)
# Determine output filename (will contain only digi/reco branches)
if not options.inputFile.find("_rec.root") < 0:
    outFile = options.inputFile
    options.inputFile = outFile.replace("_rec.root", ".root")
else:
    outFile = options.inputFile.replace(".root", "_rec.root")
    # outfile should be in local directory
    tmp = outFile.split("/")
    outFile = tmp[len(tmp) - 1]
 
if not options.geoFile:
    tmp = options.inputFile.replace("ship.", "geofile_full.")
    options.geoFile = tmp.replace("_rec", "")
 
fgeo = ROOT.TFile.Open(options.geoFile)
geoMat = ROOT.genfit.TGeoMaterialInterface()  # if only called in ShipDigiReco -> crash, reason unknown
 
from ShipGeoConfig import load_from_root_file
 
# load Shipgeo dictionary
ShipGeo = load_from_root_file(fgeo, "ShipGeo")
 
h = {}
if withHists:
    ut.bookHist(h, "distu", "distance to wire", 100, 0.0, 5.0)
    ut.bookHist(h, "distv", "distance to wire", 100, 0.0, 5.0)
    ut.bookHist(h, "disty", "distance to wire", 100, 0.0, 5.0)
    ut.bookHist(h, "nmeas", "nr measurements", 100, 0.0, 50.0)
    ut.bookHist(h, "chi2", "Chi2/DOF", 100, 0.0, 20.0)
 
import shipDet_conf
 
run = ROOT.FairRunSim()
run.SetName("TGeant4")  # Transport engine
run.SetSink(ROOT.FairRootFileSink(ROOT.TMemFile("output", "recreate")))  # Dummy output file
run.SetUserConfig("g4Config_basic.C")  # geant4 transport not used, only needed for creating VMC field
rtdb = run.GetRuntimeDb()
# -----Create geometry----------------------------------------------
modules = shipDet_conf.configure(run, ShipGeo)
# run.Init()
fgeo["FAIRGeom"]
import geomGeant4
 
if hasattr(ShipGeo.Bfield, "fieldMap"):
    fieldMaker = geomGeant4.addVMCFields(ShipGeo, "", True, withVirtualMC=False)
 
# make global variables
global_variables.debug = options.Debug
global_variables.fieldMaker = fieldMaker
global_variables.pidProton = pidProton
global_variables.withT0 = options.withT0
global_variables.realPR = options.realPR
global_variables.vertexing = vertexing
global_variables.ShipGeo = ShipGeo
global_variables.modules = modules
global_variables.withNoStrawSmearing = options.withNoStrawSmearing
global_variables.h = h
global_variables.iEvent = 0
 
# import reco tasks
import shipDigiReco
 
SHiP = shipDigiReco.ShipDigiReco(options.inputFile, outFile, fgeo)
options.nEvents = min(SHiP.sTree.GetEntries(), options.nEvents)
# main loop
for global_variables.iEvent in range(options.firstEvent, options.nEvents):
    if global_variables.iEvent % 1000 == 0 or global_variables.debug:
        print("event ", global_variables.iEvent)
    rc = SHiP.sTree.GetEvent(global_variables.iEvent)
    SHiP.digitize()
    SHiP.reconstruct()
    SHiP.recoTree.Fill()
# memory monitoring
# mem_monitor()
# end loop over events
SHiP.finish()