235 parser = argparse.ArgumentParser(
236 description=("Tune SegmentEdge thresholds using signal-muon reconstruction "
237 "efficiency relative to the regular no-ML reconstruction."))
238 parser.add_argument("--inputFile", required=True)
239 parser.add_argument("--bucketModel", default=None, help=("Optional bucket-filter ONNX model."),)
240 parser.add_argument("--bucketThreshold", "--score-threshold", dest="bucketThreshold", type=float,
241 default=None, help="Bucket-filter score threshold")
242 parser.add_argument("--bucket-output-is-logit", dest="bucketOutputIsLogit", action="store_true", default=False,
243 help=("Interpret the scalar bucket-model output as a logit"))
244 parser.add_argument("--edgeModel", required=True)
245 parser.add_argument("--nEvents", type=int, default=100)
246 parser.add_argument("--skipEvents", type=int, default=0)
247 parser.add_argument("--threads", type=int, default=1)
248 parser.add_argument("--defaultGeoFile", default="RUN4")
249 parser.add_argument("--workDir", default="edge_threshold_tuning")
250 parser.add_argument("--edgeThresholds", default="0.08,0.10,0.119,0.14,0.16", help="Comma-separated recovery edge-probability thresholds to scan",)
251 parser.add_argument("--overlapThresholds", default="0.20,0.30,0.50,0.80", help=("Comma-separated high-purity core edge-probability thresholds to scan."))
252 parser.add_argument("--targetRelativeEfficiencyLoss", "--targetLoss", dest="targetRelativeEfficiencyLoss",
253 type=float, default=0.05, help=("Maximum allowed relative loss in signal-muon efficiency"),)
254 parser.add_argument("--treeName", default="MsTrackValidTest")
255 parser.add_argument("--signalOrigin", type=int, default=13)
256 parser.add_argument("--signalType", type=int, default=6)
257 parser.add_argument("--use-cpu", dest="use_cpu", action="store_true", default=False,)
258 parser.add_argument("--skipExisting", action="store_true")
259 args = parser.parse_args()
260
261 edge_thresholds = _parse_float_list(args.edgeThresholds)
262 overlap_thresholds = _parse_float_list(args.overlapThresholds)
263
264 work_dir = Path(args.workDir).resolve()
265 work_dir.mkdir(parents=True, exist_ok=True)
266
267 baseline_root = work_dir / "baseline_noml.root"
268 baseline_log = work_dir / "baseline_noml.log"
269 baseline_command = _chain_command(args, baseline_root, edge=False)
270 if not args.skipExisting or not baseline_root.exists():
271 return_code = _run(baseline_command, baseline_log)
272 if return_code != 0:
273 raise SystemExit(
274 f"No-ML baseline job failed with rc={return_code}. See {baseline_log}"
275 )
276 if not baseline_root.exists():
277 raise SystemExit(
278 f"No-ML baseline finished but ROOT output is missing: {baseline_root}. "
279 f"See {baseline_log}"
280 )
281
282 baseline = _signal_muon_efficiency(baseline_root, args.treeName, args.signalOrigin, args.signalType,)
283
284 rows: list[dict[str, float | int | str | bool]] = []
285
286 best = None
287 for edge_threshold in edge_thresholds:
288 for overlap_threshold in overlap_thresholds:
289 tag = (
290 f"edge{edge_threshold:.6f}_overlap{overlap_threshold:.6f}"
292 )
293 out_root = work_dir / f"{tag}.root"
294 out_log = work_dir / f"{tag}.log"
295 command = _chain_command(
296 args,
297 out_root,
298 edge=True,
299 edge_threshold=edge_threshold,
300 overlap_threshold=overlap_threshold,
301 )
302 if not args.skipExisting or not out_root.exists():
303 return_code = _run(command, out_log)
304 if return_code != 0:
305 rows.append({
306 "edgeThreshold": edge_threshold,
307 "overlapThreshold": overlap_threshold,
308 "status": "failed",
309 "rootFile": str(out_root),
310 "log": str(out_log),
311 })
312 continue
313 if not out_root.exists():
314 rows.append({
315 "edgeThreshold": edge_threshold,
316 "overlapThreshold": overlap_threshold,
317 "status": "missing_output",
318 "rootFile": str(out_root),
319 "log": str(out_log),
320 })
321 continue
322
323 edge = _signal_muon_efficiency(
324 out_root,
325 args.treeName,
326 args.signalOrigin,
327 args.signalType,
328 )
329 row = _result_row(
330 edge_threshold,
331 overlap_threshold,
332 baseline,
333 edge,
334 args.targetRelativeEfficiencyLoss,
335 out_root,
336 out_log,
337 )
338
339 rows.append(row)
340
341 if row["status"] == "ok" and row["passesTarget"]:
342 key = (edge_threshold, overlap_threshold)
343 if best is None or key > (
344 best["edgeThreshold"],
345 best["overlapThreshold"],
346 ):
347 best = row
348
349 csv_path = work_dir / "edge_threshold_scan.csv"
350 fieldnames = [
351 "edgeThreshold",
352 "overlapThreshold",
353 "status",
354 "baselineSignalMuonCount",
355 "baselineMatchedSignalMuons",
356 "baselineSignalEfficiency",
357 "edgeSignalMuonCount",
358 "edgeMatchedSignalMuons",
359 "edgeSignalEfficiency",
360 "absoluteEfficiencyDifference",
361 "relativeEfficiencyDifference",
362 "relativeEfficiencyLoss",
363 "passesTarget",
364 "rootFile",
365 "log",
366 ]
367 with csv_path.open("w", newline="", encoding="utf-8") as csv_file:
368 writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
369 writer.writeheader()
370 writer.writerows(rows)
371
372 summary = {
373 "treeName": args.treeName,
374 "signalSelection": {
375 "truthOrigin": args.signalOrigin,
376 "truthType": args.signalType,
377 },
378 "matching": (
379 "Truth muon <- MsTrkSeed_truthLink -> seed <- "
380 "ActsMuons_seedLink -> reconstructed track"
381 ),
382 "baseline": baseline,
383 "bucketFilter": {
384 "model": args.bucketModel,
385 "scoreThreshold": args.bucketThreshold,
386 "scoreThresholdSource": (
387 "CLI override" if args.bucketThreshold is not None
388 else "GraphBucketFilterToolCfg default"
389 ),
390 },
391 "segmentEdgeGraph": {
392 "ReadSpacePoints": "FilteredMlBuckets",
393 "OrderingSpacePoints": "MuonSpacePoints",
394 "note": (
395 "Applied by muonEdgeRecoChain.py when bucket filtering and "
396 "edge inference are enabled."
397 ),
398 },
399 "targetRelativeEfficiencyLoss": args.targetRelativeEfficiencyLoss,
400 "selectionPolicy": (
401 "largest (edgeThreshold, overlapThreshold) pair with relative "
402 "signal-efficiency loss at or below the target"
403 ),
404 "best": best,
405 "scanCsv": str(csv_path),
406 }
407 (work_dir / "summary.json").write_text(
408 json.dumps(summary, indent=2),
409 encoding="utf-8",
410 )
411 print(json.dumps(summary, indent=2))
412
413
void print(char *figname, TCanvas *c1)
std::string replace(std::string s, const std::string &s2, const std::string &s3)