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(beginEvent, endEvent, outputDir):
0027     import acts
0028     import acts.examples
0029     from acts.examples.simulation import addParticleGun, addGeant4, EtaConfig
0030     from acts.examples.odd import getOpenDataDetector
0031
0032     u = acts.UnitConstants
0033
0034     detector, trackingGeometry, decorators = getOpenDataDetector()
0035
0036     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0037     rnd = acts.examples.RandomNumbers(seed=42)
0038
0039     s = acts.examples.Sequencer(
0040         events=endEvent - beginEvent, skip=beginEvent, numThreads=1
0041     )
0042
0043     outputDir = Path(outputDir)
0044     addParticleGun(
0045         s,
0046         EtaConfig(-2.0, 2.0),
0047         rnd=rnd,
0048         outputDirCsv=outputDir / "csv",
0049         outputDirRoot=None,
0050     )
0051     addGeant4(
0052         s,
0053         detector,
0054         trackingGeometry,
0055         field,
0056         outputDirCsv=outputDir / "csv",
0057         outputDirRoot=None,
0058         rnd=rnd,
0059     )
0060
0061     s.run()
0062     del s
0063
0064
0065 if "__main__" == __name__:
0066     n_events = 100
0067     n_jobs = 8
0068
0069     chunksize = n_events // (n_jobs - 1)
0070     begins = range(0, n_events, chunksize)
0071     ends = [min(b + chunksize, n_events) for b in begins]
0072
0073     outputDir = Path.cwd()
0074
0075     with Pool(n_jobs) as p:
0076         p.starmap(partial(runGeant4EventRange, outputDir=outputDir), zip(begins, ends))