simulation_campaign_hepmc3/scripts/parse_podio_metadata.py

0001 #!/usr/bin/env python3
0002 """Extract geometry_config and gun parameters from podio runs frame, print as JSON.
0003
0004 Usage:
0005   parse_podio_metadata.py <edm4hep.root> [--gun] [--no-beam]
0006
0007   --gun      Also extract gun parameters (for single particle runs)
0008   --no-beam  Omit beam energy and ion species fields (for background runs)
0009 """
0010
0011 import sys
0012 import math
0013 import json
0014 import re
0015 import podio.root_io
0016
0017 from shared_utils import detect_generator, detect_q2, detect_pwg, detect_dsc
0018
0019
0020 def rad_to_deg(r):
0021     return round(float(r) * 180 / math.pi, 2)
0022
0023
0024 def get_str(params, key):
0025     """Return string value for key, or None if not set or stored as 'None'."""
0026     v = params.get['std::string'](key)
0027     if not v.has_value():
0028         return None
0029     s = v.value()
0030     return None if s == "None" else s
0031
0032
0033 include_gun = "--gun" in sys.argv
0034 no_beam = "--no-beam" in sys.argv
0035 rootfile = next(a for a in sys.argv[1:] if not a.startswith("--"))
0036
0037 reader = podio.root_io.Reader(rootfile)
0038 frame = next(iter(reader.get("runs")))
0039 params = frame.get_parameters()
0040
0041 result = {}
0042
0043 # generator and q2 from inputFiles path
0044 input_files = get_str(params, "inputFiles")
0045 if input_files:
0046     input_path = input_files.strip("[]'\" ")
0047     result["generator"] = detect_generator(input_path, is_single=include_gun)
0048     if not include_gun and not no_beam:
0049         result["requester_pwg"] = detect_pwg(input_path)
0050     q2_min, q2_max = detect_q2(input_path)
0051     if q2_min is not None:
0052         result["q2_min_gev2"] = q2_min
0053     if q2_max is not None:
0054         result["q2_max_gev2"] = q2_max
0055
0056 # is_background_mixed from hepmc_merger_background_files
0057 bg_files = get_str(params, "hepmc_merger_background_files")
0058 result["is_background_mixed"] = bool(bg_files and bg_files.strip() not in ("", "None", "[]"))
0059
0060 # requester_dsc from path keywords or is_background_mixed
0061 if input_files:
0062     dsc = detect_dsc(input_path, is_background_mixed=result["is_background_mixed"])
0063     if dsc:
0064         result["requester_dsc"] = dsc
0065
0066 # data_level from outputFile extension: edm4hep -> simulation, edm4eic -> reconstruction
0067 output_file = get_str(params, "outputFile")
0068 if output_file:
0069     output_file = output_file.strip("[]'\" ")
0070     if ".edm4hep." in output_file:
0071         result["data_level"] = "simulation"
0072     elif ".edm4eic." in output_file:
0073         result["data_level"] = "reconstruction"
0074
0075 # geometry_config from compactFile: strip path and .xml, strip leading "epic_"
0076 # compactFile is stored as a list string like "['path/to/file.xml']"
0077 compact_file = get_str(params, "compactFile")
0078 if compact_file:
0079     compact_file = compact_file.strip("[]'\" ")
0080     basename = re.sub(r'\.xml$', '', compact_file.split('/')[-1])
0081     if basename.startswith("epic_"):
0082         basename = basename[len("epic_"):]
0083     result["geometry_config"] = basename
0084
0085     # Parse electron_beam_energy, ion_beam_energy, ion_species from geometry_config
0086     # Format: <detector>_<ebeam>x<pbeam>[_<species>]
0087     # Skip for singles (--gun) and backgrounds (--no-beam)
0088     if not no_beam and not include_gun:
0089         beam_match = re.search(r'_(\d+)x(\d+)(?:_(.+))?$', basename)
0090         if beam_match:
0091             result["electron_beam_energy_gev"] = int(beam_match.group(1))
0092             result["ion_beam_energy_gev"] = int(beam_match.group(2))
0093             species = beam_match.group(3)
0094             result["ion_species"] = species if species is not None else "p"
0095
0096 if include_gun:
0097     # gun_particle
0098     particle = get_str(params, "gun.particle")
0099     if particle:
0100         result["gun_particle"] = particle
0101
0102     # gun_distribution
0103     distribution = get_str(params, "gun.distribution")
0104     if distribution:
0105         result["gun_distribution"] = distribution
0106
0107     # gun_momentum: prefer gun.energy (MeV -> GeV)
0108     # gun.momentumMin/Max default to 0/10000 in npsim when not explicitly set
0109     energy = get_str(params, "gun.energy")
0110     mom_min = get_str(params, "gun.momentumMin")
0111     mom_max = get_str(params, "gun.momentumMax")
0112     npsim_defaults = (mom_min == "0.0" and mom_max == "10000.0")
0113
0114     if energy is not None:
0115         momentum_gev = round(float(energy) / 1000, 6)
0116         result["gun_momentum_min_gev"] = momentum_gev
0117         result["gun_momentum_max_gev"] = momentum_gev
0118     elif not npsim_defaults and mom_min is not None and mom_max is not None:
0119         result["gun_momentum_min_gev"] = round(float(mom_min) / 1000, 6)
0120         result["gun_momentum_max_gev"] = round(float(mom_max) / 1000, 6)
0121
0122     # gun_theta: radians -> degrees
0123     theta_min = get_str(params, "gun.thetaMin")
0124     theta_max = get_str(params, "gun.thetaMax")
0125     if theta_min is not None:
0126         result["gun_theta_min_deg"] = rad_to_deg(theta_min)
0127     if theta_max is not None:
0128         result["gun_theta_max_deg"] = rad_to_deg(theta_max)
0129
0130     # gun_phi: radians -> degrees; default 0 and 360
0131     phi_min = get_str(params, "gun.phiMin")
0132     phi_max = get_str(params, "gun.phiMax")
0133     result["gun_phi_min_deg"] = rad_to_deg(phi_min) if phi_min is not None else 0
0134     result["gun_phi_max_deg"] = rad_to_deg(phi_max) if phi_max is not None else 360
0135
0136 print(json.dumps(result))