Initial commit: grid-based multi-photon reconstruction

Code, configs, charge LUTs, best.pth and evaluation records for each
training run, the knife-edge / Siemens-star / XRF notebooks with
outputs, and the overnight run reports. Epoch checkpoints, inference
results and HTML exports stay local (see .gitignore).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

scripts/00_selftest.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""
+Self-tests for the grid formulation.  Run after touching src/.
+
+    python scripts/00_selftest.py
+
+1. encode -> decode round trip: the coordinate convention is self-consistent.
+2. D4 augmentation keeps frame and truth coordinates aligned.
+3. Translation equivariance: a model trained on SxS frames gives *identical*
+   predictions when the same content is embedded in a larger frame -- the claim
+   that D1/D2 rest on.  If this test fails, train-small/infer-large is invalid.
+4. Position-space merging removes duplicates but keeps genuinely distinct
+   photons in adjacent cells -- the property the 3x3 max-pool NMS lacked.
+"""
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE.parent / "src"))
+from datasets import apply_d4, encode_targets   # noqa: E402
+from decode import decode                       # noqa: E402
+from models import GridNet, valid_mask          # noqa: E402
+
+all_passed = True
+
+
+def check(name, condition, detail=""):
+    global all_passed
+    all_passed &= bool(condition)
+    print(f"  [{'PASS' if condition else 'FAIL'}] {name}"
+          f"{'  ' + detail if detail else ''}")
+
+
+print("1. encode -> decode round trip")
+frame_size, energy_scale = 32, 12.0
+rng = np.random.default_rng(0)
+n_photons = 40
+photons = np.zeros((n_photons, 4), np.float32)
+photons[:, 0] = rng.uniform(0, frame_size, n_photons)
+photons[:, 1] = rng.uniform(0, frame_size, n_photons)
+photons[:, 2] = rng.uniform(5, 30, n_photons)
+cells = (photons[:, 1].astype(int) * frame_size + photons[:, 0].astype(int))
+_, unique_index = np.unique(cells, return_index=True)   # drop same-cell duplicates
+photons = photons[np.sort(unique_index)]
+
+objectness, offset, energy = encode_targets(photons, frame_size, energy_scale)
+prediction = torch.zeros(1, 4, frame_size, frame_size)
+prediction[0, 0] = torch.from_numpy(
+    np.where(objectness > 0, 10.0, -10.0).astype(np.float32))
+prediction[0, 1:3] = torch.from_numpy(offset)
+prediction[0, 3:4] = torch.from_numpy(energy)
+
+points = decode(prediction, threshold=0.5, merge_distance=0.0, halo=0,
+                energy_scale=energy_scale)[0].numpy()
+check("photon count", len(points) == len(photons), f"{len(points)} vs {len(photons)}")
+order_points = np.lexsort((points[:, 0], points[:, 1]))
+order_truth = np.lexsort((photons[:, 0], photons[:, 1]))
+position_error = np.abs(points[order_points, :2] - photons[order_truth, :2]).max()
+energy_error = np.abs(points[order_points, 2] - photons[order_truth, 2]).max()
+check("position round trip", position_error < 1e-4,
+      f"max |dx| = {position_error:.2e} px")
+check("energy round trip", energy_error < 1e-3,
+      f"max |dE| = {energy_error:.2e} keV")
+
+print("2. D4 augmentation consistency")
+frame = np.zeros((frame_size, frame_size), np.float32)
+for x, y, photon_energy, _ in photons:          # put a delta at each photon
+    frame[int(y), int(x)] += photon_energy
+worst_mismatch = 0.0
+for operation in range(8):
+    transformed_frame, transformed_positions = apply_d4(frame, photons[:, :2], operation)
+    for (x, y), (x0, y0) in zip(transformed_positions, photons[:, :2]):
+        original = frame[int(y0), int(x0)]
+        worst_mismatch = max(worst_mismatch,
+                             abs(transformed_frame[int(y), int(x)] - original))
+check("charge follows the truth under all 8 ops", worst_mismatch < 1e-5,
+      f"max mismatch {worst_mismatch:.2e}")
+
+print("3. translation equivariance (train small -> infer large)")
+torch.manual_seed(0)
+network = GridNet(channels=(32, 64, 64, 64)).eval()
+radius = network.rf_radius
+small_frame = torch.randn(1, 1, frame_size, frame_size) * 0.05
+small_frame[0, 0, 12:15, 9:12] += torch.tensor([[.1, .3, .1],
+                                                [.4, 3.0, .5],
+                                                [.1, .4, .1]])
+large_frame = torch.randn(1, 1, 3 * frame_size, 3 * frame_size) * 0.05
+row_offset, column_offset = 40, 27
+large_frame[0, 0, row_offset:row_offset + frame_size,
+            column_offset:column_offset + frame_size] = small_frame[0, 0]
+with torch.no_grad():
+    small_prediction = network(small_frame)
+    large_prediction = network(large_frame)
+embedded = large_prediction[:, :, row_offset:row_offset + frame_size,
+                            column_offset:column_offset + frame_size]
+mask = valid_mask(frame_size, frame_size, radius)
+inside_difference = (small_prediction - embedded)[:, :, mask].abs().max().item()
+check(f"identical inside the halo (R={radius})", inside_difference < 1e-4,
+      f"max |diff| = {inside_difference:.2e}")
+# R must also be tight: the ring one pixel *inside* the halo must already differ
+inner_ring = valid_mask(frame_size, frame_size, radius - 1) & ~mask
+ring_difference = (small_prediction - embedded)[:, :, inner_ring].abs().max().item()
+check(f"ring at distance R-1={radius-1} already differs "
+      f"(R is tight, not over-conservative)",
+      ring_difference > 1e-6, f"|diff| = {ring_difference:.2e}")
+
+print("4. position-space merging")
+# two adjacent cells pointing at the same physical place -> one photon
+duplicate = torch.full((1, 4, 8, 8), -10.0)
+duplicate[0, 1:3] = 0.0
+duplicate[0, 3] = 1.0
+duplicate[0, 0, 4, 4] = 6.0
+duplicate[0, 1, 4, 4] = 0.45                 # x = 4.95
+duplicate[0, 0, 4, 5] = 5.0
+duplicate[0, 1, 4, 5] = -0.45                # x = 5.05, same point
+merged = decode(duplicate, threshold=0.5, merge_distance=0.3, halo=0,
+                energy_scale=12.0)[0]
+check("duplicate in adjacent cells merged to one", len(merged) == 1,
+      f"got {len(merged)} points")
+
+# two adjacent cells pointing at genuinely different places -> two photons
+distinct = torch.full((1, 4, 8, 8), -10.0)
+distinct[0, 1:3] = 0.0
+distinct[0, 3] = 1.0
+distinct[0, 0, 4, 4] = 6.0
+distinct[0, 1, 4, 4] = -0.3                  # x = 4.2
+distinct[0, 0, 4, 5] = 5.0
+distinct[0, 1, 4, 5] = 0.3                   # x = 5.8, 1.6 px away
+kept = decode(distinct, threshold=0.5, merge_distance=0.3, halo=0,
+              energy_scale=12.0)[0]
+check("distinct photons in adjacent cells both kept", len(kept) == 2,
+      f"got {len(kept)} points  (3x3 max-pool NMS would have given 1)")
+
+print("\n" + ("ALL TESTS PASSED" if all_passed else "SOME TESTS FAILED"))
+sys.exit(0 if all_passed else 1)