eic-opticks/tests/test_simg4ox_charged_multithread.py

0001 #!/usr/bin/env python3
0002 """Exercise charged MT collection, timing, and deterministic event isolation."""
0003
0004 import argparse
0005 import csv
0006 import json
0007 import os
0008 from pathlib import Path
0009 import subprocess
0010 import sys
0011
0012 import numpy as np
0013
0014
0015 def main():
0016     parser = argparse.ArgumentParser(description=__doc__)
0017     parser.add_argument("binary", type=Path)
0018     parser.add_argument("case", choices=("cerenkov", "scintillation", "zero"))
0019     args = parser.parse_args()
0020     repo = Path(__file__).resolve().parents[1]
0021     particle = "geantino" if args.case == "zero" else "e-"
0022     scintillation = args.case == "scintillation"
0023     geometry = repo / "tests/geom" / (
0024         "opticks_raindrop_with_scintillation.gdml" if scintillation else "pfrich_min_FINAL.gdml"
0025     )
0026     baseline = {}
0027     for threads in (2, 4):
0028         output = Path.cwd() / f"{args.case}-t{threads}"
0029         output.mkdir(exist_ok=True)
0030         config_dir = output / "config"
0031         config_dir.mkdir(exist_ok=True)
0032         config = json.loads((repo / "config/dev_timing.json").read_text())
0033         config["torch"].update(pos=[0, 0, 0], mom=[0, 0.2, 0.8])
0034         config["event"].update(event_mode="Minimal", output_dir=str(output))
0035         (config_dir / "charged.json").write_text(json.dumps(config))
0036         macro = output / "run.mac"
0037         # Repeating a run must reset event-local buffers, GPU ordering, and timing.
0038         artifacts = ("events.csv", "events.manifest.json", "s_hits.npy", "g_hits.npy")
0039         run_outputs = [output / f"run-{run}" for run in range(2)]
0040         for directory in (output, *run_outputs):
0041             directory.mkdir(exist_ok=True)
0042             for name in artifacts:
0043                 (directory / name).unlink(missing_ok=True)
0044         macro_lines = ["/run/verbose 0", "/run/initialize", "/run/eventModulo 1"]
0045         if scintillation:
0046             # Geant4 11.4.2 multiple scattering changes first-run counts with
0047             # worker count. Isolate collection so every event can be compared.
0048             macro_lines.append("/process/inactivate msc")
0049         for run_output in run_outputs:
0050             macro_lines.extend((
0051                 "/run/beamOn 12",
0052                 f"/control/shell mv {' '.join(artifacts)} {run_output.name}",
0053             ))
0054         macro.write_text("\n".join(macro_lines) + "\n")
0055         env = dict(os.environ, OPTICKS_HOME=str(repo), SIMPHONY_CONFIG_DIR=str(config_dir))
0056         if scintillation:
0057             # Positive counts must come from the scintillation collector.
0058             env["U4__CollectGenstep_G4Cerenkov_modified_DISABLE"] = "1"
0059         command = [str(args.binary.resolve()), "--gdml", str(geometry), "--config", "charged",
0060                    "--macro", str(macro), "--threads", str(threads), "--particle", particle,
0061                    "--momentum-gev-c", "0.0014219697254161214" if scintillation else "20",
0062                    "--multiplicity", "2" if scintillation else "10", "--seed", "42"]
0063         with (output / "stdout.log").open("w") as log:
0064             result = subprocess.run(command, cwd=output, env=env, stdout=log,
0065                                     stderr=subprocess.STDOUT, timeout=90)
0066         if result.returncode:
0067             raise AssertionError(f"Charged {threads}-worker run failed:\n"
0068                                  + (output / "stdout.log").read_text()[-5000:])
0069         for run, run_output in enumerate(run_outputs):
0070             validation = [sys.executable, str(repo / "tests/check_simg4ox_timing.py"),
0071                           str(run_output / "events.csv"), "--events", "12", "--threads", str(threads),
0072                           "--particle", particle]
0073             if args.case != "zero":
0074                 validation.append("--require-photons")
0075             subprocess.run(validation, check=True)
0076             with (run_output / "events.csv").open() as stream:
0077                 rows = list(csv.DictReader(stream))
0078             counts = [tuple(int(row[key]) for key in
0079                             ("num_gensteps", "num_photons", "num_gpu_hits", "num_g4_hits"))
0080                       for row in rows]
0081             assert all(count[3] == 0 for count in counts), "Optical secondaries reached the CPU stack"
0082             if args.case == "zero":
0083                 assert all(count == (0, 0, 0, 0) for count in counts), "Zero-work event leaked gensteps"
0084             else:
0085                 assert all(count[0] > 0 and count[1] > 0 for count in counts)
0086             # The timing validator checks configured threads and worker IDs. Geant4
0087             # need not assign an event to every worker in a short, zero-work run.
0088             hits = np.load(run_output / "s_hits.npy")
0089             assert hits.shape == (sum(count[2] for count in counts), 4, 4)
0090             assert np.load(run_output / "g_hits.npy").shape == (0, 4, 4)
0091             if run not in baseline:
0092                 baseline[run] = counts, hits.view(np.uint32)
0093             else:
0094                 assert counts == baseline[run][0], f"Run {run}: counts changed with MT worker count"
0095                 assert np.array_equal(hits.view(np.uint32), baseline[run][1]), (
0096                     f"Run {run}: GPU event inputs were mixed"
0097                 )
0098     print(f"Charged MT {args.case}: repeat-run counts, hit arrays, and timing passed")
0099
0100
0101 if __name__ == "__main__":
0102     main()