#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
build_overlays_and_list.py
From image_run_log.csv (latest run only), create per-source overlay HDF5 files
(named with the next run number) and per-(source,event) lst files:
run_{next}/event_{EV6}_{HASH8}/lst_{EV6}_{HASH8}.lst -> "<overlay.h5> //EV"
Key fixes vs single-H5 version:
- Disambiguate same event ids coming from different .h5 sources by appending a
short hash of the absolute source path to the event folder and lst file names.
- Keep one overlay per source; write all proposed shifts for that source into it.
- Persist overlay→original maps (JSON + TSV) for traceability.
Notes:
- Skips entries whose next_* are "done".
- Overlays are created in runs/run_{next:03d}/
- Each overlay filename is <src_basename>_overlay_{next:03d}.h5
Requires overlay_elink.py (create_overlay, write_shifts_mm).
"""
from __future__ import annotations
from coseda.logging_utils import log_print
import argparse, os, sys, math, json, hashlib
from pathlib import Path
DEFAULT_ROOT = ""
try:
from coseda.ici.overlay_elink import create_overlay, write_shifts_mm
except Exception:
log_print(
"ERROR: Could not import coseda.ici.overlay_elink (create_overlay, write_shifts_mm). "
"Ensure COSEDA is on PYTHONPATH.",
file=sys.stderr,
)
raise
[docs]
def load_state(state_path: str) -> dict:
"""
Load the JSON sidecar produced by propose_next_shifts.py.
Expected structure:
{
"last_global_run": int,
"events": {
"<abs_h5_path>::<event_id>": {
"trials": [...],
"latest_status": [next_dx_str, next_dy_str],
"last_run": int
},
...
}
}
"""
try:
with open(state_path, "r", encoding="utf-8") as f:
state = json.load(f)
if not isinstance(state, dict):
raise ValueError("state not a dict")
if "events" not in state:
state["events"] = {}
if "last_global_run" not in state:
state["last_global_run"] = -1
return state
except FileNotFoundError:
# No state yet
return {"last_global_run": -1, "events": {}}
except Exception:
# Corrupt state → fall back to empty
return {"last_global_run": -1, "events": {}}
def _abs(p: str) -> str:
return os.path.abspath(os.path.expanduser(p))
def _short_hash(s: str, n: int = 8) -> str:
return hashlib.blake2s(s.encode("utf-8"), digest_size=8).hexdigest()[:n]
[docs]
def parse_log(log_path: str):
"""
Returns:
entries: list of tuples (type, payload)
type == 'section' -> payload = (real_path, event:int)
type == 'row' -> payload = (run_n:int, dx:str, dy:str, idx:str, wr:str, ndx:str, ndy:str)
latest_run: int
"""
entries = []
latest_run = -1
with open(log_path, "r", encoding="utf-8") as f:
_ = f.readline() # header
for ln in f:
if ln.startswith("#/"):
# header for a (source,event) section
try:
path_part, ev_part = ln[1:].rsplit(" event ", 1)
ev = int(ev_part.strip())
entries.append(("section", (_abs(path_part.strip()), ev)))
except Exception:
pass
continue
parts = [p.strip() for p in ln.rstrip("\n").split(",")]
while len(parts) < 7:
parts.append("")
run_s, dx_s, dy_s, idx_s, wr_s, ndx_s, ndy_s = parts[:7]
try:
run_n = int(run_s)
if run_n > latest_run:
latest_run = run_n
except Exception:
continue
entries.append(("row", (run_n, dx_s, dy_s, idx_s, wr_s, ndx_s, ndy_s)))
return entries, latest_run
[docs]
def collect_latest_numeric_proposals(entries, latest_run: int):
"""
From latest run only, collect numeric proposals per (real_path, event).
Returns mapping: real_path -> { event_idx: (next_dx_float, next_dy_float) }
"""
by_src = {}
current_key = None
last_next = {} # (src,event) -> (ndx_str, ndy_str)
for typ, payload in entries:
if typ == "section":
current_key = payload # (real_path, event)
continue
if typ != "row" or current_key is None:
continue
run_n, _dx_s, _dy_s, _idx_s, _wr_s, ndx_s, ndy_s = payload
if run_n != latest_run:
continue
# keep last row's proposal in this section
last_next[current_key] = (ndx_s, ndy_s)
# Convert strings to floats and filter "done" / blanks / non-finite
for (rp, ev), (ndx_s, ndy_s) in last_next.items():
if not ndx_s or not ndy_s:
continue
if ndx_s == "done" or ndy_s == "done":
continue
try:
ndx = float(ndx_s); ndy = float(ndy_s)
if not (math.isfinite(ndx) and math.isfinite(ndy)):
continue
except Exception:
continue
by_src.setdefault(rp, {})[int(ev)] = (ndx, ndy)
return by_src
[docs]
def tag_overlay_with_original(overlay_path: str, original_path: str) -> None:
"""Optionally store the original path as an attribute in the overlay HDF5 file."""
try:
import h5py
except Exception:
return
try:
with h5py.File(overlay_path, "r+") as h5:
h5.attrs["overlay_original_path"] = original_path
except Exception:
pass # non-fatal
[docs]
def ensure_overlay_for_run(src_path: str, run_dir: str, next_run: int) -> str:
"""
Create (or overwrite) an overlay for src_path under run_dir.
Name: <basename>_overlay_{next:03d}.h5
Returns absolute overlay path.
"""
os.makedirs(run_dir, exist_ok=True)
base = os.path.basename(src_path)
root, _ = os.path.splitext(base)
overlay_name = f"{root}_overlay_{next_run:03d}.h5"
overlay_path = os.path.join(run_dir, overlay_name)
create_overlay(src_path, overlay_path)
tag_overlay_with_original(overlay_path, _abs(src_path))
return _abs(overlay_path)
[docs]
def write_all_shifts(run_dir: str, next_run: int, proposals_by_src: dict):
"""
For each source, write its proposed shifts to its overlay in run_dir.
Returns: dict src_path -> overlay_path
"""
overlay_paths = {}
for src_path, ev2shift in proposals_by_src.items():
if not ev2shift:
continue
overlay_path = ensure_overlay_for_run(src_path, run_dir, next_run)
overlay_paths[src_path] = overlay_path
# Prepare vectors (CrystFEL overlay writer wants parallel arrays)
indices = sorted(ev2shift.keys())
dx = [ev2shift[i][0] for i in indices]
dy = [ev2shift[i][1] for i in indices]
write_shifts_mm(overlay_path, indices, dx, dy)
return overlay_paths
[docs]
def write_event_lists(run_dir: str, proposals_by_src: dict, overlay_paths: dict) -> int:
"""
Create unique per-(source,event) folders + one-line .lst:
run_dir/event_{EV6}_{H8}/lst_{EV6}_{H8}.lst -> "<overlay.h5> //EV"
Returns number of lists written.
"""
n = 0
Path(run_dir).mkdir(parents=True, exist_ok=True)
for src_path, ev2shift in sorted(proposals_by_src.items()):
if not ev2shift:
continue
ov = overlay_paths.get(src_path)
if not ov:
continue
h8 = _short_hash(_abs(src_path), 8)
for ev in sorted(ev2shift.keys()):
ev6 = f"{int(ev):06d}"
ev_dir = os.path.join(run_dir, f"event_{ev6}_{h8}")
Path(ev_dir).mkdir(parents=True, exist_ok=True)
lst_path = os.path.join(ev_dir, f"lst_{ev6}_{h8}.lst")
with open(lst_path, "w", encoding="utf-8") as f:
f.write(f"{ov} //{int(ev)}\n")
n += 1
return n
[docs]
def write_overlay_mapping(run_dir: str, overlay_paths: dict) -> str:
"""
Persist mapping files relating overlays to originals within the run folder.
Writes:
- overlay_to_original.json : {"/abs/overlay.h5": "/abs/original.h5", ...}
- overlay_to_original.tsv : tab-delimited table for quick inspection
Returns the path to the JSON file.
"""
os.makedirs(run_dir, exist_ok=True)
map_json = os.path.join(run_dir, "overlay_to_original.json")
map_tsv = os.path.join(run_dir, "overlay_to_original.tsv")
overlay_to_src = { _abs(ov): _abs(src) for src, ov in overlay_paths.items() }
with open(map_json, "w", encoding="utf-8") as f:
json.dump(overlay_to_src, f, indent=2)
with open(map_tsv, "w", encoding="utf-8") as f:
f.write("# overlay_path\toriginal_path\n")
for ov, src in sorted(overlay_to_src.items()):
f.write(f"{ov}\t{src}\n")
return map_json
[docs]
def main(argv=None) -> int:
ap = argparse.ArgumentParser(
description=(
"Create run-named overlays and unique per-(source,event) lst folders "
"from latest proposals stored in image_run_state.json. Also writes "
"overlay→original mapping files."
)
)
ap.add_argument("--run-root", default=None,
help="Path to run root that contains 'runs/'. Defaults to DEFAULT_ROOT if omitted.")
args = ap.parse_args(argv)
run_root = _abs(args.run_root or DEFAULT_ROOT)
# --- Pull latest proposals from sidecar state rather than scanning CSV ---
state_path = os.path.join(run_root, "image_run_state.json")
state = load_state(state_path)
events_state = state.get("events", {})
latest_run = int(state.get("last_global_run", -1))
if latest_run < 0 or not events_state:
log_print("[overlay] No state or no runs found in image_run_state.json; nothing to do.")
return 0
# Build mapping: src_path -> { event_idx: (ndx, ndy) } using latest_status
proposals_by_src: dict[str, dict[int, tuple[float, float]]] = {}
for key, ev_state in events_state.items():
# key format: "<abs_h5_path>::<event_id>"
try:
src_path, ev_str = key.rsplit("::", 1)
ev = int(ev_str)
except Exception:
continue
latest_status = ev_state.get("latest_status", ["", ""])
if not isinstance(latest_status, (list, tuple)) or len(latest_status) < 2:
continue
ndx_s, ndy_s = latest_status
# Skip events that are explicitly done or have no numeric proposal
if (not ndx_s) or (not ndy_s):
continue
if str(ndx_s).lower() == "done" or str(ndy_s).lower() == "done":
continue
try:
ndx = float(ndx_s)
ndy = float(ndy_s)
if not (math.isfinite(ndx) and math.isfinite(ndy)):
continue
except Exception:
continue
proposals_by_src.setdefault(_abs(src_path), {})[int(ev)] = (ndx, ndy)
if not proposals_by_src:
log_print(f"[overlay] No numeric proposals for latest run={latest_run} in state; nothing to do.")
return 0
# Next run index is latest_run + 1
next_run = latest_run + 1
next_run_dir = os.path.join(run_root, f"run_{next_run:03d}")
os.makedirs(next_run_dir, exist_ok=True)
# 1) Per-source overlays (one per HDF5 origin)
overlay_paths = write_all_shifts(next_run_dir, next_run, proposals_by_src)
# 2) Unique per-(source,event) lsts under event_{EV6}_{HASH8}/
_ = write_event_lists(next_run_dir, proposals_by_src, overlay_paths)
# 3) Persist overlay→original mapping (JSON + TSV)
_ = write_overlay_mapping(next_run_dir, overlay_paths)
return 0
if __name__ == "__main__":
sys.exit(main())