Scripts/Python/hepmc3_merge.py

0001 #!/usr/bin/env python3
0002
0003 import argparse
0004 from pathlib import Path
0005 import re
0006
0007 import acts
0008 import acts.examples
0009 import acts.examples.hepmc3
0010
0011
0012 def float_with_unit(s: str) -> float:
0013     m = re.match(r"(\d+\.?\d*)(.*)", s)
0014     assert m is not None, f"Invalid length string: {s}"
0015
0016     num, unit = m.groups()
0017
0018     scale = getattr(acts.UnitConstants, unit, 1)
0019
0020     return float(num) * scale
0021
0022
0023 def main():
0024     parser = argparse.ArgumentParser()
0025     parser.add_argument(
0026         "--hard-scatter", "--hs", type=Path, help="Hard scatter file", required=True
0027     )
0028     parser.add_argument(
0029         "--hs-input-events",
0030         type=int,
0031         default=None,
0032         help="If None, the HepMC3 will read the entire file to determine the number of available events.",
0033     )
0034     parser.add_argument(
0035         "--pileup", "--pu", type=Path, help="Pileup file", required=True
0036     )
0037     parser.add_argument("--output", "-o", type=Path, help="Output file", required=True)
0038     parser.add_argument("--force", "-f", action="store_true", help="Force overwrite")
0039     parser.add_argument(
0040         "--pileup-multiplicity",
0041         "--npu",
0042         type=int,
0043         help="Pileup multiplicity",
0044         default=1,
0045     )
0046
0047     parser.add_argument(
0048         "--vtx-pos",
0049         "--vp",
0050         type=float_with_unit,
0051         help="Vertex position, use suffix like mm for units",
0052         default=[0.0, 0.0, 0.0, 0.0],
0053         nargs=4,
0054     )
0055
0056     parser.add_argument(
0057         "--vtx-stddev",
0058         "--vs",
0059         type=float_with_unit,
0060         help="Vertex stddev, use suffix like mm for units",
0061         default=[0.0, 0.0, 0.0, 0.0],
0062         nargs=4,
0063     )
0064
0065     parser.add_argument("--jobs", "-j", type=int, help="Number of jobs", default=-1)
0066
0067     args = parser.parse_args()
0068
0069     if not args.hard_scatter.exists():
0070         raise FileNotFoundError(f"Hard scatter file {args.hard_scatter} does not exist")
0071     if not args.pileup.exists():
0072         raise FileNotFoundError(f"Pileup file {args.pileup} does not exist")
0073
0074     extensions = {
0075         acts.examples.hepmc3.compressionExtension(c): c
0076         for c in acts.examples.hepmc3.availableCompressionModes()
0077         if c != acts.examples.hepmc3.Compression.none
0078     }
0079
0080     compression = extensions.get(
0081         args.output.suffix, acts.examples.hepmc3.Compression.none
0082     )
0083
0084     if compression != acts.examples.hepmc3.Compression.none:
0085         output_path = Path(args.output.stem)
0086     else:
0087         output_path = args.output
0088
0089     if output_path.exists():
0090         if args.force:
0091             output_path.unlink()
0092         else:
0093             raise FileExistsError(
0094                 f"Output file {args.output} already exists, run with --force to overwrite"
0095             )
0096
0097     s = acts.examples.Sequencer(numThreads=args.jobs, logLevel=acts.logging.WARNING)
0098
0099     rng = acts.examples.RandomNumbers(seed=42)
0100     s.addReader(
0101         acts.examples.hepmc3.HepMC3Reader(
0102             inputPaths=[
0103                 (args.hard_scatter, 1),
0104                 (args.pileup, args.pileup_multiplicity),
0105             ],
0106             level=acts.logging.INFO,
0107             outputEvent="hepmc3_event",
0108             checkEventNumber=False,  # This is not generally guaranteed for arbitrary inputs
0109             numEvents=args.hs_input_events,
0110             randomNumbers=rng,
0111             vertexGenerator=acts.examples.GaussianVertexGenerator(
0112                 stddev=acts.Vector4(
0113                     args.vtx_stddev[0],
0114                     args.vtx_stddev[1],
0115                     args.vtx_stddev[2],
0116                     args.vtx_stddev[3],
0117                 ),
0118                 mean=acts.Vector4(
0119                     args.vtx_pos[0],
0120                     args.vtx_pos[1],
0121                     args.vtx_pos[2],
0122                     args.vtx_pos[3],
0123                 ),
0124             ),
0125         )
0126     )
0127
0128     s.addWriter(
0129         acts.examples.hepmc3.HepMC3Writer(
0130             inputEvent="hepmc3_event",
0131             outputPath=output_path,
0132             level=acts.logging.INFO,
0133             compression=compression,
0134             writeEventsInOrder=False,
0135         )
0136     )
0137
0138     s.run()
0139
0140
0141 if __name__ == "__main__":
0142     main()