Create default simulation scripts

Control/CalypsoExample/Simulation/scripts/faser_simulate.py

+#!/usr/bin/env python
+"""
+Produce simulated hits from input 4-vectors
+Derived from G4FaserAlgConfigNew
+
+Usage:
+faser_simulate.py filepath outfile
+
+filepath - full path, including url if needed, to the input 4-vector file
+outfile - output filename, parameters will be inferred from this name
+
+Copyright (C) 2002-2021 CERN for the benefit of the ATLAS and FASER collaborations
+"""
+
+if __name__ == '__main__':
+
+    import sys
+    import time
+    a = time.time()
+#
+# Parse command-line options
+#
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run FASER simulation")
+
+    parser.add_argument("file_path",
+                        help="Fully qualified path of the raw input file")
+
+    parser.add_argument("output",
+                        help="Output file name")
+
+    #parser.add_argument("--run", type=int, default=123456,
+    #                    help="Specify run number to use in simulated data")
+
+    parser.add_argument("--geom", default="TI12MC",
+                        help="Specify geomtery to simulation (default: TI12MC, alt: TestBeamMC)")
+
+
+    parser.add_argument("--xangle", type=float, default=0.0,
+                        help="Specify H crossing angle (in Radians)")
+    parser.add_argument("--yangle", type=float, default=0.0,
+                        help="Specify V crossing angle (in Radians)")
+    parser.add_argument("--xshift", type=float, default=0.0,
+                        help="Specify H shift of events wrt FASER (in mm)")
+    parser.add_argument("--yshift", type=float, default=0.0,
+                        help="Specify V shift of events wrt FASER (in mm)")
+
+    parser.add_argument("-t", "--tag", default="",
+                        help="Specify sim tag (to append to output filename)")
+    parser.add_argument("-s", "--skip", type=int, default=0,
+                        help="Specify number of events to skip (default: none)")
+    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)
+
+#
+# Parse input file
+#
+    print(f"Starting digitization of {filepath.name}")
+
+    filestem = filepath.stem
+    if len(args.tag) > 0:
+        if args.tag in filestem:
+            print(f"Not adding tag {args.tag} to {filestem}")
+        else:
+            filestem += f"-{args.tag}"
+
+    if args.output:
+        # Just directly specify the output filename
+        outfile = args.output
+
+        spl = outfile.split('-')
+        if len(spl) < 4:
+            print(f"Can't get run number from {outfile}!")
+            sys.exit(1)
+
+        runnum = int(spl[2])
+        segnum = int(spl[3])
+
+    else:
+        outfile = f"{filestem}-HITS.root"
+        print(f"Output file name not specified")
+        sys.exit(1)
+
+    print(f"Outfile: {outfile}")
+
+#
+# Figure out events to run
+#
+    if args.skip > 0:
+        print(f"Skipping {args.skip} events by command-line option")
+
+    if args.nevents > 0:
+        print(f"Reconstructing {args.nevents} events by command-line option")
+#
+# Set up logging and config behaviour
+#
+    from AthenaCommon.Logging import log
+    from AthenaCommon.Constants import DEBUG, VERBOSE
+    from AthenaCommon.Configurable import Configurable
+    log.setLevel(DEBUG)
+    Configurable.configurableRun3Behavior = 1
+#
+# Import and set config flags
+#
+    from CalypsoConfiguration.AllConfigFlags import ConfigFlags
+    from AthenaConfiguration.Enums import ProductionStep
+    ConfigFlags.Common.ProductionStep = ProductionStep.Simulation
+#
+# All these must be specified to avoid auto-configuration
+#
+    ConfigFlags.Input.RunNumber = [ runnum ] 
+    ConfigFlags.Input.OverrideRunNumber = True
+    ConfigFlags.Input.LumiBlockNumber = [ (segnum+1) ]
+    ConfigFlags.Input.isMC = True
+#
+# Input file name
+# 
+    # Path mangles // in url, so use direct file_path here
+    ConfigFlags.Input.Files = [ args.file_path ]
+#
+# Skip events
+#
+    ConfigFlags.Exec.SkipEvents = args.skip
+    ConfigFlags.Exec.MaxEvents = args.nevents
+#
+# Output file name
+# 
+    ConfigFlags.Output.HITSFileName = outfile
+#
+# Sim ConfigFlags
+#
+    ConfigFlags.Sim.Layout = "FASER"
+    ConfigFlags.Sim.PhysicsList = "FTFP_BERT"
+    ConfigFlags.Sim.ReleaseGeoModel = False
+    ConfigFlags.Sim.IncludeParentsInG4Event = True # Controls whether BeamTruthEvent is written to output HITS file
+
+#
+# Figure out configuration
+#
+    if args.geom == "TI12MC":
+        # 2022 TI12 geometry
+        ConfigFlags.GeoModel.FaserVersion = "FASERNU-03"  # Geometry set-up
+        ConfigFlags.IOVDb.GlobalTag = "OFLCOND-FASER-03"  # Conditions set-up
+        # TI12 detectors
+        detectors = ['Veto', 'VetoNu', 'Preshower', 'FaserSCT', 'Ecal', 
+                     'Trigger', 'Dipole', 'Emulsion', 'Trench']
+
+    elif args.geom == "TestBeamMC":
+        # Define 2021 test beam geometry
+        ConfigFlags.GeoModel.FaserVersion = "FASER-TB00"   # Geometry set-up
+        ConfigFlags.IOVDb.GlobalTag = "OFLCOND-FASER-TB00" # Conditions set-up
+        # Testbeam detectors (trigger layers are actually veto counters)
+        detectors = ['Veto', 'Preshower', 'FaserSCT', 'Ecal']
+
+    else:
+        print(f"Unknown geometry {args.geom}!")
+        sys.exit(1)
+
+
+    #
+    # Units are radians and mm
+    # Probably should only be allowed in TI12, but leave it here for now
+    ConfigFlags.addFlag("Sim.Beam.xangle", args.xangle)  # Potential beam crossing angles
+    ConfigFlags.addFlag("Sim.Beam.yangle", args.yangle)    
+    ConfigFlags.addFlag("Sim.Beam.xshift", args.xshift)  # Potential beam shift
+    ConfigFlags.addFlag("Sim.Beam.yshift", args.yshift)    
+
+    ConfigFlags.addFlag("Input.InitialTimeStamp", 0)     # To avoid autoconfig 
+    ConfigFlags.GeoModel.Align.Dynamic = False
+
+    # import sys
+    # ConfigFlags.fillFromArgs(sys.argv[1:])
+
+    doShiftLOS = (ConfigFlags.Sim.Beam.xangle or ConfigFlags.Sim.Beam.yangle or
+                  ConfigFlags.Sim.Beam.xshift or ConfigFlags.Sim.Beam.yshift)
+#
+# Setup detector flags
+#
+    from CalypsoConfiguration.DetectorConfigFlags import setupDetectorsFromList
+    setupDetectorsFromList(ConfigFlags, detectors, toggle_geometry=True)
+#
+# Finalize flags
+#
+    ConfigFlags.lock()
+#
+# Initialize a new component accumulator
+#
+    from CalypsoConfiguration.MainServicesConfig import MainServicesCfg
+    cfg = MainServicesCfg(ConfigFlags)
+#
+# Check whether a known input file was specified
+#
+    if ConfigFlags.Input.Files[0].endswith(".events") or ConfigFlags.Input.Files[0].endswith(".hepmc"):
+
+        from HEPMCReader.HepMCReaderConfig import HepMCReaderCfg
+
+        if doShiftLOS:
+            cfg.merge(HepMCReaderCfg(ConfigFlags, McEventKey = "BeamTruthEvent_ATLASCoord"))
+        else:
+            cfg.merge(HepMCReaderCfg(ConfigFlags))
+
+        from McEventSelector.McEventSelectorConfig import McEventSelectorCfg
+        cfg.merge(McEventSelectorCfg(ConfigFlags))
+#
+# Else, set up to read it as a pool.root file
+#
+    else:
+        from AthenaPoolCnvSvc.PoolReadConfig import PoolReadCfg
+        cfg.merge(PoolReadCfg(ConfigFlags))
+
+        if doShiftLOS:
+            from SGComps.AddressRemappingConfig import InputOverwriteCfg
+            # Rename old truth collection to add ATLAS coord to can still use BeamTruthEvent for the one in FASER Coords
+            cfg.merge(InputOverwriteCfg("McEventCollection", "BeamTruthEvent", "McEventCollection", "BeamTruthEvent_ATLASCoord"))
+
+#
+# Output file
+#
+    from AthenaPoolCnvSvc.PoolWriteConfig import PoolWriteCfg
+    cfg.merge(PoolWriteCfg(ConfigFlags))
+
+#
+# Shift LOS
+#
+    if doShiftLOS:
+
+        import McParticleEvent.Pythonizations
+        from GeneratorUtils.ShiftLOSConfig import ShiftLOSCfg
+        cfg.merge(ShiftLOSCfg(ConfigFlags, 
+                              xcross = ConfigFlags.Sim.Beam.xangle, 
+                              ycross = ConfigFlags.Sim.Beam.yangle,
+                              xshift = ConfigFlags.Sim.Beam.xshift,
+                              yshift = ConfigFlags.Sim.Beam.yshift))
+
+
+#
+# Add the G4FaserAlg
+#
+    from G4FaserAlg.G4FaserAlgConfigNew import G4FaserAlgCfg
+    cfg.merge(G4FaserAlgCfg(ConfigFlags))
+#
+# Dump config
+#
+    from AthenaConfiguration.ComponentFactory import CompFactory
+    cfg.addEventAlgo(CompFactory.JobOptsDumperAlg(FileName="G4FaserTestConfig.txt"))
+    cfg.getService("StoreGateSvc").Dump = True
+    cfg.getService("ConditionStore").Dump = True
+    cfg.printConfig(withDetails=True, summariseProps = False)  # gags on ParticleGun if summariseProps = True?
+
+    ConfigFlags.dump()
+    #f = open("test.pkl","wb")
+    #cfg.store(f)
+    #f.close()
+#
+# Execute and finish
+#
+
+    # This fails with ShiftLOSCfg...
+    #if args.verbose:
+    #    cfg.foreach_component("*").OutputLevel = "DEBUG" 
+    #else:
+    #    cfg.foreach_component("*").OutputLevel = "INFO"  
+
+    sc = cfg.run(maxEvents=args.nevents)
+
+    b = time.time()
+    log.info("Finish execution in " + str(b-a) + " seconds")
+#
+# Success should be 0
+#
+    sys.exit(int(sc.isFailure()))
+