Scripts/Python/guntam_transformer_seeder.py

0001 """
0002 guntam_transformer_seeder.py
0003
0004 Integrates the GUNTAM ONNX transformer into ACTS as a PythonCallable seeder.
0005
0006 Requires ``acts.SpacePointContainer2``, ``acts.SeedContainer2``, and the
0007 associated Python bindings for mutable space points and seeds.
0008
0009 Usage:
0010     from guntam_transformer_seeder import guntam_transformer_seeder
0011
0012     addSeeding(
0013         s,
0014         trackingGeometry,
0015         field,
0016         seedingAlgorithm=SeedingAlgorithm.PythonCallable,
0017         customSeeder=guntam_transformer_seeder,
0018         customSeederConfig={
0019             "model_path": "/path/to/model.onnx",
0020             "r_max": 500.0,
0021             "z_max": 1000.0,
0022             "score_threshold": 0.35,
0023             "providers": ["CUDAExecutionProvider", "CPUExecutionProvider"],
0024         },
0025         ...
0026     )
0027 """
0028
0029 import numpy as np
0030 import onnxruntime as ort
0031
0032 import acts
0033 import acts.examples
0034
0035 # Numerical stability guard: seeds where two spacepoints are identical in the transverse
0036 # plane crash the CKF. This is not a physics threshold, it exists solely to prevent
0037 # crashes from pathological seeds.
0038 _SP_DEDUP_GUARD_MM = 1e-3
0039
0040
0041 def _apply_model_acceptance(
0042     sp, xyz: np.ndarray, r_max: float, z_max: float
0043 ) -> tuple[np.ndarray, np.ndarray]:
0044     """Return (filtered_sp_xyz, original_indices) after applying the acceptance cut.
0045
0046     Indices in original_indices map positions in filtered_sp_xyz back to rows in xyz.
0047     sp.r is read directly from the stored column rather than recomputed from x, y.
0048     """
0049     r = np.asarray(sp.r)
0050     mask = (r < r_max) & (np.abs(xyz[:, 2]) < z_max)
0051     return xyz[mask].astype(np.float32), np.where(mask)[0].astype(np.uint32)
0052
0053
0054 def _filter_valid_seeds(
0055     seeds: np.ndarray, scores: np.ndarray, sp_xyz: np.ndarray | None = None
0056 ) -> tuple[np.ndarray, np.ndarray]:
0057     """Keep only seeds with a finite score, all three spacepoint indices filled, and no
0058     duplicate indices. When sp_xyz is provided, also drops seeds with degenerate
0059     pairwise transverse separation (crash guard, see _SP_DEDUP_GUARD_MM)."""
0060     valid = np.isfinite(scores) & (np.sum(seeds >= 0, axis=1) == 3)
0061     seeds_v = seeds[valid]
0062     scores_v = scores[valid]
0063     if len(seeds_v) > 0:
0064         i0, i1, i2 = seeds_v[:, 0], seeds_v[:, 1], seeds_v[:, 2]
0065         no_dups = (i0 != i1) & (i1 != i2) & (i0 != i2)
0066         seeds_v = seeds_v[no_dups]
0067         scores_v = scores_v[no_dups]
0068     if sp_xyz is not None and len(seeds_v) > 0:
0069         i0, i1, i2 = seeds_v[:, 0], seeds_v[:, 1], seeds_v[:, 2]
0070         xy = sp_xyz[:, :2]
0071         d01 = np.linalg.norm(xy[i0] - xy[i1], axis=1)
0072         d12 = np.linalg.norm(xy[i1] - xy[i2], axis=1)
0073         d02 = np.linalg.norm(xy[i0] - xy[i2], axis=1)
0074         non_degen = np.minimum(np.minimum(d01, d12), d02) >= _SP_DEDUP_GUARD_MM
0075         seeds_v = seeds_v[non_degen]
0076         scores_v = scores_v[non_degen]
0077     return seeds_v, scores_v
0078
0079
0080 class _GuntamAlgorithm(acts.examples.IAlgorithm):
0081     """Per-event algorithm: reads spacepoints, runs GUNTAM ONNX, writes seeds.
0082     The ONNX session is created once at construction. Configure num_threads
0083     conservatively when running the sequencer with multiple event threads.
0084     """
0085
0086     def __init__(
0087         self,
0088         model_path: str,
0089         sp_key: str,
0090         seeds_key: str,
0091         log_level,
0092         num_threads: int = 1,
0093         r_max: float = 500.0,
0094         z_max: float = 1000.0,
0095         score_threshold: float = 0.35,
0096         providers: list[str] | None = None,
0097     ):
0098         acts.examples.IAlgorithm.__init__(self, "GuntamTransformerSeeder", log_level)
0099
0100         self._r_max = r_max
0101         self._z_max = z_max
0102         self._score_threshold = score_threshold
0103
0104         self._sp_handle = acts.examples.ReadDataHandle(
0105             self, acts.SpacePointContainer2, "InputSpacePoints"
0106         )
0107         self._sp_handle.initialize(sp_key)
0108
0109         self._seeds_handle = acts.examples.WriteDataHandle(
0110             self, acts.SeedContainer2, "OutputSeeds"
0111         )
0112         self._seeds_handle.initialize(seeds_key)
0113
0114         sess_options = ort.SessionOptions()
0115         sess_options.intra_op_num_threads = num_threads
0116         sess_options.inter_op_num_threads = num_threads
0117
0118         self._session = ort.InferenceSession(
0119             model_path,
0120             sess_options=sess_options,
0121             providers=providers or ["CPUExecutionProvider"],
0122         )
0123
0124     def execute(self, ctx) -> acts.examples.ProcessCode:
0125         sp = self._sp_handle(ctx.eventStore)
0126         xyz = np.stack([np.asarray(sp.x), np.asarray(sp.y), np.asarray(sp.z)], axis=1)
0127
0128         filtered_sp_xyz, orig_idx = _apply_model_acceptance(
0129             sp, xyz, self._r_max, self._z_max
0130         )
0131
0132         seeds_raw, scores_raw = self._session.run(
0133             ["seeds", "seed_scores"],
0134             {"hits": filtered_sp_xyz},
0135         )
0136
0137         seeds_v, scores_v = _filter_valid_seeds(seeds_raw, scores_raw, filtered_sp_xyz)
0138
0139         score_mask = scores_v > self._score_threshold
0140         seeds_v = seeds_v[score_mask]
0141         scores_v = scores_v[score_mask]
0142
0143         container = acts.SeedContainer2()
0144         container.assignSpacePointContainer(sp)
0145
0146         sp_indices_v = orig_idx[seeds_v]
0147
0148         seed_proxies = [container.createSeed() for _ in range(len(seeds_v))]
0149         for seed, sp_indices, score in zip(seed_proxies, sp_indices_v, scores_v):
0150             seed.quality = float(score)
0151             seed.vertexZ = 0.0
0152             seed.assignSpacePointIndices(sp_indices.tolist())
0153
0154         # MutableSeedProxy2 holds raw pointers into the container; drop them before
0155         # the whiteboard write which transfers container ownership to C++.
0156         del seed_proxies
0157
0158         self._seeds_handle(ctx, container)
0159         return acts.examples.ProcessCode.SUCCESS
0160
0161
0162 def guntam_transformer_seeder(
0163     s,
0164     spacePoints: str,
0165     outputSeeds: str,
0166     config: dict,
0167     **kwargs,
0168 ) -> str:
0169     """PythonCallable entry point for the GUNTAM ONNX seeder.
0170
0171     config keys:
0172         ``model_path``      (str, required)            — path to the ONNX model file.
0173         ``num_threads``     (int, default 1)            — ONNX intra/inter-op thread count.
0174             When running events in parallel (Sequencer.numThreads > 1), ensure
0175             numThreads * num_threads does not exceed available CPU cores.
0176         ``r_max``           (float, default 500.0)     — radial acceptance cut in mm.
0177         ``z_max``           (float, default 1000.0)    — longitudinal acceptance cut in mm.
0178         ``score_threshold`` (float, default 0.35)      — minimum score to keep a seed.
0179         ``providers``       (list[str], default None)  — ONNX execution providers in
0180             priority order. None falls back to ["CPUExecutionProvider"]. Pass
0181             ["CUDAExecutionProvider", "CPUExecutionProvider"] for GPU with CPU fallback.
0182
0183     **kwargs absorbs trackingGeometry, logLevel, and any future addSeeding additions.
0184     Returns the whiteboard key for the output seeds.
0185     """
0186     try:
0187         model_path = config["model_path"]
0188     except KeyError as exc:
0189         raise ValueError("customSeederConfig must contain 'model_path'") from exc
0190     num_threads = int(config.get("num_threads", 1))
0191     r_max = float(config.get("r_max", 500.0))
0192     z_max = float(config.get("z_max", 1000.0))
0193     score_threshold = float(config.get("score_threshold", 0.35))
0194     providers = config.get("providers", None)
0195     log_level = kwargs.get("logLevel", acts.logging.INFO)
0196     s.addAlgorithm(
0197         _GuntamAlgorithm(
0198             model_path,
0199             spacePoints,
0200             outputSeeds,
0201             log_level,
0202             num_threads=num_threads,
0203             r_max=r_max,
0204             z_max=z_max,
0205             score_threshold=score_threshold,
0206             providers=providers,
0207         )
0208     )
0209     return outputSeeds