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>
177 lines
7.7 KiB
Python
177 lines
7.7 KiB
Python
#!/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"<occ>={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()
|