dd4hepplugins/examples/validate_drich.py

0001 #!/usr/bin/env python3
0002 """Run GPU and CPU with simplified dRICH geometry (no non-optical PDU components)."""
0003 import math, os, sys, numpy as np
0004
0005 _SCRIPT_DIR = "/tmp/eic-opticks/dd4hepplugins/examples"
0006 _GEOM_DIR = os.path.join(_SCRIPT_DIR, "geometry")
0007 sys.path.insert(0, _SCRIPT_DIR)
0008
0009 NUM_EVENTS = int(os.environ.get("NUM_EVENTS", "3"))
0010 MULTIPLICITY = int(os.environ.get("MULTIPLICITY", "20"))
0011
0012
0013 def setup_kernel(mode):
0014     import cppyy, DDG4
0015     from g4units import GeV
0016
0017     cppyy.include("G4OpticalParameters.hh")
0018     from cppyy.gbl import G4OpticalParameters
0019
0020     compact = os.path.join(_GEOM_DIR, "epic_drich_simple.xml")
0021     if "DRICH_SIMPLE_XML" not in os.environ:
0022         os.environ["DRICH_SIMPLE_XML"] = os.path.join(_GEOM_DIR, "drich_1sector_simple.xml")
0023
0024     if mode == "gpu":
0025         os.environ["OPTICKS_EVENT_MODE"] = "Minimal"
0026         os.environ["OPTICKS_INTEGRATION_MODE"] = "1"
0027         # Ensure DDG4 can find the plugin
0028         import ctypes
0029         ctypes.CDLL("libddeicopticks.so", ctypes.RTLD_GLOBAL)
0030     else:
0031         os.environ.pop("OPTICKS_EVENT_MODE", None)
0032         os.environ.pop("OPTICKS_INTEGRATION_MODE", None)
0033
0034     kernel = DDG4.Kernel()
0035     kernel.loadGeometry(str("file:" + compact))
0036     geant4 = DDG4.Geant4(kernel)
0037     geant4.setupUI(typ="tcsh", vis=False, ui=False)
0038
0039     eta = 2.0
0040     theta = 2.0 * math.atan(math.exp(-eta))
0041     phi = math.radians(-167.0)
0042     geant4.setupGun("Gun", particle="pi+", energy=10*GeV,
0043                     position=(0, 0, 0), isotrop=False,
0044                     direction=(math.sin(theta)*math.cos(phi),
0045                                math.sin(theta)*math.sin(phi),
0046                                math.cos(theta)),
0047                     multiplicity=MULTIPLICITY)
0048
0049     geant4.setupPhysics("QGSP_BERT")
0050     ph = DDG4.PhysicsList(kernel, "Geant4PhysicsList/OpticalPhys")
0051     ph.addPhysicsConstructor(str("G4OpticalPhysics"))
0052     kernel.physicsList().adopt(ph)
0053
0054     # Detector setup needed for both modes (GPU injects hits into DD4hep)
0055     seq, act = geant4.setupDetector("DRICH", "Geant4OpticalTrackerAction")
0056     filt = DDG4.Filter(kernel, "ParticleSelectFilter/OpticalPhotonSelector")
0057     filt.particle = "opticalphoton"
0058     seq.adopt(filt)
0059
0060     if mode == "gpu":
0061         step_action = DDG4.SteppingAction(kernel, "OpticsSteppingAction/OpticsStep1")
0062         kernel.steppingAction().adopt(step_action)
0063         run_action = DDG4.RunAction(kernel, "OpticsRun/OpticsRun1")
0064         kernel.runAction().adopt(run_action)
0065         # OpticsEvent must be registered BEFORE ROOT output so hits are
0066         # injected into the collection before serialization.
0067         evt_action = DDG4.EventAction(kernel, "OpticsEvent/OpticsEvt1")
0068         kernel.eventAction().adopt(evt_action)
0069
0070     if mode == "gpu":
0071         geant4.setupROOTOutput("RootOutput", "/tmp/drich_simple_gpu")
0072     else:
0073         geant4.setupROOTOutput("RootOutput", "/tmp/drich_simple_cpu")
0074
0075     kernel.NumEvents = NUM_EVENTS
0076     kernel.configure()
0077
0078     # In GPU mode, tell G4Cerenkov/G4Scintillation to NOT push optical
0079     # photon secondaries onto the Geant4 stack.  The genstep is still
0080     # computed (numPhotons, BetaInverse, etc.) and collected by
0081     # OpticsSteppingAction, but no CPU photon tracks are created.
0082     if mode == "gpu":
0083         G4OpticalParameters.Instance().SetCerenkovStackPhotons(False)
0084         G4OpticalParameters.Instance().SetScintStackPhotons(False)
0085
0086     kernel.initialize()
0087     G4OpticalParameters.Instance().SetBoundaryInvokeSD(True)
0088     return kernel
0089
0090
0091 def run_gpu():
0092     kernel = setup_kernel("gpu")
0093     kernel.run()
0094     kernel.terminate()
0095
0096     # Read hits from npy files
0097     base = os.path.join(os.environ.get("TMP", "/tmp/opticks_out"),
0098                         "GEOM", os.environ.get("GEOM", "simple"),
0099                         "python3.13", "ALL0_no_opticks_event_name")
0100     total = 0
0101     if os.path.isdir(base):
0102         for d in sorted(os.listdir(base)):
0103             hp = os.path.join(base, d, "hit.npy")
0104             if os.path.exists(hp):
0105                 n = len(np.load(hp))
0106                 total += n
0107                 print(f"  {d}: {n} hits")
0108     print(f"GPU total: {total} hits")
0109     return total
0110
0111
0112 def run_cpu():
0113     kernel = setup_kernel("cpu")
0114     kernel.run()
0115     kernel.terminate()
0116
0117     # Extract hits from ROOT
0118     import ROOT
0119     ROOT.gSystem.Load("libDDG4Plugins")
0120     ROOT.gSystem.Load("libDDG4")
0121     f = ROOT.TFile.Open("/tmp/drich_simple_cpu.root")
0122     tree = f.Get("EVENT")
0123     total = 0
0124     for i in range(tree.GetEntries()):
0125         tree.GetEntry(i)
0126         nhits = tree.DRICHHits.size()
0127         total += nhits
0128         print(f"  event {i}: {nhits} hits")
0129     f.Close()
0130     print(f"CPU total: {total} hits")
0131     return total
0132
0133
0134 if __name__ == "__main__":
0135     mode = sys.argv[1] if len(sys.argv) > 1 else "gpu"
0136     if mode == "gpu":
0137         run_gpu()
0138     elif mode == "cpu":
0139         run_cpu()
0140     else:
0141         print(f"Usage: {sys.argv[0]} [gpu|cpu]")