faser_digi.py

#!/usr/bin/env python
#
# Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
# Run with:
# ./faser_digi.py filepath runtype
# 
# filepath - fully qualified path, including url if needed, to the input HITS file
#   example: "root://eospublic.cern.ch//eos/experiment/faser/sim/GeniePilot/HITS/1/faser.150fbInv.1.001.HITS.pool.root"
# 
# runtype - MANDATORY flag to specify the data type (TI12OldMC or TI12MC or TestBeamMC).
#   Not extracted (yet) from file path for MC data 
#
import sys
import argparse

parser = argparse.ArgumentParser(description="Run FASER reconstruction")

parser.add_argument("file_path",
                    help="Fully qualified path of the raw input file")
parser.add_argument("run_type", nargs="?", default="",
                    help="Specify run type (if it can't be parsed from path)")
parser.add_argument("-t", "--tag", default="",
                    help="Specify tag (to append to output filename)")
parser.add_argument("--highCaloGain", action='store_true',
                    help="Use high gain settings for calo PMTs")
parser.add_argument("-n", "--nevents", type=int, default=-1,
                    help="Specify number of events to process (default: all)")
parser.add_argument("-v", "--verbose", action='store_true', 
                    help="Turn on DEBUG output")

args = parser.parse_args()

from pathlib import Path

filepath=Path(args.file_path)

# runtype has been provided
if len(args.run_type) > 0:
    runtype=args.run_type

# Extract runtype from path
# Should be directory above run
# i.e.: TestBeamData/Run-004150/Faser-Physics-004150-00000.raw"
else:
    if True or len(filepath.parts) < 3:
        print("Can't determine run type from path - specify on command line ")
        sys.exit(-1)

#    runtype = filepath.parts[-3]

print(f"Starting digitization of {filepath.name} with type {runtype}")
if args.nevents > 0:
    print(f"Reconstructing {args.nevents} events by command-line option")

# Start digitization

from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator
from AthenaConfiguration.ComponentFactory import CompFactory
from AthenaCommon.Constants import VERBOSE, INFO

from AthenaCommon.Configurable import Configurable
from CalypsoConfiguration.AllConfigFlags import ConfigFlags

Configurable.configurableRun3Behavior = True
    
# Flags for this job
ConfigFlags.Input.isMC = True                    # Needed to bypass autoconfig
ConfigFlags.IOVDb.DatabaseInstance = "OFLP200"   # Use MC conditions for now

ConfigFlags.Input.ProjectName = "mc20"
ConfigFlags.GeoModel.Align.Dynamic    = False
ConfigFlags.Beam.NumberOfCollisions = 0.
ConfigFlags.Digitization.TruthOutput = True

# TI12 old geometry
if runtype == "TI12OldMC":
    ConfigFlags.GeoModel.FaserVersion = "FASER-01" 
    ConfigFlags.IOVDb.GlobalTag = "OFLCOND-FASER-01"

# Testbeam setup 
elif runtype == "TestBeamMC" :
    ConfigFlags.GeoModel.FaserVersion = "FASER-TB00" 
    ConfigFlags.IOVDb.GlobalTag = "OFLCOND-FASER-TB00"

# New TI12 geometry (ugh)
elif runtype == "TI12MC":
    ConfigFlags.GeoModel.FaserVersion = "FASERNU-03" 
    ConfigFlags.IOVDb.GlobalTag = "OFLCOND-FASER-02"

else:
    print("Invalid run type found:", runtype)
    print("Specify correct type or update list")
    sys.exit(-1)


# Must use original input string here, as pathlib mangles double // in path names
ConfigFlags.Input.Files = [ args.file_path ]

filestem = filepath.stem
if len(args.tag) > 0:
    filestem += f"-{args.tag}"

# ConfigFlags.addFlag("Output.xAODFileName", f"{filestem}-xAOD.root")
ConfigFlags.Output.RDOFileName = f"{filestem}-RDO.root"

#
# Play around with this?
# ConfigFlags.Concurrency.NumThreads = 2
# ConfigFlags.Concurrency.NumConcurrentEvents = 2
ConfigFlags.lock()

#
# Configure components
from CalypsoConfiguration.MainServicesConfig import MainServicesCfg
from AthenaPoolCnvSvc.PoolReadConfig import PoolReadCfg
from AthenaPoolCnvSvc.PoolWriteConfig import PoolWriteCfg
    
acc = MainServicesCfg(ConfigFlags)
acc.merge(PoolReadCfg(ConfigFlags))
acc.merge(PoolWriteCfg(ConfigFlags))

#
# Needed, or move to MainServicesCfg?
from FaserGeoModel.FaserGeoModelConfig import FaserGeometryCfg
acc.merge(FaserGeometryCfg(ConfigFlags))

# Set up algorithms
from FaserSCT_Digitization.FaserSCT_DigitizationConfigNew import FaserSCT_DigitizationCfg
acc.merge(FaserSCT_DigitizationCfg(ConfigFlags))

from CaloDigiAlgs.CaloDigiAlgsConfig import CaloWaveformDigitizationCfg
if args.highCaloGain:
    calo_norm = 25.
else:
    calo_norm = 5.
acc.merge(CaloWaveformDigitizationCfg(ConfigFlags, CB_norm=calo_norm))

from ScintDigiAlgs.ScintDigiAlgsConfig import ScintWaveformDigitizationCfg
acc.merge(ScintWaveformDigitizationCfg(ConfigFlags))

# from WaveRecAlgs.WaveRecAlgsConfig import WaveformReconstructionCfg    
# acc.merge(WaveformReconstructionCfg(ConfigFlags))

# # Not ready for primetime
# # from CaloRecAlgs.CaloRecAlgsConfig import CalorimeterReconstructionCfg
# # acc.merge(CalorimeterReconstructionCfg(ConfigFlags))

# # Tracker clusters
# from TrackerPrepRawDataFormation.TrackerPrepRawDataFormationConfig import FaserSCT_ClusterizationCfg
# acc.merge(FaserSCT_ClusterizationCfg(ConfigFlags))

# # SpacePoints
# from TrackerSpacePointFormation.TrackerSpacePointFormationConfig import TrackerSpacePointFinderCfg
# acc.merge(TrackerSpacePointFinderCfg(ConfigFlags))

# # Try Dave's fitter
# from TrackerClusterFit.TrackerClusterFitConfig import ClusterFitAlgCfg
# acc.merge(ClusterFitAlgCfg(ConfigFlags))

#
# Configure output
# from OutputStreamAthenaPool.OutputStreamConfig import OutputStreamCfg
# itemList = [ "xAOD::EventInfo#*"
#              , "xAOD::EventAuxInfo#*"
#              , "xAOD::FaserTriggerData#*"
#              , "xAOD::FaserTriggerDataAux#*"
#              , "FaserSCT_RDO_Container#*"
#             #  , "Tracker::FaserSCT_ClusterContainer#*"
#             #  , "FaserSCT_SpacePointContainer#*"
#             #  , "FaserSCT_SpacePointOverlapCollection#*"
#             #  , "TrackCollection#*"
# ]
# acc.merge(OutputStreamCfg(ConfigFlags, "xAOD", itemList))

# Waveform reconstruction output
# from WaveRecAlgs.WaveRecAlgsConfig import WaveformReconstructionOutputCfg    
# acc.merge(WaveformReconstructionOutputCfg(ConfigFlags))

# Calorimeter reconstruction output
# from CaloRecAlgs.CaloRecAlgsConfig import CalorimeterReconstructionOutputCfg
# acc.merge(CalorimeterReconstructionOutputCfg(ConfigFlags))

# Check what we have
# print( "Writing out xAOD objects:" )
# print( acc.getEventAlgo("OutputStreamxAOD").ItemList )

# Hack to avoid problem with our use of MC databases when isMC = False
# replicaSvc = acc.getService("DBReplicaSvc")
# replicaSvc.COOLSQLiteVetoPattern = ""
# replicaSvc.UseCOOLSQLite = True
# replicaSvc.UseCOOLFrontier = False
# replicaSvc.UseGeomSQLite = True

# Configure verbosity    
if args.verbose:
    acc.foreach_component("*").OutputLevel = VERBOSE
    ConfigFlags.dump()

else:
    acc.foreach_component("*").OutputLevel = INFO

acc.foreach_component("*ClassID*").OutputLevel = INFO

acc.getService("MessageSvc").Format = "% F%40W%S%7W%R%T %0W%M"

# Execute and finish
sys.exit(int(acc.run(maxEvents=args.nevents).isFailure()))