MLXIDDemo/scripts/05_infer_measurement.py

"""
Run the pretrained 12 keV models on real measurement data living on group
storage (not shipped with the demo -- the full target is tens of GB). By
default this only reads one chunk out of sixteen, capped at a small number
of events per chunk, so it runs quickly; widen with --n-chunks /
--max-events-per-chunk for a fuller (slower) reconstruction.

Output mirrors DeepLearning/Infer_{1,2,3}Photon.py's accumulate_hits() /
save_results(): a super-resolution reconstructed frame, a per-pixel
hit-count frame, and a sub-pixel position histogram, each saved as both
.npy and .png. There's no ground truth for real data, so this is
qualitative -- the knife-edge target should show a sharp edge in the
super-resolution frame, the flat-field target should show a uniform
sub-pixel histogram.

If group storage isn't reachable (e.g. running outside this machine), this
falls back to the small illustrative sample shipped in
data/measurement_samples/ (~30k-65k events, one pseudo-chunk only --
built by scripts/prepare_measurement_fallback.py). --n-chunks /
--max-events-per-chunk are capped accordingly in that case; the
reconstruction will look noticeably noisier than with real group-storage
access.

Usage: python 05_infer_measurement.py --photons 1 --target Edge2Filters_12keV
       python 05_infer_measurement.py --photons 2 --target Flat2Filters_12keV --n-chunks 4

Needs: torch, numpy, h5py + CUDA (double/triple-photon models). Use the
`mlxid_demo` conda env.
"""
import sys
import argparse
from pathlib import Path
import numpy as np
import torch
import h5py
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

HERE = Path(__file__).resolve().parent.parent
sys.path.append(str(HERE / 'src'))
from model_zoo import load_pretrained

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

SOURCE_DIR = Path('/mnt/sls_det_storage/moench_data/MLXID/Samples/Measurement')
FALLBACK_DIR = HERE / 'data/measurement_samples'  ### small local sample, see prepare_measurement_fallback.py
TARGETS = {
    'Edge2Filters_12keV': dict(folder='2603MaxIV_Edge2Filters_12keV', NX=101, NY=101, roi=[0, 101, 0, 101]),
    'Flat2Filters_12keV': dict(folder='2603MaxIV_Flat2Filters_12keV', NX=101, NY=101, roi=[0, 101, 0, 101]),
}
CLUSTER_PREFIX = {1: '1Photon_CS3', 2: '2Photon_CS6', 3: '3Photon_CS9'}  ### CS must match each model's input size


def resolve_base_dir(target):
    """Group storage if reachable, else the small locally-shipped fallback sample."""
    group_folder = SOURCE_DIR / TARGETS[target]['folder']
    try:
        if group_folder.is_dir() and any(group_folder.iterdir()):
            return group_folder, True
    except OSError:
        pass
    print(f'  No access to {group_folder} -- falling back to data/measurement_samples/{target}/')
    return FALLBACK_DIR / target, False


def load_clusters(target, n_photons, n_chunks, max_events_per_chunk):
    folder, using_group_storage = resolve_base_dir(target)
    if not using_group_storage:
        n_chunks = 1  ### the fallback sample only ships one pseudo-chunk per photon count
    all_clusters, all_refs = [], []
    for i in range(n_chunks):
        f_path = folder / f'{CLUSTER_PREFIX[n_photons]}_chunk{i}.h5'
        if not f_path.exists():
            print(f'  SKIP (not found): {f_path}')
            continue
        with h5py.File(f_path, 'r') as f:
            n = min(max_events_per_chunk, f['clusters'].shape[0])
            all_clusters.append(f['clusters'][:n])
            all_refs.append(f['referencePoint'][:n])
        print(f'  loaded {n:,} events from {f_path.name}')
    return np.concatenate(all_clusters, axis=0), np.concatenate(all_refs, axis=0)


def run_inference(model, clusters, n_photons, batch_size=8192):
    """clusters: (N, size, size) raw pixel arrays -> (M, 2) predicted sub-pixel positions
    (within the crop), M = N for 1-photon, N*n_photons for 2/3-photon (flattened)."""
    size = clusters.shape[-1]
    if n_photons == 1:
        ### singlePhotonNet expects the (x, y) coordinate grid concatenated as extra channels
        x = np.linspace(-0.5, 0.5, size)
        x_grid, y_grid = np.meshgrid(x, x, indexing='ij')
        x_grid = torch.tensor(x_grid).float().to(DEVICE)
        y_grid = torch.tensor(y_grid).float().to(DEVICE)

    all_preds = []
    with torch.no_grad():
        for i in range(0, len(clusters), batch_size):
            batch = torch.from_numpy(clusters[i:i + batch_size]).float().unsqueeze(1).to(DEVICE)  # (b,1,size,size)
            if n_photons == 1:
                b = batch.size(0)
                inp = torch.cat([batch, x_grid.expand(b, 1, size, size), y_grid.expand(b, 1, size, size)], dim=1)
                pred = model(inp)[:, :2]
            else:
                pred = model(batch).view(-1, 2)  ### flatten the n_photons dimension
            all_preds.append(pred.cpu())
    preds = torch.cat(all_preds, dim=0).numpy()
    preds += size / 2.  ### shift from centered-at-0 (training convention) back to position-within-crop
    return preds


def accumulate_hits(predictions, reference_points, NX, NY, binning_factor=10):
    ml_super_frame = np.zeros((NY * binning_factor, NX * binning_factor), dtype=np.int64)
    count_frame = np.zeros((NY, NX), dtype=np.int64)
    subpixel_dist = np.zeros((binning_factor, binning_factor), dtype=np.int64)

    absolute = predictions + reference_points

    hit_x = np.floor(absolute[:, 0] * binning_factor).astype(int)
    hit_y = np.floor(absolute[:, 1] * binning_factor).astype(int)
    valid = (hit_x >= 0) & (hit_x < NX * binning_factor) & (hit_y >= 0) & (hit_y < NY * binning_factor)
    np.add.at(ml_super_frame, (hit_y[valid], hit_x[valid]), 1)

    ref_x = (reference_points[:, 0] + 1).astype(int)  ### reference point is the crop's lower-left corner
    ref_y = (reference_points[:, 1] + 1).astype(int)
    valid_ref = (ref_x >= 0) & (ref_x < NX) & (ref_y >= 0) & (ref_y < NY)
    np.add.at(count_frame, (ref_y[valid_ref], ref_x[valid_ref]), 1)

    sub_x = np.floor((absolute[:, 0] % 1) * binning_factor).astype(int)
    sub_y = np.floor((absolute[:, 1] % 1) * binning_factor).astype(int)
    np.add.at(subpixel_dist, (sub_y, sub_x), 1)

    return ml_super_frame, count_frame, subpixel_dist


def save_results(ml_super_frame, count_frame, subpixel_dist, roi, binning_factor, n_photons, output_dir):
    output_dir.mkdir(parents=True, exist_ok=True)
    x_st, x_ed, y_st, y_ed = roi
    prefix = f'{n_photons}Photon'

    np.save(output_dir / f'{prefix}_ML_superFrame.npy', ml_super_frame)
    plt.figure(figsize=(8, 8))
    plt.imshow(ml_super_frame[y_st * binning_factor:y_ed * binning_factor, x_st * binning_factor:x_ed * binning_factor],
               origin='lower', extent=[x_st, x_ed, y_st, y_ed])
    plt.colorbar(label='Counts')
    plt.title('ML Super-Resolution Frame')
    plt.xlabel('X (pixel)')
    plt.ylabel('Y (pixel)')
    plt.savefig(output_dir / f'{prefix}_ML_superFrame.png', dpi=300, bbox_inches='tight')
    plt.close()

    np.save(output_dir / f'{prefix}_count_Frame.npy', count_frame)
    plt.figure(figsize=(8, 8))
    plt.imshow(count_frame[y_st:y_ed, x_st:x_ed], origin='lower', extent=[x_st, x_ed, y_st, y_ed])
    plt.colorbar(label='Counts')
    plt.title('Photon Count Frame')
    plt.xlabel('X (pixel)')
    plt.ylabel('Y (pixel)')
    plt.savefig(output_dir / f'{prefix}_count_Frame.png', dpi=300, bbox_inches='tight')
    plt.close()

    np.save(output_dir / f'{prefix}_subpixel_Distribution.npy', subpixel_dist)
    plt.figure(figsize=(8, 8))
    plt.imshow(subpixel_dist, origin='lower', extent=[0, 1, 0, 1])
    plt.colorbar(label='Counts')
    plt.title('Sub-pixel Distribution')
    plt.xlabel('Sub-pixel X')
    plt.ylabel('Sub-pixel Y')
    plt.savefig(output_dir / f'{prefix}_subpixel_Distribution.png', dpi=300, bbox_inches='tight')
    plt.close()
    rms, mean = np.std(subpixel_dist), np.mean(subpixel_dist)
    print(f'[Plotting]: Sub-pixel distribution RMS/mean: {rms / mean:.4f}, '
          f'expected {1 / np.sqrt(mean):.4f} for a uniform distribution')
    print(f'[Plotting]: Results saved to: {output_dir}')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--photons', type=int, choices=[1, 2, 3], required=True)
    parser.add_argument('--target', choices=list(TARGETS), required=True)
    parser.add_argument('--n-chunks', type=int, default=1, help='out of 16 available per target')
    parser.add_argument('--max-events-per-chunk', type=int, default=50_000)
    parser.add_argument('--binning-factor', type=int, default=10)
    args = parser.parse_args()

    print(f'[05] Using device: {DEVICE}')
    print(f'[05] Loading {args.target}, up to {args.n_chunks} chunk(s) '
          f'(out of 16 if group storage is reachable), {args.max_events_per_chunk:,} events each...')
    clusters, ref_points = load_clusters(args.target, args.photons, args.n_chunks, args.max_events_per_chunk)
    print(f'[05] {clusters.shape[0]:,} clusters loaded, shape {clusters.shape[1:]}')

    model = load_pretrained(args.photons, DEVICE)
    predictions = run_inference(model, clusters, args.photons)

    if args.photons > 1:
        ref_points = np.repeat(ref_points, args.photons, axis=0)  ### align with the flattened per-photon predictions

    target_conf = TARGETS[args.target]
    ml_super_frame, count_frame, subpixel_dist = accumulate_hits(
        predictions, ref_points, NX=target_conf['NX'], NY=target_conf['NY'], binning_factor=args.binning_factor)

    output_dir = HERE / 'InferenceResults' / args.target / f'{args.photons}ph'
    save_results(ml_super_frame, count_frame, subpixel_dist, target_conf['roi'],
                 args.binning_factor, args.photons, output_dir)


if __name__ == '__main__':
    main()