xie_x1andClaude Opus 5.5
0f9a3f496f
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>
139 lines
6.1 KiB
Python
139 lines
6.1 KiB
Python
#!/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)
|