Scripts/Python/geant4_parallel.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004 from multiprocessing import Pool
0005 from functools import partial
0006
0007 # This script runs a Geant4 simulation in parallel by passing chunks of events to subprocesses.
0008 # This is a workaround to achieve parallel processing even though Geant4 is not thread-save
0009 # and thus the internal parallelism of the ACTS examples framework cannot be used.
0010 #
0011 # Note that:
0012 # ==========
0013 #
0014 # * This should give equivalent results to a sequential run if the RNG is initialized with
0015 #   the same seed in all runs
0016 #
0017 # * So far this works only for csv outputs, since they write the results in one file per event
0018 #   (the naming of the csv-files should be equivalent to a sequential run)
0019 #
0020 # * In principle it is not difficult to extend this for ROOT files as well. One would need to
0021 #   write the root-files into separate directory per chunk, and then use ROOT's hadd to combine
0022 #   the output files.
0023 #
0024
0025
0026 def runGeant4EventRange(detector, trackingGeometry, beginEvent, endEvent, outputDir):
0027     import acts
0028     import acts.examples
0029     from acts.examples.simulation import addParticleGun, addGeant4, EtaConfig
0030
0031     u = acts.UnitConstants
0032
0033     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0034     rnd = acts.examples.RandomNumbers(seed=42)
0035
0036     s = acts.examples.Sequencer(
0037         events=endEvent - beginEvent, skip=beginEvent, numThreads=1
0038     )
0039
0040     outputDir = Path(outputDir)
0041     addParticleGun(
0042         s,
0043         EtaConfig(-2.0, 2.0),
0044         rnd=rnd,
0045         outputDirCsv=outputDir / "csv",
0046         outputDirRoot=None,
0047     )
0048     addGeant4(
0049         s,
0050         detector,
0051         trackingGeometry,
0052         field,
0053         outputDirCsv=outputDir / "csv",
0054         outputDirRoot=None,
0055         rnd=rnd,
0056     )
0057
0058     s.run()
0059
0060
0061 if "__main__" == __name__:
0062     from acts.examples.odd import getOpenDataDetector
0063
0064     detector = getOpenDataDetector()
0065     trackingGeometry = detector.trackingGeometry()
0066     decorators = detector.contextDecorators()
0067
0068     n_events = 100
0069     n_jobs = 8
0070
0071     chunksize = n_events // (n_jobs - 1)
0072     begins = range(0, n_events, chunksize)
0073     ends = [min(b + chunksize, n_events) for b in begins]
0074
0075     outputDir = Path.cwd()
0076
0077     with Pool(n_jobs) as p:
0078         p.starmap(partial(runGeant4EventRange, outputDir=outputDir), zip(begins, ends))