Tracking/svt_rate_per_source/plot_hit_rate_per_source.py

0001 #!/usr/bin/env python3
0002 """Plot SVT reconstructed-hit rates by MC source from chunked EDM4eic events.
0003
0004 Run with --help. All project imports are local to this folder; ROOT and PODIO
0005 Python bindings are unnecessary. See README.md for normalization assumptions.
0006 """
0007
0008 import argparse
0009 import csv
0010 import hashlib
0011 import json
0012 import math
0013 import re
0014 import time
0015 from collections import Counter, defaultdict
0016 from datetime import datetime, timezone
0017 from pathlib import Path
0018
0019 import awkward as ak
0020 import matplotlib
0021
0022 matplotlib.use("Agg")  # Batch jobs and remote terminals do not need a display.
0023 import matplotlib.pyplot as plt
0024 import numpy as np
0025 import uproot
0026 from matplotlib.backends.backend_pdf import PdfPages
0027
0028 from svt_sources import (BACKGROUND_COMPONENTS, RELATIONS, SOURCE_ORDER,
0029                          STATUS_TO_SOURCE, SURFACES, assign_sources, read_branch)
0030
0031
0032 HERE = Path(__file__).resolve().parent
0033 SUMMARY_COLUMNS = ["lname", "source", "rec_branch", "raw_hits",
0034                    "scaled_total_rate_khz", "rsu_or_tile_max_rate_khz",
0035                    "raw_rsu_or_tile_max", "events_with_hits",
0036                    "max_raw_hits_per_event_tile", "outside_map_hits"]
0037
0038
0039 def parse_args():
0040     parser = argparse.ArgumentParser(description=__doc__)
0041     inputs = parser.add_mutually_exclusive_group(required=True)
0042     inputs.add_argument("--input-files", nargs="+", help="Local ROOT files or full root:// URLs; optional :events suffix.")
0043     inputs.add_argument("--file-list", type=Path, help="One input per line; blank lines and # comments are ignored.")
0044     parser.add_argument("--label", default="svt_rate", help="Name of a fresh output subdirectory (default: svt_rate).")
0045     parser.add_argument("--output-dir", type=Path, default=HERE / "outputs", help="Parent output directory (default: outputs beside this script).")
0046     parser.add_argument("--entry-start", type=int, default=0, help="First event per file, inclusive (default: 0).")
0047     parser.add_argument("--entry-stop", type=int, help="Last event per file, exclusive (default: end of file).")
0048     parser.add_argument("--max-files", type=int, help="Use only the first N files for a short check.")
0049     parser.add_argument("--chunk-size", type=int, default=10, help="Events read at once (default: 10); does not change normalization.")
0050     parser.add_argument("--frame-us", type=float, default=2.0, help="Physical duration of each equal-weight event frame in us (default: 2).")
0051     parser.add_argument("--tile", action="store_true", help="Use 3.5 x 9.8 mm bins; otherwise use 20 x 20 mm RSU bins.")
0052     parser.add_argument("--min-surface-hits", type=int, default=1, help="Minimum in-map source hits for a PDF page; CSV counts are unaffected.")
0053     args = parser.parse_args()
0054     if args.entry_start < 0 or (args.entry_stop is not None and args.entry_stop <= args.entry_start):
0055         parser.error("Require 0 <= entry-start < entry-stop (when supplied)")
0056     if args.chunk_size <= 0 or (args.max_files is not None and args.max_files <= 0):
0057         parser.error("chunk-size and max-files must be positive")
0058     if not math.isfinite(args.frame_us) or args.frame_us <= 0 or args.min_surface_hits < 1:
0059         parser.error("frame-us must be finite and positive; min-surface-hits must be >= 1")
0060     if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", args.label):
0061         parser.error("label must start with a letter/digit and contain only letters, digits, _, -, or .")
0062     return args
0063
0064
0065 def build_inputs(args):
0066     """Resolve local paths and preserve full remote URLs, in the supplied order."""
0067     if args.file_list:
0068         base = args.file_list.resolve().parent
0069         items = [line.strip() for line in args.file_list.read_text().splitlines()
0070                  if line.strip() and not line.lstrip().startswith("#")]
0071     else:
0072         base = Path.cwd()
0073         items = args.input_files
0074     if args.max_files:
0075         items = items[:args.max_files]
0076     paths = []
0077     for item in items:
0078         path = item[:-len(":events")] if item.endswith(":events") else item
0079         if not path.endswith(".root"):
0080             raise ValueError(f"Expected a ROOT file with optional :events suffix: {item}")
0081         if "://" not in path:
0082             path = str((base / Path(path).expanduser()).resolve())
0083         paths.append(path)
0084     if not paths:
0085         raise ValueError("The input list is empty")
0086     if len(set(paths)) != len(paths):
0087         raise ValueError("Repeated input files would double-count events; remove duplicate entries")
0088     return paths
0089
0090
0091 def make_surface_maps(tile=False):
0092     """Build fixed grids; keep complete nominal-size bins at the positive edge."""
0093     dx, dy = (3.5, 9.8) if tile else (20.0, 20.0)
0094     maps = {}
0095     for name, branch, barrel, window in SURFACES:
0096         extent = (int(window[1] * np.pi // 100) + 1) * 100 if barrel else 500
0097         # The original workflow stopped one bin short; include the full extent.
0098         xedges = -extent + dx * np.arange(math.ceil(2 * extent / dx) + 1)
0099         yedges = -extent + dy * np.arange(math.ceil(2 * extent / dy) + 1)
0100         shape = (len(xedges) - 1, len(yedges) - 1)
0101         maps[name] = {"branch": branch, "barrel": barrel, "window": window,
0102                       "xedges": xedges, "yedges": yedges,
0103                       "counts": {"ALL": np.zeros(shape, dtype=np.int64)},
0104                       "counters": defaultdict(Counter)}
0105     return maps
0106
0107
0108 def fill_event(surface_maps, branch, hits, sources):
0109     """Fill ALL, exclusive sources, and the inclusive other-background group."""
0110     if len(hits) != len(sources):
0111         raise ValueError(f"{branch}: rec-hit and raw-hit relation lengths disagree")
0112     x, y, z = (ak.to_numpy(hits[field]) for field in ("position.x", "position.y", "position.z"))
0113     if not (np.all(np.isfinite(x)) and np.all(np.isfinite(y)) and np.all(np.isfinite(z))):
0114         raise ValueError(f"{branch}: nonfinite reconstructed-hit coordinates")
0115     radius = np.hypot(x, y)
0116     for surface in surface_maps.values():
0117         if surface["branch"] != branch:
0118             continue
0119         low, high = surface["window"]
0120         selector = radius if surface["barrel"] else z
0121         selected = (selector > low) & (selector < high)
0122         if not np.any(selected):
0123             continue
0124         px = z[selected] if surface["barrel"] else x[selected]
0125         py = (radius * np.arctan2(y, x))[selected] if surface["barrel"] else y[selected]
0126         labels = sources[selected]
0127         masks = {"ALL": np.ones(len(labels), dtype=bool)}
0128         masks.update({source: labels == source for source in np.unique(labels)})
0129         background = np.isin(labels, list(BACKGROUND_COMPONENTS))
0130         if np.any(background):
0131             masks["other background"] = background
0132         for source, mask in masks.items():
0133             counts = np.histogram2d(px[mask], py[mask], bins=(surface["xedges"], surface["yedges"]))[0].astype(np.int64)
0134             if source not in surface["counts"]:
0135                 surface["counts"][source] = np.zeros_like(counts)
0136             surface["counts"][source] += counts
0137             counter = surface["counters"][source]
0138             in_map = int(counts.sum())
0139             counter["events_with_hits"] += int(in_map > 0)
0140             counter["max_raw_hits_per_event_tile"] = max(counter["max_raw_hits_per_event_tile"], int(counts.max()))
0141             counter["outside_map_hits"] += int(mask.sum()) - in_map
0142
0143
0144 def process_files(paths, args):
0145     """Keep one event chunk and accumulated maps in memory; share MC truth."""
0146     maps = make_surface_maps(args.tile)
0147     diagnostics = {config["rec"]: Counter() for config in RELATIONS}
0148     inputs = []
0149     total_events = 0
0150     started = time.perf_counter()
0151     for file_id, path in enumerate(paths):
0152         print(f"[{file_id + 1}/{len(paths)}] {path}", flush=True)
0153         # Disable array caching: completed chunks need not remain in memory.
0154         with uproot.open(path, array_cache=None) as root_file:
0155             tree = root_file["events"]
0156             stop = tree.num_entries if args.entry_stop is None else min(args.entry_stop, tree.num_entries)
0157             start = min(args.entry_start, tree.num_entries)
0158             inputs.append({"file_id": file_id, "input_url": path + ":events",
0159                            "tree_events": tree.num_entries, "entry_start": start,
0160                            "entry_stop": max(start, stop), "processed_events": max(0, stop - start)})
0161             for chunk_start in range(start, stop, args.chunk_size):
0162                 chunk_stop = min(chunk_start + args.chunk_size, stop)
0163                 print(f"  events [{chunk_start}, {chunk_stop})", flush=True)
0164                 mc = read_branch(tree, "MCParticles", chunk_start, chunk_stop,
0165                                  ["generatorStatus", "parents_begin", "parents_end"])
0166                 parents = read_branch(tree, "_MCParticles_parents", chunk_start, chunk_stop, ["index"])
0167                 for config in RELATIONS:
0168                     branch = config["rec"]
0169                     hits = read_branch(tree, branch, chunk_start, chunk_stop,
0170                                        ["position.x", "position.y", "position.z"])
0171                     relations = {key: read_branch(tree, config[key], chunk_start, chunk_stop,
0172                                                   ["weight"] if key == "association" else ["index"])
0173                                  for key in ("rec_raw", "association", "association_raw", "association_sim", "sim_particle")}
0174                     for event in range(chunk_stop - chunk_start):
0175                         labels, counts = assign_sources(
0176                             relations["rec_raw"][event], relations["association"][event]["weight"],
0177                             relations["association_raw"][event], relations["association_sim"][event],
0178                             relations["sim_particle"][event], mc[event], parents[event])
0179                         diagnostics[branch].update(counts)
0180                         fill_event(maps, branch, hits[event], labels)
0181                 total_events += chunk_stop - chunk_start
0182     if total_events == 0:
0183         raise ValueError("No events were selected; check input sizes and the requested event range")
0184     # ALL must conserve hit counts bin-by-bin, including unresolved categories.
0185     for name, surface in maps.items():
0186         exclusive = np.zeros_like(surface["counts"]["ALL"])
0187         for source, counts in surface["counts"].items():
0188             if source not in ("ALL", "other background"):
0189                 exclusive += counts
0190         if not np.array_equal(exclusive, surface["counts"]["ALL"]):
0191             raise RuntimeError(f"Source-count conservation failed on {name}")
0192     return maps, inputs, diagnostics, total_events, time.perf_counter() - started
0193
0194
0195 def summarize(maps, scale):
0196     """Return per-surface/source counts and rates, without integer truncation."""
0197     rows = []
0198     for name, surface in maps.items():
0199         for source in SOURCE_ORDER:
0200             if source not in surface["counts"]:
0201                 continue
0202             counts = surface["counts"][source]
0203             raw_hits, raw_peak = int(counts.sum()), int(counts.max())
0204             row = {"lname": name, "source": source, "rec_branch": surface["branch"],
0205                    "raw_hits": raw_hits, "scaled_total_rate_khz": raw_hits * scale,
0206                    "rsu_or_tile_max_rate_khz": raw_peak * scale, "raw_rsu_or_tile_max": raw_peak}
0207             for key in SUMMARY_COLUMNS[-3:]:
0208                 row[key] = int(surface["counters"][source][key])
0209             rows.append(row)
0210     return rows
0211
0212
0213 def write_csv(path, rows, columns, metadata):
0214     with path.open("w", newline="") as output:
0215         for key, value in metadata.items():
0216             output.write(f"# {key}: {value}\n")
0217         writer = csv.DictWriter(output, fieldnames=columns)
0218         writer.writeheader()
0219         writer.writerows(rows)
0220
0221
0222 def plot_maps(maps, scale, args, output):
0223     """Plot already accumulated counts; no rebinning or integer rate rounding."""
0224     pdfs = []
0225     for source in SOURCE_ORDER:
0226         populated = [s for s in maps.items() if source in s[1]["counts"]
0227                      and s[1]["counts"][source].sum() >= args.min_surface_hits]
0228         if not populated:
0229             continue
0230         filename = f"{args.label}_{source.replace(' ', '_')}_svt_hit_rate.pdf"
0231         with PdfPages(output / filename) as pdf:
0232             for name, surface in populated:
0233                 rates = surface["counts"][source] * scale
0234                 fig, ax = plt.subplots(figsize=(9, 8), constrained_layout=True)
0235                 image = ax.pcolormesh(surface["xedges"], surface["yedges"],
0236                                      np.ma.masked_where(rates.T == 0, rates.T), shading="flat")
0237                 fig.colorbar(image, ax=ax, label="Hit rate per bin [kHz]")
0238                 ax.set_aspect("equal")
0239                 ax.set_xlabel("z [mm]" if surface["barrel"] else "x [mm]")
0240                 ax.set_ylabel("rφ [mm]" if surface["barrel"] else "y [mm]")
0241                 ax.set_title(f"{name} | {source} | {args.label}")
0242                 bin_label = "tile" if args.tile else "RSU"
0243                 ax.text(0.02, 0.98, f"Total: {rates.sum():.4g} kHz\nMaximum per {bin_label}: {rates.max():.4g} kHz",
0244                         transform=ax.transAxes, va="top", bbox={"facecolor": "white", "alpha": 0.85})
0245                 pdf.savefig(fig)
0246                 plt.close(fig)
0247         pdfs.append(filename)
0248     return pdfs
0249
0250
0251 def save_maps(maps, output, label):
0252     """Archive raw histograms and mm bin edges with an explicit JSON key index."""
0253     arrays, index = {}, []
0254     for surface_id, (name, surface) in enumerate(maps.items()):
0255         prefix = f"surface_{surface_id:02d}"
0256         arrays[prefix + "_xedges_mm"] = surface["xedges"]
0257         arrays[prefix + "_yedges_mm"] = surface["yedges"]
0258         for source in SOURCE_ORDER:
0259             if source not in surface["counts"]:
0260                 continue
0261             key = prefix + "_" + source.replace(" ", "_")
0262             arrays[key] = surface["counts"][source]
0263             index.append({"key": key, "lname": name, "source": source,
0264                           "xedges_key": prefix + "_xedges_mm", "yedges_key": prefix + "_yedges_mm"})
0265     np.savez_compressed(output / f"{label}_maps.npz", **arrays)
0266     return index
0267
0268
0269 def main():
0270     args = parse_args()
0271     paths = build_inputs(args)
0272     output = args.output_dir / args.label
0273     if output.exists() and any(output.iterdir()):
0274         raise FileExistsError(f"Output directory is not empty: {output}. Use a new --label to retain earlier results.")
0275     # Create output only after processing succeeds; failed reads leave no partial tables.
0276     maps, inputs, diagnostics, events, elapsed = process_files(paths, args)
0277     output.mkdir(parents=True, exist_ok=True)
0278     scale = 1000.0 / (args.frame_us * events)
0279     metadata = {"label": args.label, "processed_events": events, "frame_us": args.frame_us,
0280                 "bin_mode": "tile" if args.tile else "RSU", "rate_scale_khz_per_count": scale}
0281     write_csv(output / f"{args.label}_summary.csv", summarize(maps, scale), SUMMARY_COLUMNS, metadata)
0282     write_csv(output / f"{args.label}_input_files.csv", inputs, list(inputs[0]), metadata)
0283     map_index = save_maps(maps, output, args.label)
0284     pdfs = plot_maps(maps, scale, args, output)
0285     record = {
0286         **metadata, "created_utc": datetime.now(timezone.utc).isoformat(),
0287         "processing_seconds": elapsed, "input_files": inputs,
0288         "arguments": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()},
0289         "versions": {"numpy": np.__version__, "awkward": ak.__version__, "uproot": uproot.__version__, "matplotlib": matplotlib.__version__},
0290         "code_sha256": {name: hashlib.sha256((HERE / name).read_bytes()).hexdigest()
0291                         for name in ("plot_hit_rate_per_source.py", "svt_sources.py")},
0292         "status_to_source": STATUS_TO_SOURCE, "relations": RELATIONS,
0293         "surface_windows_mm": [{"lname": n, "rec_branch": b, "is_barrel": barrel, "range": window}
0294                                for n, b, barrel, window in SURFACES],
0295         "diagnostics": {branch: dict(counts) for branch, counts in diagnostics.items()},
0296         "maps": map_index, "pdfs": pdfs,
0297     }
0298     (output / f"{args.label}_run.json").write_text(json.dumps(record, indent=2) + "\n")
0299     print(f"Wrote {output}: {events} events, {scale:.6g} kHz/count, processing {elapsed:.2f} s", flush=True)
0300
0301
0302 if __name__ == "__main__":
0303     main()