examples/drich/drich_utils.py

0001 """dRICH optimization utilities and AID2E workflow components."""
0002
0003 import json
0004 import re
0005 from pathlib import Path
0006 from types import SimpleNamespace
0007
0008 import numpy as np
0009 import uncertainties
0010
0011 from aid2e.utilities.configurations import load_config, load_raw_config
0012
0013
0014 def make_paths(output_dir):
0015     """Return the shared output directories used by the dRICH example."""
0016
0017     output_root = Path(output_dir).resolve()
0018     log_dir = output_root / "log"
0019     return SimpleNamespace(
0020         output_root=output_root,
0021         log_dir=log_dir,
0022         results_dir=log_dir / "results",
0023         sim_dir=log_dir / "sim_files",
0024     )
0025
0026
0027 def build_sim_arguments(npart, point, particle):
0028     """Return npsim gun arguments for one dRICH scan point and particle."""
0029
0030     return [
0031         "-G",
0032         f"-N {npart}",
0033         f"--gun.etaMax {point['eta_max']}",
0034         f"--gun.etaMin {point['eta_min']}",
0035         "--gun.phiMin 0",
0036         "--gun.phiMax 6.2831853",
0037         f"--gun.momentumMax '{point['p']}*GeV'",
0038         f"--gun.momentumMin '{point['p']}*GeV'",
0039         f"--gun.particle {particle}",
0040         "--gun.distribution eta",
0041     ]
0042
0043
0044 def build_reco_arguments(trial_xml):
0045     """Return dRICH reconstruction arguments."""
0046
0047     reco_collections = (
0048         "DRICHHits,MCParticles,DRICHRawHits,DRICHRawHitsAssociations,"
0049         "DRICHAerogelTracks,DRICHGasTracks,"
0050         "DRICHAerogelIrtCherenkovParticleID,DRICHGasIrtCherenkovParticleID,"
0051         "DRICHMergedIrtCherenkovParticleID"
0052     )
0053     return [
0054         f"-Pdd4hep:xml_files={trial_xml}",
0055         f"-Ppodio:output_include_collections={reco_collections}",
0056     ]
0057
0058
0059 def build_analysis_arguments(point, eval_config):
0060     """Return dRICH analysis arguments for one scan point."""
0061
0062     return [
0063         str(point["radiator"]),
0064         str(eval_config["bootstrap_samples"]),
0065         str(eval_config["nbootstraps"]),
0066     ]
0067
0068
0069 def sim_reco_files(sim_dir, npart, trial_tag, point, particle):
0070     """Return dRICH simulation and reconstruction file paths for one scan point."""
0071
0072     tag = f"{npart}_{trial_tag}_{particle}_p_{point['p']}_eta_{point['eta_min']}_{point['eta_max']}"
0073     return Path(sim_dir) / f"scan_{tag}.root", Path(sim_dir) / f"recon_scan_{tag}.root"
0074
0075
0076 def stage_scan_work(layer_names, job_index, eval_config):
0077     """Map one worker job index to the dRICH scan point and particles it should run."""
0078
0079     particles = eval_config["particles"]
0080     scan_points = eval_config["scan_points"]
0081     if "ana" in layer_names:
0082         return scan_points[job_index], particles
0083
0084     point_index, particle_index = divmod(job_index, len(particles))
0085     return scan_points[point_index], [particles[particle_index]]
0086
0087
0088 def apply_overlap_policy(overlap_log, penalty_file, failure_policy, ok_value=0, error=None):
0089     """Apply overlap failure policy and return a penalty marker if needed."""
0090
0091     overlap_log = Path(overlap_log)
0092     penalty_file = Path(penalty_file)
0093     overlap_text = overlap_log.read_text() if overlap_log.exists() else ""
0094     match = re.search(r"Number of illegal overlaps/extrusions\s*:\s*(\d+)", overlap_text)
0095     overlaps = int(match.group(1)) if match is not None else None
0096     use_penalty = failure_policy == "penalty"
0097
0098     def penalty(**payload):
0099         # Later stages check this marker and skip expensive work for failed geometry.
0100         penalty_file.parent.mkdir(parents=True, exist_ok=True)
0101         penalty_file.write_text(json.dumps({"penalty": True, **payload}, indent=2))
0102         return {"ok": 1.0}
0103
0104     if error is not None:
0105         if not use_penalty:
0106             raise error
0107         payload = {"reason": "checkOverlaps command failed", "return_code": getattr(error, "returncode", 1)}
0108         if overlaps is not None and overlaps != ok_value:
0109             payload = {"overlaps": overlaps}
0110         return penalty(**payload)
0111
0112     if overlaps is None:
0113         if use_penalty:
0114             return penalty(reason="checkOverlaps did not print overlap count")
0115         raise RuntimeError("Overlap check failed: no overlap count found")
0116
0117     if overlaps != ok_value:
0118         if use_penalty:
0119             return penalty(overlaps=overlaps)
0120         raise RuntimeError(f"Overlap check failed: overlaps={overlaps}")
0121
0122     return None
0123
0124
0125 def compute_drich_objectives(
0126     results_dir,
0127     trial_tag,
0128     eval_config,
0129     penalty=False,
0130 ):
0131     """Compute dRICH objective metrics from per-scan analysis outputs."""
0132
0133     failed_metrics = {
0134         key: metric_value
0135         for name, failed_value in eval_config["failed_objectives"].items()
0136         for key, metric_value in ((name, float(failed_value)), (f"{name}_sem", 0.0))
0137     }
0138     if penalty:
0139         return failed_metrics
0140
0141     results_dir = Path(results_dir)
0142     npart = eval_config["npart"]
0143     nsigma, eff, momenta = [], [], []
0144     # dRICHAna_bootstrap writes one text file per scan point.
0145     for point in eval_config["scan_points"]:
0146         p, eta_min, eta_max = point["p"], point["eta_min"], point["eta_max"]
0147         result = np.loadtxt(results_dir / f"recon_scan_{npart}_{trial_tag}_p_{p}_eta_{eta_min}_{eta_max}.txt")
0148         nsigma.append(uncertainties.ufloat(result[2], result[3]))
0149         eff.append(uncertainties.ufloat(result[0], result[1]))
0150         momenta.append(p)
0151
0152     nsigma = np.array(nsigma)
0153     eff = np.array(eff)
0154     momenta = np.array(momenta)
0155
0156     def metric(value):
0157         return float(value.n), float(value.s)
0158
0159     # These final objective definitions are dRICH-specific physics choices.
0160     piKsep_etalow, piKsep_etalow_sem = metric(np.mean(nsigma[momenta == 15]))
0161     piKsep_etahigh, piKsep_etahigh_sem = metric(np.mean(nsigma[momenta == 45]))
0162     acceptance, acceptance_sem = metric(np.mean(eff[1:]))
0163     metrics = {
0164         "piKsep_etalow": piKsep_etalow,
0165         "piKsep_etalow_sem": piKsep_etalow_sem,
0166         "piKsep_etahigh": piKsep_etahigh,
0167         "piKsep_etahigh_sem": piKsep_etahigh_sem,
0168         "acceptance": acceptance,
0169         "acceptance_sem": acceptance_sem,
0170     }
0171
0172     if any(np.isnan(value) for value in metrics.values()):
0173         if eval_config["failure_policy"] != "penalty":
0174             raise RuntimeError(f"retrieve_results failed: NaN objective for trial {trial_tag}")
0175         return {name: failed_metrics.get(name, value) for name, value in metrics.items()}
0176
0177     return metrics
0178
0179
0180 def load_drich_config(config_path):
0181     """Load typed AID2E config plus dRICH evaluation_config."""
0182
0183     config_path = Path(config_path).resolve()
0184     raw_cfg = load_raw_config(str(config_path))
0185     return config_path, load_config(str(config_path)), raw_cfg["problem"]["evaluation_config"]