optiphy/tools/run_performance.py

0001 import argparse
0002 import subprocess
0003 import re
0004 import os
0005 from pathlib import Path
0006
0007 run_mac_template = """
0008 /run/numberOfThreads {threads}
0009 /run/verbose 1
0010 /process/optical/cerenkov/setStackPhotons {flag}
0011 /run/initialize
0012 /run/beamOn 500
0013 """
0014
0015 os.environ["OPTICKS_EVENT_MODE"] = "Minimal"
0016
0017 def get_opticks_home():
0018     """Get OPTICKS_HOME from environment, warn if not set."""
0019     opticks_home = os.environ.get("OPTICKS_HOME")
0020     if opticks_home is None:
0021         print("Warning: $OPTICKS_HOME is not defined, so this script should be called from the eic-opticks directory.")
0022         return Path(".")
0023     return Path(opticks_home)
0024
0025 def parse_real_time(time_str):
0026     # Parses 'real\t0m41.149s' to seconds
0027     match = re.search(r'real\s+(\d+)m([\d.]+)s', time_str)
0028     if match:
0029         minutes = int(match.group(1))
0030         seconds = float(match.group(2))
0031         return minutes * 60 + seconds
0032     return None
0033
0034 def parse_sim_time(output):
0035     match = re.search(r"Simulation time:\s*([\d.]+)\s*seconds", output)
0036     if match:
0037         return float(match.group(1))
0038     return None
0039
0040 def main():
0041     opticks_home = get_opticks_home()
0042
0043     parser = argparse.ArgumentParser()
0044     parser.add_argument('-g', '--gdml', type=Path, default=None, help="Path to a custom GDML geometry file (relative to current directory)")
0045     parser.add_argument('-o', '--outpath', type=Path, default=Path('./'), help="Path where the output file will be saved")
0046     args = parser.parse_args()
0047
0048     # If gdml not provided, use default path relative to OPTICKS_HOME
0049     if args.gdml is None:
0050         gdml_path = opticks_home / 'tests/geom/opticks_raindrop.gdml'
0051     else:
0052         # User-provided path is used as-is (relative to current directory or absolute)
0053         gdml_path = args.gdml
0054
0055     # run.mac is created in OPTICKS_HOME
0056     run_mac_path = opticks_home / "run.mac"
0057
0058     geant_file = args.outpath / "timing_geant.txt"
0059     optix_file = args.outpath / "timing_optix.txt"
0060     log_file = args.outpath / "g4logs.txt"
0061
0062     with open(geant_file, "w") as gfile, open(optix_file, "w") as ofile, open(log_file, "w") as logfile:
0063         for threads in range(50, 0, -1):
0064             times = {}
0065             sim_time_true = None
0066             for flag in ['true', 'false']:
0067                 # Write run.mac with current flag
0068                 with open(run_mac_path, "w") as rm:
0069                     rm.write(run_mac_template.format(threads=threads, flag=flag))
0070
0071                 # Run with time in bash to capture real/user/sys
0072                 cmd = f"time GPUCerenkov -g {gdml_path} -m {run_mac_path}"
0073                 print(f"Running {threads} threads: {cmd}")
0074                 result = subprocess.run(
0075                     ["bash", "-c", cmd],
0076                     capture_output=True, text=True
0077                 )
0078                 stdout = result.stdout
0079                 stderr = result.stderr
0080                 # Save full output to log file
0081                 logfile.write(f"\n{'='*60}\n")
0082                 logfile.write(f"Threads: {threads}, StackPhotons: {flag}\n")
0083                 logfile.write(f"{'='*60}\n")
0084                 logfile.write(f"--- STDOUT ---\n{stdout}\n")
0085                 logfile.write(f"--- STDERR ---\n{stderr}\n")
0086                 logfile.flush()
0087
0088                 # Save simulation time for true run only
0089                 if flag == 'true':
0090                     sim_time_true = parse_sim_time(stdout + stderr)
0091                     if sim_time_true is not None:
0092                         ofile.write(f"{threads} {sim_time_true}\n")
0093                         ofile.flush()
0094
0095                 # Extract real time
0096                 real_match = re.search(r"real\s+\d+m[\d.]+s", stderr)
0097                 if real_match:
0098                     real_sec = parse_real_time(real_match.group())
0099                     times[flag] = real_sec
0100                 else:
0101                     print(f"[!] Could not find 'real' time for threads={threads} flag={flag}")
0102
0103             # Write the difference to timing_geant.txt (true - false)
0104             if 'true' in times and 'false' in times:
0105                 diff = times['true'] - times['false']
0106                 gfile.write(f"{threads} {diff}\n")
0107                 gfile.flush()
0108             else:
0109                 print(f"[!] Missing times for threads={threads}")
0110
0111     print("Done.")
0112
0113 if __name__ == '__main__':
0114     main()