#!/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)