#!/usr/bin/env python3 """ WP1.2 -- generate grid-based MC frames from the existing single-photon stamps. python scripts/02_generate_frames.py --config Configs/generate_frames_12keV.yaml python scripts/02_generate_frames.py --config ... --splits train val python scripts/02_generate_frames.py --config ... --splits eval_scan python scripts/02_generate_frames.py --config ... --smoke # tiny, fast Each shard is one HDF5 file: frames (N, S, S) float32 keV photons (M, 4) float32 x, y, E[keV], z[um] (frame coordinates) frame_ptr (N+1,) int64 photons of frame i = photons[ptr[i]:ptr[i+1]] occupancy (N,) float32 the occupancy each frame was thrown at """ import argparse import os import sys import time from multiprocessing import Pool import h5py import numpy as np from omegaconf import OmegaConf sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src")) from frame_generator import generate_shard # noqa: E402 SMOKE_FRAMES = 400 def stamp_path(conf, index): return os.path.join(conf.stamps.folder, conf.stamps.pattern.format(index)) def run_job(job): started = time.time() frames, photons, frame_pointer, occupancy, stats = generate_shard( stamp_files=job["stamp_files"], n_frames=job["n_frames"], frame_size=job["frame_size"], occupancy_spec=job["occupancy"], seed=job["seed"], margin=job["margin"], noise_keV=job["noise_keV"], cm_row_keV=job["cm_row_keV"], cm_col_keV=job["cm_col_keV"], edge_fraction=job["edge_fraction"], mode=job["mode"], lut_path=job.get("lut_path"), energy_spec=job.get("energy_spec"), pixel_size_um=job["pixel_size_um"], depletion_voltage=job.get("depletion_voltage"), temperature=job.get("temperature"), ) os.makedirs(os.path.dirname(job["output"]), exist_ok=True) with h5py.File(job["output"], "w") as handle: handle.create_dataset("frames", data=frames, dtype="f4") handle.create_dataset("photons", data=photons, dtype="f4") handle.create_dataset("frame_ptr", data=frame_pointer, dtype="i8") handle.create_dataset("occupancy", data=occupancy, dtype="f4") for key, value in job["attrs"].items(): handle.attrs[key] = value for key, value in stats.items(): handle.attrs[key] = value return dict(output=job["output"], seconds=time.time() - started, **stats) def build_jobs(conf, splits, smoke=False): jobs = [] base_seed = int(conf.runtime.seed) generation = conf.get("generation", {}) mode = str(generation.get("mode", "stamps")) common = dict( frame_size=int(conf.frame.size), margin=int(conf.frame.margin), noise_keV=float(conf.noise.pixel_keV), cm_row_keV=float(conf.noise.common_mode_row_keV), cm_col_keV=float(conf.noise.common_mode_col_keV), mode=mode, pixel_size_um=float(conf.meta.pixel_size_um), ) if mode == "lut": common.update(lut_path=str(generation.lut), depletion_voltage=float(conf.meta.depletion_voltage), temperature=float(conf.meta.temperature)) attrs = dict(energy_keV=float(conf.meta.get("energy", 0.0)), sensor=str(conf.meta.sensor), hv=int(conf.meta.hv), pixel_size_um=float(conf.meta.pixel_size_um), frame_size=int(conf.frame.size), margin=int(conf.frame.margin), noise_keV=float(conf.noise.pixel_keV)) for split in splits: if split == "eval_scan": scan = conf.eval_scan stamp_files = ([] if mode == "lut" else [stamp_path(conf, index) for index in scan.stamp_files]) n_frames = SMOKE_FRAMES if smoke else int(scan.frames_each) for point_index, occupancy in enumerate(scan.occupancies): tag = f"occ{float(occupancy) * 100:g}pct".replace(".", "p") jobs.append(dict( output=os.path.join(conf.output.folder, "eval_scan", f"{tag}.h5"), stamp_files=([] if mode == "lut" else [stamp_files[point_index % len(stamp_files)]]), energy_spec=(OmegaConf.to_container(conf.spectrum["eval"], resolve=True) if mode == "lut" else None), n_frames=n_frames, occupancy={"mode": "fixed", "value": float(occupancy)}, seed=base_seed + 9000 + point_index, edge_fraction=float(scan.edge_fraction), attrs=dict(attrs, split="eval_scan", occupancy_nominal=float(occupancy)), **common)) continue split_conf = conf.splits[split] if mode == "lut": first_file, last_file = 0, int(split_conf.n_shards) - 1 else: first_file, last_file = conf.stamps[f"{split}_files"] n_frames = SMOKE_FRAMES if smoke else int(split_conf.frames_per_shard) for shard_index, file_index in enumerate(range(int(first_file), int(last_file) + 1)): jobs.append(dict( output=os.path.join(conf.output.folder, split, f"shard_{shard_index:02d}.h5"), stamp_files=([] if mode == "lut" else [stamp_path(conf, file_index)]), energy_spec=(OmegaConf.to_container(conf.spectrum[split], resolve=True) if mode == "lut" else None), n_frames=n_frames, occupancy=OmegaConf.to_container(split_conf.occupancy, resolve=True), seed=base_seed + 1000 * (splits.index(split) + 1) + shard_index, edge_fraction=float(conf.edge_augmentation.fraction), attrs=dict(attrs, split=split, mode=mode), **common)) return jobs def main(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True) parser.add_argument("--splits", nargs="+", default=["train", "val", "test", "eval_scan"]) parser.add_argument("--smoke", action="store_true", help=f"{SMOKE_FRAMES} frames per shard, for a quick check") parser.add_argument("--workers", type=int, default=None) args = parser.parse_args() conf = OmegaConf.load(args.config) jobs = build_jobs(conf, list(args.splits), smoke=args.smoke) n_workers = args.workers or int(conf.runtime.n_workers) print(f"[gen] {len(jobs)} shards, {n_workers} workers, " f"frame {conf.frame.size}x{conf.frame.size}") print(f"[gen] output -> {conf.output.folder}") started = time.time() with Pool(processes=min(n_workers, len(jobs))) as pool: results = pool.map(run_job, jobs) total_photons = sum(r["n_photons"] for r in results) total_frames = sum(r["n_frames"] for r in results) total_same_cell = sum(r["n_same_cell"] for r in results) for result in sorted(results, key=lambda r: r["output"]): print(f" {os.path.relpath(result['output'], conf.output.folder):32s} " f"{result['n_frames']:7d} frames {result['n_photons']:9d} photons " f"={result['mean_occupancy']*100:5.2f}% " f"same-cell={result['n_same_cell']:6d} ({result['seconds']:.1f}s)") print(f"[gen] total {total_frames} frames, {total_photons} photons, " f"same-cell loss {total_same_cell} " f"({100.0*total_same_cell/max(total_photons,1):.2f}%), " f"{time.time()-started:.1f}s") if __name__ == "__main__": main()