diff --git a/fix_header_race.md b/fix_header_race.md new file mode 100644 index 0000000..2bf4e0e --- /dev/null +++ b/fix_header_race.md @@ -0,0 +1,306 @@ +# Fix Header Race: Aligned STXM + Ptycho Transition Plan (Revised) + +## Goal +Prevent the header/preemption race that can crash ptychography while keeping STXM and ptycho branches aligned across scan transitions. + +## Decisions Locked In +- Header-preemption completion and end-of-scan completion are both treated as full completion. +- Branch alignment is strict: both branches resume from the same first post-transition batch. +- Data arriving during transition should be preserved where possible. +- Overflow during blocked transition is fail-fast. +- Overflow fail-fast behavior: ControlOp publishes a final error signal first, then raises. +- max_blocked_frames is config-driven (default policy value: 10 x batch_size). + +## Why The Old Draft Was Not Sufficient +- It did not explicitly separate ptycho flush request from ptycho flush execution. +- It released the transition barrier too early. +- It did not define overflow behavior and signaling. +- It did not define a deterministic transition phase model. + +## Transition State Model + +### Shared state (thread-safe) +- **File:** pipeline/pipeline.py +- **Where:** StxmApp shared state initialization +- **Add:** + - transition_blocked_event (threading.Event) + - transition_phase: idle | waiting_quiesce | waiting_flush_exec + - ptycho_accum_flushed (bool) + - ptycho_recon_flushed (bool) + - max_blocked_frames (config value) + - transition_error (optional text for diagnostics) + +**Reason** +- Multiple scheduler worker threads can read/write transition state concurrently. +- Event plus explicit phase prevents ambiguous behavior. + +## Authoritative Transition Sequence + +1. **Header is received** in pipeline/header_io.py + - Validate header. + - Stage pending_geometry. + - Set preempt_requested. + - Set transition_blocked_event. + - Set transition_phase = waiting_quiesce. + - Clear ptycho flush ack flags. + - Emit header token to ControlOp. + +2. **ControlOp receives header** in pipeline/control.py + - Execute STXM-only flush. + - Do not request ptycho flush here. + - Keep transition blocked. + +3. **Ptycho recon sees preempt_requested** in pipeline/ptychography_ops.py + - Save partial result if needed. + - Emit recon_complete. + - Enter quiesced preemption path. + +4. **ControlOp receives recon_complete while phase=waiting_quiesce** + - Request ptycho flush now (request only): + - call ptycho_accum.flush() + - call ptycho_recon.flush() + - Set transition_phase = waiting_flush_exec. + +5. **Actual ptycho flush execution occurs later in operator compute** + - Accumulator: _perform_flush executes at top of next compute if requested. + - Recon: _perform_flush executes at top of next compute if requested. + - Each sets its corresponding ack flag when _perform_flush has actually run. + +6. **Barrier release condition** + - Only release transition when BOTH are true: + - ptycho_accum_flushed + - ptycho_recon_flushed + - Then: + - clear transition_blocked_event + - set transition_phase = idle + - reset ack flags for next transition + +This is the key timing rule: +- Ptycho flush is requested in ControlOp on recon_complete during waiting_quiesce. +- Ptycho flush is executed inside ptycho operators in _perform_flush during their next compute tick. +- Barrier is released only after both executions are acknowledged. + +## Concrete Code Changes + +### 1) Shared transition primitives +- **File:** pipeline/pipeline.py +- **Where:** StxmApp.__init__ and compose wiring +- **Change:** + - Add thread-safe transition fields to shared state. + - Pass shared state into GatherOp and ControlOp. + +### 2) Gather alignment gate and overflow accounting +- **File:** pipeline/data_io.py +- **Where:** GatherOp.__init__, setup, compute +- **Change:** + - Store shared transition state reference. + - Before emit, block when transition_blocked_event is set. + - Keep matched data cached while blocked. + - Track blocked cached frame count. + - Load max_blocked_frames from config/state. + +### 3) Gather fail-fast path +- **File:** pipeline/data_io.py and pipeline/control.py +- **Where:** GatherOp.compute overflow branch + ControlOp input handling +- **Change:** + - When blocked cache exceeds max_blocked_frames: + - set transition_error in shared state + - emit control message, e.g. transition_overflow + - ControlOp handles transition_overflow by: + - publishing final error signal + - raising RuntimeError + +### 4) Header starts transition cleanly +- **File:** pipeline/header_io.py +- **Where:** after pending geometry staging and preempt_requested +- **Change:** + - set transition_blocked_event + - set transition_phase=waiting_quiesce + - clear prior ack flags + - keep current scan_state geometry update behavior + +### 5) Split ControlOp flush ownership +- **File:** pipeline/control.py +- **Where:** constructor + helpers +- **Change:** + - Replace single flushable_ops with stxm_flush_ops and ptycho_flush_ops. + - Add helpers: + - do_stxm_flush() + - request_ptycho_flush() + - do_full_flush() (for non-transition full completion path) + +### 6) Control header branch +- **File:** pipeline/control.py +- **Where:** msg == header +- **Change:** + - STXM-only flush. + - Never request ptycho flush in this branch. + +### 7) Control recon_complete branch (phase-aware) +- **File:** pipeline/control.py +- **Where:** msg == recon_complete +- **Change:** + - If transition_phase == waiting_quiesce: + - request ptycho flush + - set waiting_flush_exec + - do not release barrier + - Else: + - existing full-completion behavior + +### 8) Control flush/start branch hardening +- **File:** pipeline/control.py +- **Where:** msg == flush +- **Change:** + - If transition blocked, do not request ptycho flush from this path. + - Use STXM-only or no-op policy to avoid race reintroduction. + +### 9) Ptycho flush ack on execution +- **File:** pipeline/ptychography_ops.py +- **Where:** + - PtychoAccumulatorOp._perform_flush + - PtychoReconstructionOp._perform_flush +- **Change:** + - Set ack flags when each _perform_flush has actually executed. + +### 10) Barrier release at safe point only +- **File:** pipeline/ptychography_ops.py or pipeline/control.py (single owner chosen) +- **Where:** after both ack flags observed true +- **Change:** + - Release transition barrier only when both ptycho flush executions are confirmed. + - Do not release solely at end of _apply_pending_geometry. + +### 11) Optional extra preemption guard +- **File:** pipeline/ptychography_ops.py +- **Where:** immediately before reconstruction_data/combine launch +- **Change:** + - Add second preempt check to reduce chance of one extra iteration starting. + +## Config Changes +- **File:** pipeline/config_test.yaml and pipeline/config_prod.yaml +- **Add under scheduler or a new transition section:** + - max_blocked_frames: integer +- **Default policy suggestion:** + - max_blocked_frames = 10 x image_src.batch_size (computed if unset) + +## Logging and Observability +- **File:** pipeline/header_io.py + - Log transition start, phase, and header id/shape. +- **File:** pipeline/control.py + - Log phase transitions. + - Log ptycho flush request moment. + - Log final error signal publication before raise. +- **File:** pipeline/ptychography_ops.py + - Log each actual _perform_flush execution and ack set. + - Log barrier release with both ack flags. +- **File:** pipeline/data_io.py + - Log blocked-cache growth and threshold crossing. + +## Expected Outcome +- No ptycho mid-iteration invalid-state reset from header/start control paths. +- Deterministic and explicit timing for ptycho flush request vs execution. +- Strict STXM/ptycho alignment through transition barrier. +- Bounded blocked buffering with explicit fail-fast and error signaling. + +## Notes And Remaining Non-goals +- Multiple rapid header bursts are still out of scope for this pass. +- Best-effort data preservation is targeted, but overflow path intentionally stops the run. + +## PR-sized Implementation Plan + +### PR 1: Control flush split (refactor, behavior-preserving) +- **Purpose:** Separate STXM vs ptycho flush ownership in ControlOp with minimal behavior change. +- **Includes:** + - Split ControlOp operator groups into `stxm_flush_ops` and `ptycho_flush_ops`. + - Add helper methods for STXM-only flush, ptycho-only flush request, and full flush. + - Update compose wiring to pass split groups. +- **Files:** + - pipeline/control.py + - pipeline/pipeline.py +- **Verification:** + - Existing single-scan behavior remains unchanged. + - Existing flush topics still publish as before. + +### PR 2: Transition state primitives + header phase start +- **Purpose:** Introduce thread-safe transition state and start transition on header. +- **Includes:** + - Add transition fields in shared state (`transition_blocked_event`, `transition_phase`, ack flags, `transition_error`). + - Add `max_blocked_frames` state value from config/default policy. + - Header path sets `waiting_quiesce` and blocks transition on valid header when ptycho is enabled. +- **Files:** + - pipeline/pipeline.py + - pipeline/header_io.py +- **Verification:** + - Header during active recon sets blocked event and phase to `waiting_quiesce`. + - No data-path behavior change yet. + +### PR 3: Phase-aware ControlOp handling +- **Purpose:** Make header and recon_complete handling deterministic by phase. +- **Includes:** + - Header branch performs STXM-only flush and never requests ptycho flush. + - recon_complete branch: + - if `waiting_quiesce`: request ptycho flush and move to `waiting_flush_exec` + - else: retain existing full-completion behavior + - flush/start branch hardening while transition is blocked. +- **Files:** + - pipeline/control.py +- **Verification:** + - Header no longer triggers immediate ptycho flush request. + - recon_complete in preemption path triggers ptycho flush request exactly once. + +### PR 4: Ptycho flush execution ack + safe barrier release +- **Purpose:** Release transition only after ptycho flush has actually executed. +- **Includes:** + - Set `ptycho_accum_flushed` in accumulator `_perform_flush`. + - Set `ptycho_recon_flushed` in recon `_perform_flush`. + - Release blocked event only when both ack flags are true. + - Reset phase to `idle` and clear acks for next transition. +- **Files:** + - pipeline/ptychography_ops.py + - pipeline/control.py (if release ownership is centralized) +- **Verification:** + - Logs clearly show: request -> execution ack -> barrier release. + - No early release before both ptycho flush executions. + +### PR 5: Gather alignment gate + fail-fast overflow signaling +- **Purpose:** Enforce strict branch alignment and bounded blocked buffering. +- **Includes:** + - Pass shared state into GatherOp. + - Block Gather emit while transition is blocked; keep matched data cached. + - Track blocked cached frame count. + - Overflow path (`> max_blocked_frames`): + - Gather emits `transition_overflow` control message and sets transition error context. + - ControlOp publishes final error signal first, then raises RuntimeError. +- **Files:** + - pipeline/data_io.py + - pipeline/control.py + - pipeline/pipeline.py + - pipeline/config_test.yaml + - pipeline/config_prod.yaml +- **Verification:** + - Normal transition preserves and resumes aligned batches. + - Forced overflow path publishes final error signal before raise. + +### PR 6: Hardening and observability +- **Purpose:** Improve resilience and diagnosability. +- **Includes:** + - Optional second preemption check just before PIE iteration launch. + - Log phase transitions, flush requests, flush execution acks, and barrier release. + - Tighten comments to document state machine semantics. +- **Files:** + - pipeline/ptychography_ops.py + - pipeline/control.py + - pipeline/header_io.py + - pipeline/data_io.py +- **Verification:** + - Repeated header-during-recon runs show deterministic ordering. + - Logs are sufficient to reconstruct transition timeline end-to-end. + +## Suggested Merge Order +1. PR 1 +2. PR 2 +3. PR 3 +4. PR 4 +5. PR 5 +6. PR 6 + +This order minimizes risk by landing structure first, then phase logic, then barrier correctness, then buffered alignment/fail-fast, then hardening. diff --git a/pipeline/DECISIONS_header_tomo.md b/pipeline/DECISIONS_header_tomo.md new file mode 100644 index 0000000..48b0b94 --- /dev/null +++ b/pipeline/DECISIONS_header_tomo.md @@ -0,0 +1,34 @@ +# Decision log — header / dynamic-geometry / tomography feature + +Design decisions made while implementing the `feat/scan-header-tomo` branch (PR0→PR3), +recorded for review. Dates are when the decision was made. + +## Git / workflow +| # | Decision | Rationale | +|---|----------|-----------| +| 1 | Keep the beamline sim-test tweaks OUT of the feature commits — held in a reversible patch `test_data/sim_code_tweaks.patch` (position mapping `/FMC_IN.VAL1`, `PTYCHO_CENTER` override), applied for testing, reverted before committing. | The tweaks are local test scaffolding (localhost endpoints, fixed scan centre) that must not ship in production code. | +| 2 | Single feature branch `feat/scan-header-tomo` for the whole effort (renamed from `feat/pr1-flush-plumbing`), not stacked per-PR branches. | Simpler to manage/review as one branch. | +| 3 | Dropped the pre-existing `736faf4 "Commit before merge"` via rebase; `Dockerfile_bwell` kept locally (untracked), not pushed. | That commit mixed a Blackwell Dockerfile with debug prints later removed; not part of this feature. | +| 4 | Commit sign-off trailer is `Assisted-By: Claude Opus 4.8 (1M context)`, not `Co-Authored-By:`. | Reflects Claude's assistant role. Applies to all future commits. | + +## Architecture (from the plan `dls-holoscan-header-tomo-plan.md`) +| # | Decision | Rationale | +|---|----------|-----------| +| 5 | Header transport = a **dedicated ZMQ SUB socket** (`header_src`), separate from images/positions. | Keeps the geometry channel independent; always listening. | +| 6 | GPU buffers allocated **once at max capacity** (`max_npoints_h/v`), never realloced at runtime (R-6). A header requesting more frames than capacity is **rejected**. | Runtime cupy realloc under live op references risks fragmentation/leaks and a swap-under-recon race. | +| 7 | Header preemption uses a **quiescence handshake** (R-4): header stages geometry + sets `preempt_requested`; recon finishes the in-flight iteration → saves the partial → signals complete → quiesces → applies the new geometry → re-inits GPU. | The recon holds buffer *views* across a PIE iteration; geometry can only change safely while it's idle. | +| 8 | Flush model = flush-on-completion + a skip-if-clean safety flush at scan start (PR1). | Idempotent; removes stale-buffer bugs without double-flushing. | + +## Tomography (PR3) +| # | Decision | Rationale | +|---|----------|-----------| +| 9 | One `arm`/`start` wraps **all projections**; `num_projections × no_frames` frames stream continuously between a single start/end. Projection boundary is segmented **by frame count**. | Matches the acquisition model; `series_id` (Dectris series counter, one per arm) is shared across all projections. | +| 10 | Each projection **early-stops** as soon as its `no_frames` are accumulated + the in-flight iteration finishes (does NOT run full `total_iterations` per projection). Single-projection scans still run to `total_iterations`. | Throughput: keep up with a continuous multi-projection stream. | +| 11 | **Projection-boundary advance is scoped** (2026-07-02): flush ONLY the accumulator (reset GPU buffer) + recon (reset object/iters) + advance the STXM sink and `current_projection`. **Do NOT flush GatherOp** — its cached next-projection frames must survive. Only the FINAL projection triggers a full scan-end flush (incl. GatherOp). | A full flush at every boundary would drop the next projection's frames piling up in GatherOp's cache (R-5). | +| 12 | Per-projection output files named `{series_id}_proj{NN}.h5` for both STXM and ptycho recon (2026-07-02). Theta omitted from the name for now (≈0 in current test data; easy to add). | All projections share one `series_id`, so the projection index is mandatory to avoid overwrite (M2). | +| 13 | Single-buffer first with a bounded GatherOp cache; add double-buffering only if a load test shows the boundary backlog overflows (PR4). | Avoid paying 2× memory + complexity before evidence it's needed. | + +## Open questions (not blocking — revisit in PR4) +| # | Question | Context | +|---|----------|---------| +| O1 | Should tomography projections get a fixed number of **post-stream refinement iterations** before advancing, instead of stopping the instant all frames arrive? | For tomography (`num_projections > 1`), `is_last = all_data_arrived` (`ptychography_ops.py`), so a projection ends as soon as its `no_frames` are in — in the container test proj-0 got 7 PIE iterations and proj-1 only 4 (its frames were already cached by GatherOp during proj-0, so `all_data_arrived` tripped almost immediately). The single-scan branch grants `post_stream_iterations`; the tomography branch ignores them entirely. On a fast stream this means per-projection reconstructions are coarse. Lever: let each projection run N post-stream iterations after its frames arrive but before advancing, trading throughput for per-projection recon quality. | diff --git a/pipeline/DESIGN_pr4_double_buffer.md b/pipeline/DESIGN_pr4_double_buffer.md new file mode 100644 index 0000000..451f5fd --- /dev/null +++ b/pipeline/DESIGN_pr4_double_buffer.md @@ -0,0 +1,173 @@ +# Design — PR4 double-buffering (projection ping-pong) + +Branch `feat/tomo-pr4-loadtest`. Motivated by moving to a **much faster detector**; +the single-buffer design measured sufficient at 3 kHz on the sim but with shrinking +margin as frame rate / frame size grow. + +## Why (measured, 2026-07-02) + +Load test on daqsim (scan409907, 2 projections, A400). Instrumentation tagged +`PR4-MEASURE` (`data_io.py` GatherOp cache HWM; `ptychography_ops.py` finalize-window +timing + `PR4_ACCUM_QUEUE_CAP` / `PR4_FINALIZE_DELAY_MS` knobs). + +- **Run A (baseline, 500 fps, queue=128):** finalize window **~490 ms** (≈ one PIE + iteration + HDF save). GatherOp cache HWM **256 frames**. Both projections exact. +- **Run B (queue=8, ~2 s injected finalize):** GatherOp cache HWM **336 frames**, + both projections exact, **no drops, no deadlock abort** at 4× the 500 ms timeout. + +**Conclusions:** +1. The current single-buffer design degrades **gracefully** — backpressure is lossless, + and the `stop_on_deadlock` (500 ms) tripwire does **not** fire during a long finalize + because the recon thread is busy (counts as progress). So PR4 does **not** need to + touch the deadlock timeout. +2. Zero loss on the sim is a **sim artifact**: daqsim streams images (PUSH) and positions + from one loop, so pipeline backpressure on the PUSH socket stalls the whole sim, + incidentally pausing positions. **A real detector + PandA are independent free-running + sources** — backpressure during the finalize window can't throttle them, so frames/ + positions land in a full socket and drop (positions especially: PUB/SUB, no flow + control, no CONFLATE, ~1000 RCVHWM). +3. So the single-buffer failure mode against a real fast detector is **source throttling / + silent position loss during the ~1-iteration finalize window**, not a crash. Backlog + ≈ `acq_rate × finalize_window`; the finalize window grows with object/frame size. + +Double-buffering removes the finalize-window backpressure: the accumulator never stops +draining, so a free-running detector is never throttled and positions never back up. + +## How single-buffer works today (baseline) + +One shared GPU buffer set in `ptycho_state`: `raw_gpu (capacity,H,W)`, `positions_full`, +`tilts_full`, one `filled_until` counter. The PtyREX model (`pty_model.scan.positions`, +`pty_data.raw_expanded`) holds **views** into these buffers. + +- `PtychoAccumulatorOp` writes batches into `raw_gpu[filled:new_end]`, advances + `filled_until`. On a projection boundary it splits the straddling batch (head fills the + projection, tail → `self._carry`). +- Once `filled_until >= no_frames` (tomography), the accumulator **backpressures**: + `compute` returns before `receive()` (`ptychography_ops.py:147`). Incoming batches wait + in the `gather → accum` `DOUBLE_BUFFER` queue (capacity 128 ≈ 8k frames). +- `PtychoReconstructionOp` runs the final iteration on the full buffer, saves the + per-projection HDF (~490 ms), emits `projection_complete`. +- `ControlOp` → `accum.advance_projection()` → next accum tick resets `filled_until=0`, + writes the carry, resumes draining. Recon observes `filled < no_frames`, resets object + (probe carried) for the next projection. + +The gap: for the whole finalize window the accumulator is **not draining**, so the source +is backpressured. That is what double-buffering eliminates. + +## Double-buffer design (2-buffer ping-pong) + +### State (`ptycho_state`, allocated once at `capacity` per R-6 — now ×2) +``` +raw_gpu: [bufA, bufB] # two (capacity,H,W) arrays +positions_full: [posA, posB] +tilts_full: [tiltA, tiltB] +filled_until: [fillA, fillB] # per-buffer fill level (under lock) +write_idx: 0 # buffer the accumulator writes (accum owns) +read_idx: 0 # buffer the recon reads (recon owns) +buffer_ready: [Event, Event] # buffer i full & handed to recon +buffer_free: [Event, Event] # buffer i free for the accumulator (both set at init) +``` +GPU cost: 2× `raw_gpu` (~64 MB each at 1024×128×128×f32 → ~128 MB total). Fine on the +A400; scales 2× with frame size — watch on the faster detector. + +### Accumulator (`PtychoAccumulatorOp`) +- Write into `raw_gpu[write_idx]` until `filled_until[write_idx] == no_frames`. +- **On full (tomography): flip instead of backpressure.** + 1. `buffer_ready[write_idx].set()` (hand this projection to the recon). + 2. If `buffer_free[1-write_idx]` is set → flip `write_idx`, `filled_until[write_idx]=0`, + clear `buffer_free[write_idx]`, write the carried straddle-tail, keep draining. + 3. Else (recon still on the other buffer — we've lapped it) → **backpressure `return` + as today.** This is the graceful fallback for sustained slowness: double-buffer buys + exactly one projection of runway, not infinite throughput. +- Single projection (`num_projections == 1`): unchanged — always buffer 0, no flip. + +### Reconstruction (`PtychoReconstructionOp`) +- Read views from `raw_gpu[read_idx]` / `positions_full[read_idx]` (index the existing + view re-point that already runs each `compute`). +- Run iterations; when `filled_until[read_idx] == no_frames` finalize + save the + projection HDF (as today). +- **After finalizing:** `buffer_free[read_idx].set()`, clear `buffer_ready[read_idx]`, + flip `read_idx`, reset object (probe carried) for the next projection. +- Per-projection object reset is the existing advance logic, keyed off the flip. + +### Coordination +`write_idx` is touched only by the accumulator, `read_idx` only by the recon; they observe +each other through `buffer_ready` / `buffer_free` events and per-buffer `filled_until` +(under the existing `ptycho_state["lock"]`). Ping-pong invariant: the accumulator never +writes a buffer whose `buffer_free` is clear (recon still reading it) — enforced by step +2/3 above. + +### Flush / preempt / advance +- **Full flush (scan end / R-2 start safety / PR2 header preempt):** reset BOTH buffers + (`filled_until=[0,0]`, zero both `raw_gpu`), `write_idx=read_idx=0`, `buffer_free` both + set, `buffer_ready` both clear, drop the carry. Fold into `_perform_flush` + + `_reset_for_new_geometry`. +- **Per-projection advance** is now implicit in the flip — the explicit + `advance_projection()` / `projection_complete → control → advance` round-trip can be + simplified or kept as the buffer-free signal. Decision point below. + +## Open design decisions (resolve before implementing) +- **D1 — keep or retire the `projection_complete → ControlOp → advance` round-trip?** + With the flip owning the advance, ControlOp's role shrinks to flush-on-final only. Keep + it for the final-projection flush; the per-projection advance becomes buffer-local. +- **D2 — 2 buffers vs an N-ring.** Start with 2 (one projection of runway, matches the + plan). Parameterize `num_buffers` if a load test later shows one projection isn't enough + headroom for the faster detector. +- **D3 — event objects vs simple int flags under the lock.** Events are clean but add + threading objects to `ptycho_state`; two ints (`ready_mask`, `free_mask`) under the + existing lock may be simpler and match the existing deferred-flag style. +- **D4 — does the STXM path need the same treatment?** `SinkAndPublishOp` saves per + projection but doesn't hold a GPU buffer across a long finalize; it likely doesn't need + double-buffering, but confirm it isn't coupled to the ptycho backpressure via the shared + `gather` output. + +## Testing plan +- **Regression:** the PR3 2-projection tomo test still produces exact 1024-frame + projections + all 4 files (`{series}_proj00/01` × STXM + recon). +- **Ping-pong proof:** instrument the flip (`write_idx`/`read_idx` transitions) and confirm + the accumulator keeps draining (no backpressure `return`) through a boundary while the + recon finalizes — i.e. GatherOp cache HWM stays flat and the accum input queue does not + back up during the finalize window. +- **Lapping fallback:** inject a long finalize (`PR4_FINALIZE_DELAY_MS`) longer than one + projection's accumulation so the recon is lapped; confirm the accumulator falls back to + clean backpressure (no data loss) rather than overwriting a buffer the recon is reading. +- **Preempt + flush:** header preemption mid-tomo resets both buffers correctly. + +## Resolved decisions (as implemented) +- **D1** — kept ControlOp for `recon_complete`/`header`/`flush` only. The per-projection + `projection_complete` round-trip is **gone**: the accumulator self-flips write buffers, + and the **recon owns `current_projection`** + the read-buffer flip (`_flip_read`), which + removes the save-vs-bump filename race. `current_projection` still resets to 0 on a new + header (`header_io.py`). +- **D2** — exactly 2 buffers (`NUM_BUFFERS = 2`), one projection of runway. `num_buffers` + is parameterised so an N-ring is a one-constant change if the faster detector needs more. +- **D3** — int-under-lock coordination (`write_idx`/`read_idx`/`filled_until[]`/`buf_free[]` + under `ptycho_state["lock"]`); no new Event objects. +- **D4** — STXM path untouched (own `_projection` counter, no iterative recon, no buffer + held across a finalize). + +## Validation results (container, A400, 2026-07-02) +All three design tests pass (`test_data/run_pr4c/d/e.log`): +- **Test 1 — regression** (2 proj, no stress): both projections reconstruct to exactly + 1024, all 4 files, same iteration budget as single-buffer (proj0=7, proj1=5). Ping-pong + logs confirm the flips (`Accumulator flipped to buffer 1` → `flipped read buffer to 1`). +- **Test 2 — overlap** (2 proj, 1.5 s injected finalize): the accumulator flipped to buf1 + and **fully accumulated proj-1 during proj-0's 2 s finalize** — never backpressured, + GatherOp cache flat. This is the win vs single-buffer Run B (which stalled and relied on + the sim's PUSH throttle). +- **Test 3 — lapping fallback** (3 proj, 3 s injected finalize): when the recon was lapped + by a full projection, the accumulator logged `Recon lagging … backpressuring upstream` + and **recovered** once the recon released the buffer — clean backpressure, no clobbering, + all 3 projections at 1024, no drops/deadlock. + +**O1 is now more urgent (observed):** when the recon lags acquisition, later projections +arrive **pre-filled** and finalize with near-zero refinement iterations (Test 2 proj-1 got +1 iteration vs 5 when filled concurrently). Double-buffering makes this visible; the fix is +per-projection post-stream iterations (deferred O1). On the faster detector this will matter. + +## Note +All `PR4-MEASURE` instrumentation + the sim tweaks (position mapping, `PTYCHO_CENTER`, +`dummy_img_index=True`) were working-tree-only test scaffolding, reverted before this commit +(the one kept production log is the `Recon lagging …` backpressure warning). The +`num_buffers`/ping-pong state lives in `ptycho_state`; `config_sim.yaml` + `Dockerfile_bwell` +remain untracked local test scaffolding. diff --git a/pipeline/TESTING_header_tomo.md b/pipeline/TESTING_header_tomo.md new file mode 100644 index 0000000..7a93b1c --- /dev/null +++ b/pipeline/TESTING_header_tomo.md @@ -0,0 +1,62 @@ +# Testing log — header / dynamic-geometry / tomography feature + +Branch `feat/scan-header-tomo` (PR0 → PR3). Tests run in the `ptycho-holoscan:ptyrex` +container on the **A400** (`CUDA_VISIBLE_DEVICES=1`), driven by the daqsim simulator +(`daqsim:latest`, Dectris SIMPLON emulator) streaming scan **409907** (32×32 = 1024 +frames, 515×515 uint32). Sim tweaks applied via `test_data/sim_code_tweaks.patch`; +geometry headers sent with `test_data/send_header.py`; scans triggered with +`dectris-hackathon/daqsim/trigger.py --stream both --nimages 1024`. + +## PR0 — Unwire PublishToCloudOp +- Regression: STXM + ptycho scans produce identical output with the op unwired. Pipeline + composes and runs. (Validated during the PR0/PR1 session.) + +## PR1 — Flush plumbing + hardening +- Single ptycho scan reconstructs to `total_iterations` and idles; `recon_complete` logged. +- Back-to-back second scan: start-flush safety net fires, clean second reconstruction, two + output HDFs. No-double-flush confirmed (`Start-flush skipped — already flushed on completion`). + +## PR2 — Live header operator + dynamic geometry — ALL PASS (2026-07-02) +Container run, logs captured to `test_data/run34.log`. + +| # | Test | Result | +|---|------|--------| +| 1 | Launch with `npoints_*` removed from config | ✅ `buffers allocated at capacity: 1024 frames`, `Scan geometry configured: 1024 frames … object size (…,523,523)` | +| 2 | Baseline scan + back-to-back regression | ✅ both reconstruct to iter 25 and flush; no-double-flush holds | +| 3 | Header before data → full scan on reconfigured geometry | ✅ full handshake (`Received header`→`Header received`→`Recon quiesced`→`Applied new scan geometry`→`Recon reset`); scan then reconstructs to completion, **all 1024 positions valid**. 2nd `before_reconstruction_stream` stable. | +| 4 | Header mid-reconstruction (preemption) | ✅ `Saved partial result before preemption (iter 24)` emitted **before** flush/reconfigure, then quiesce → reconfigure → reset | +| 4b | Recovery after preemption | ✅ `GPU initialisation complete` → fresh scan `Reconstruction complete at iteration 25` | +| 5 | Oversized grid (64×64 = 4096 > capacity 1024) | ✅ `Header grid 64x64 … exceeds capacity 1024 — rejected`; rejected before staging, pipeline stays alive | + +**Caveat:** reconfigure tested only to the *same* geometry (32×32 / 1.5 µm, object stayed +523×523). The reconfigure + GPU-reinit code path is fully exercised; a change to a +*different* object size with matching stream data is unverified (needs a 2nd dataset). + +## PR3 — Tomography / multi-projection — PASS (2026-07-02) +Container run (A400, `config_sim.yaml`), 2-projection header sent, 2048-frame stream +(`trigger.py --stream both --nimages 2048`, sequence_id 29). Logs in `test_data/run_pr3c.log`. + +| # | Test | Result | +|---|------|--------| +| 1 | Header sets tomography mode | ✅ `Received header … num_projections=2` → handshake → geometry applied | +| 2 | Projection 0 completes at frame boundary | ✅ `Wrote projection 0 (1024 frames) → 29_proj00.h5`, `Saved projection recon → 29_proj00_recon.h5`, `Projection 0/2 complete … signal=projection_complete` (early-stopped at iter 7) | +| 3 | Scoped advance (accum + recon reset, GatherOp cache preserved) | ✅ `Projection complete — advanced to projection 1`; accumulator advanced with carry preserved, recon object reset / probe carried, GatherOp **not** flushed | +| 4 | Projection 1 completes | ✅ `Wrote projection 1 (1024 frames) → 29_proj01.h5`, `Saved projection recon → 29_proj01_recon.h5`, `Projection 1/2 complete … signal=recon_complete` (final; early-stopped at iter 4) | +| 5 | Final flush | ✅ `Reconstruction complete — flushing for next scan` | + +All four files on disk (`series_id=29`, both projections × STXM + recon), 1024-frame counts +exact, and the non-final vs final completion signals correctly distinguished +(`projection_complete` for proj-0, `recon_complete` for proj-1). No stall. + +**Test-tweak note:** the run used `dummy_img_index=True` (continuous synthetic image IDs) +to work around a daqsim looped-replay artifact where image IDs reset per loop while +position IDs continue. Testing exposed an off-by-one in that path +(`data_io.py`: `series_frame_count - 1` produced IDs `-1..N-2`, leaving one frame +unmatched per scan) — fixed to `series_frame_count`. `dummy_img_index` reverted to +`False` before commit; a real single-series tomography acquisition (monotonic IDs) +would not hit the artifact. + +**Open item (see DECISIONS O1):** each projection early-stops the moment its frames +arrive, so proj-1 got fewer PIE iterations than proj-0 (cached frames tripped +`all_data_arrived` immediately). Per-projection post-stream refinement iterations +deferred to PR4. diff --git a/pipeline/config_prod.yaml b/pipeline/config_prod.yaml index d225900..f59c0d2 100644 --- a/pipeline/config_prod.yaml +++ b/pipeline/config_prod.yaml @@ -18,6 +18,13 @@ position_src: zmq_endpoint: "tcp://172.23.82.204:6666" receive_timeout_ms: 1000 +header_src: + # Dedicated ZMQ SUB socket for live scan-geometry headers (PR2). A JSON + # object with npoints_h/npoints_v/step_size_h/step_size_v/num_projections + # reconfigures scan geometry on the fly and preempts an in-flight recon. + zmq_endpoint: "tcp://172.23.82.204:6667" # production endpoint (placeholder — adjust) + receive_timeout_ms: 100 + masking_op: center_x: 257 center_y: 515 @@ -38,6 +45,14 @@ ptychography: ptyrex_config: "/workdir/test_data/pty_config_15keV_streamTest_6mm.json" scan_ID: [1, 1, 1] ID: [1, 1, 1] + # Scan grid now comes from the live header (header_src). max_npoints_* set the + # GPU buffer capacity (allocated once, never realloced — R-6); default_step_* + # is the startup geometry before any header arrives. (Legacy no_frames/R/ + # scan_range below are unused — dead keys, S3, left for a separate cleanup.) + max_npoints_h: 32 # capacity = 32*32 = 1024 frames (matches prior no_frames: 1024) + max_npoints_v: 32 # bump these (with GPU RAM in mind) if larger scans are needed + default_step_size_h: 0.25 + default_step_size_v: 0.25 no_frames: 1024 total_iterations: 20 post_stream_iterations: 1 diff --git a/pipeline/config_test.yaml b/pipeline/config_test.yaml index b516910..ee525b1 100644 --- a/pipeline/config_test.yaml +++ b/pipeline/config_test.yaml @@ -5,25 +5,33 @@ scheduler: worker_threads: 4 image_src: - zmq_endpoint: "tcp://172.23.82.48:31001" # Production endpoint - #zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator + #zmq_endpoint: "tcp://172.23.82.131:31001" # Production endpoint + zmq_endpoint: "tcp://172.23.82.77:5555" # Local simulator receive_timeout_ms: 1000 # Increased timeout for testing - batch_size: 100 # Larger batch for testing + batch_size: 180 # Larger batch for testing decompress_op: - # data_size: [192, 192] # Match selun test data (192x192 images) - data_size: [514, 1030] # Match selun test data (192x192 images) - data_dtype: uint16 + data_size: [190, 190] # Match selun test data (192x192 images) + #data_size: [514, 1030] # Match selun test data (192x192 images) + data_dtype: uint32 position_src: - # zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator - zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint + zmq_endpoint: "tcp://172.23.82.77:5556" # Local simulator + #zmq_endpoint: "tcp://172.23.82.204:6666" # production endpoint receive_timeout_ms: 1000 # Increased timeout for testing +header_src: + # Dedicated ZMQ SUB socket for live scan-geometry headers (PR2). A JSON + # object with npoints_h/npoints_v/step_size_h/step_size_v/num_projections + # reconfigures scan geometry on the fly and preempts an in-flight recon. + zmq_endpoint: "tcp://172.23.82.77:5558" # Local simulator + #zmq_endpoint: "tcp://172.23.82.32:6667" # production endpoint (placeholder) + receive_timeout_ms: 100 # short so the blocking recv doesn't hold a worker thread + masking_op: - center_x: 524 # Adjusted for 192x192 images (center) - center_y: 287 # Adjusted for 192x192 images (center) - radius: 15 # Adjusted proportionally + center_x: 94 # Adjusted for 192x192 images (center) + center_y: 94 # Adjusted for 192x192 images (center) + radius: 20 # Adjusted proportionally sink_and_publish_op: publish_tensors: ["positions", "positions_ids", "inner", "outer", "intensity_ids"] @@ -37,15 +45,20 @@ sink_and_publish_op: ptychography: enabled: true - ptyrex_config: "/workdir/PtyREX/config_409907.json" + ptyrex_config: "/workdir/PtyREX/config_414287.json" scan_ID: [1, 1, 1] ID: [1, 1, 1] - npoints_h: 100 # horizontal scan points - npoints_v: 100 # vertical scan points - step_size_h: 0.25 # horizontal step size (microns) - step_size_v: 0.25 # vertical step size (microns) - total_iterations: 25 + # Scan grid now comes from the live header (header_src). These set the GPU + # buffer capacity (allocated once, never realloced — R-6) and the default + # geometry used at startup before any header arrives. + max_npoints_h: 200 # buffer capacity + startup default grid (horizontal) + max_npoints_v: 100 # buffer capacity + startup default grid (vertical) + # capacity = 100*100 = 10000 frames (matches the prior committed 100x100 grid; + # a header requesting more frames is rejected — bump these with GPU RAM in mind) + default_step_size_h: 0.2 # startup step size (microns), until a header arrives + default_step_size_v: 0.2 # startup step size (microns), until a header arrives + total_iterations: 35 post_stream_iterations: 1 housekeeping_interval: 1 publish_interval: 1 - reset_probe: true # false = carry previous scan's probe forward (warm start); true = full probe reset each scan + reset_probe: false # false = carry previous scan's probe forward (warm start); true = full probe reset each scan diff --git a/pipeline/control.py b/pipeline/control.py index 36aa720..6097dd3 100644 --- a/pipeline/control.py +++ b/pipeline/control.py @@ -12,52 +12,144 @@ class ControlOp(Operator): """ Control operator for managing pipeline flow. - Handles control messages like flush and processing_end, - coordinating state across multiple operators. + Handles the flush control message, coordinating flush state across + multiple operators. """ def __init__(self, fragment, *args, - flushable_ops: list[Operator] = None, + stxm_flush_ops: list[Operator] = None, + ptycho_flush_ops: list[Operator] = None, publish_backend = None, + ptycho_accum = None, + ptycho_recon = None, + scan_state: dict = None, **kwargs): """ Initialize control operator. - + Args: fragment: Holoscan fragment - flushable_ops: List of operators that can be flushed + stxm_flush_ops: STXM-side operators that can be flushed + ptycho_flush_ops: Ptycho-side operators that can be flushed publish_backend: Backend instance for publishing flush messages + ptycho_accum: PtychoAccumulatorOp (for the scoped projection advance) + ptycho_recon: PtychoReconstructionOp (for the scoped projection advance) + scan_state: shared holder whose current_projection is advanced at a + tomography projection boundary (PR3) """ super().__init__(fragment, *args, **kwargs) self.logger = logging.getLogger(kwargs.get("name", "ControlOp")) - self.flushable_ops = flushable_ops + self.stxm_flush_ops = stxm_flush_ops or [] + self.ptycho_flush_ops = ptycho_flush_ops or [] self.publish_backend = publish_backend - + self.ptycho_accum = ptycho_accum + self.ptycho_recon = ptycho_recon + self.scan_state = scan_state + # True once a completion (recon_complete) flush has run and no new scan + # has started since. Lets the scan-start flush skip when the buffers are + # already clean, so we don't double-flush (Task 3 flush-check-at-start). + self._flushed = False + def setup(self, spec: OperatorSpec): spec.input("input").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128) - spec.output("output") + + def _do_stxm_flush(self): + """Flush STXM-side operators only.""" + for op in self.stxm_flush_ops: + op.flush() + + def _do_ptycho_flush(self): + """Flush ptycho-side operators only.""" + for op in self.ptycho_flush_ops: + op.flush() + + def _request_ptycho_flush(self): + """Request ptycho-side flush without doing any STXM work.""" + self._do_ptycho_flush() + + def _do_full_flush(self): + """Flush STXM + ptycho operators and broadcast flush signals.""" + self._do_stxm_flush() + self._do_ptycho_flush() + if self.publish_backend is not None: + import numpy as np + self.publish_backend.publish("stxm_flush", np.array([1])) # Simple signal + # Also signal ptycho consumers; harmless when ptycho is disabled + # (no subscriber listens on this subject). + self.publish_backend.publish("ptycho_flush", np.array([1])) def compute(self, op_input, op_output, context): """Handle control messages.""" msg = op_input.receive("input") - - if msg == "flush": - # Flush all flushable operators - for op in self.flushable_ops: - op.flush() - - # Publish flush message through the backend if available + transition_blocked = False + transition_phase = "idle" + if self.scan_state is not None: + transition_blocked = bool(self.scan_state.get("transition_blocked_event")) + if transition_blocked: + transition_blocked = self.scan_state["transition_blocked_event"].is_set() + transition_phase = self.scan_state.get("transition_phase", "idle") + self.logger.info( + "Control message=%s (transition_blocked=%s phase=%s)", + msg, + transition_blocked, + transition_phase, + ) + + if msg == "recon_complete": + # The recon finished its final iteration and has ALREADY saved + # (after_iteration -> pty_out) and published the result before emitting + # this, so flushing now is safe (Task 3: flush after the last iteration). + if transition_blocked and transition_phase == "waiting_quiesce": + self.logger.info( + "Recon quiesced for header transition — requesting ptycho flush" + ) + self._request_ptycho_flush() + if self.scan_state is not None: + self.scan_state["transition_phase"] = "waiting_flush_exec" + self.logger.info("Transition state -> waiting_flush_exec") + return + + self.logger.info("Reconstruction complete — flushing for next scan") + self._do_full_flush() + self._flushed = True + + # PR4: tomography projection boundaries no longer round-trip through + # ControlOp. With double-buffering the accumulator flips write buffers + # itself and the recon owns current_projection + the read-buffer flip + # (avoiding a save-vs-bump race), so there is no "projection_complete" + # signal any more — the recon emits "recon_complete" only on the FINAL + # projection, handled above. ControlOp is kept for recon_complete / header + # / flush. + + elif msg == "header": + # A live header reconfigures the scan for a new dataset. Flush so the + # STXM path saves+clears its current buffer before reconfiguration + # (SinkAndPublishOp.flush writes any unwritten scan). Ptycho flush is + # deferred until the recon quiesces and emits recon_complete. + self.logger.info("Header received — flushing for reconfigured scan") + self._do_stxm_flush() if self.publish_backend is not None: import numpy as np - self.publish_backend.publish("stxm_flush", np.array([1])) # Simple signal - # Also signal ptycho consumers; harmless when ptycho is disabled - # (no subscriber listens on this subject). - self.publish_backend.publish("ptycho_flush", np.array([1])) - - elif msg == "processing_end": - # Forward processing_end signal - op_output.emit("processing_end", "output") - + self.publish_backend.publish("stxm_flush", np.array([1])) + self._flushed = True + + elif msg == "flush": + # Scan-start safety flush: only flush if the buffers aren't already + # clean from a completion flush. If the previous scan completed, this + # no-ops (no double flush); if it was interrupted, this cleans up. + if transition_blocked: + self.logger.info( + "Start-flush deferred for blocked transition — STXM-only flush now, ptycho waits for recon quiesce" + ) + self._do_stxm_flush() + self._flushed = True + return + if self._flushed: + self.logger.info("Start-flush skipped — already flushed on completion") + self._flushed = False + else: + self._do_full_flush() + else: self.logger.info(f"Received unknown message: {msg}") diff --git a/pipeline/data_io.py b/pipeline/data_io.py index 73052d9..373eaff 100644 --- a/pipeline/data_io.py +++ b/pipeline/data_io.py @@ -45,7 +45,7 @@ def receive_cbor_message(zmq_message) -> tuple[str, cbor2.CBORTag, int, dict]: if msg_type == "image": compressed_image, image_id, msg_content = msg["data"]["threshold_1"], msg["image_id"], None elif msg_type == "start": - print(f"{msg_type} message content: {msg}") + print(f"Received {msg_type} message.") # content: {msg} compressed_image, image_id, msg_content = None, None, msg elif msg_type == "end": print(f"{msg_type} message content: {msg}") @@ -203,7 +203,7 @@ def compute(self, op_input, op_output, context): # Handle data message datasets = msg["datasets"] #print(datasets) - + # Extract position data x = np.array(datasets["/pi_x"]["data"]) #FMC_IN.VAL1.Mean y = np.array(datasets["/FMC_IN.VAL2.Mean"]["data"]) @@ -376,7 +376,7 @@ def compute(self, op_input, op_output, context): self.batch[self.current_index] = data if self.dummy_img_index: - self.batch_ids[self.current_index] = self.series_frame_count - 1 + self.batch_ids[self.current_index] = self.series_frame_count else: # self.logger.info(f"Received image with id: {data_id}") self.batch_ids[self.current_index] = data_id @@ -459,7 +459,14 @@ class GatherOp(Operator): gather -> masking_op -> publish """ - def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): + def __init__( + self, + fragment, + *args, + batch_size: int = 1, + scan_state: dict = None, + **kwargs, + ): """ Initialize gather operator. @@ -472,6 +479,18 @@ def __init__(self, fragment, *args, batch_size: int = 1, **kwargs): self.position_ids = np.zeros((0,), dtype=int) self.count = 0 self.batch_size = int(batch_size) + self.scan_state = scan_state + # R-1 (PR3): latched once the series-end metadata is seen, so the final + # partial batch (< batch_size) is drained instead of stranded. Reset on flush. + self._series_finished = False + # Deferred flush: flush() sets this flag and the actual cache clear + # happens at the top of the next compute(), so it never mutates the + # caches while compute() is mid-synchronise. This avoids the boolean- + # index race (data_io.py "size of axis is 0 but ... 64") when a flush + # arrives mid-stream — e.g. PR2's header preemption. + self._flush_requested = False + # PR5: track blocked matched-frame growth for observability. + self._last_blocked_common = -1 self.logger = logging.getLogger(kwargs.get("name", "GatherOp")) super().__init__(fragment, *args, **kwargs) @@ -481,24 +500,38 @@ def setup(self, spec: OperatorSpec): spec.input("positions").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) spec.output("output") + def _transition_blocked(self): + if self.scan_state is None: + return False + blocked_event = self.scan_state.get("transition_blocked_event") + return bool(blocked_event is not None and blocked_event.is_set()) + def flush(self): - """Reset all cached data on flush.""" + """Request a cache reset. Deferred to the top of the next compute() so it + never clears the caches while compute() is mid-synchronise (thread-safe).""" + self._flush_requested = True + + def _perform_flush(self): + """Actually clear the caches — only ever called from compute().""" self.images = None self.image_ids = np.zeros((0,), dtype=int) self.positions = np.zeros((0, 4)) self.position_ids = np.zeros((0,), dtype=int) self.count = 0 + self._series_finished = False + self._last_blocked_common = -1 self.logger.info( - f"[FLUSH VERIFY] GatherOp cleared: " - f"images={None if self.images is None else self.images.shape}, " - f"image_ids={self.image_ids.size}, " - f"positions={self.positions.shape}, " - f"position_ids={self.position_ids.size}, " - f"count={self.count}" + "[FLUSH VERIFY] GatherOp cleared: images=None, image_ids=0, " + "positions=(0, 4), position_ids=0, count=0" ) - def compute(self, op_input, op_output, context): + def compute(self, op_input, op_output, context): """Gather and synchronize image and position data.""" + # Perform any requested flush here — single-threaded w.r.t. the caches. + if self._flush_requested: + self._flush_requested = False + self._perform_flush() + # Receive image data images_dict = op_input.receive("images") if images_dict is not None: @@ -523,12 +556,42 @@ def compute(self, op_input, op_output, context): self.positions = np.concatenate([self.positions, positions]) self.position_ids = np.concatenate([self.position_ids, position_ids]) + # R-1 (PR3): once the series has ended, drain the remaining matched IDs + # even if fewer than batch_size, so a final partial batch (no_frames not a + # multiple of batch_size) reaches the accumulator instead of stranding and + # idling the pipeline forever. series_finished flows from the image source + # via metadata on the final batch; latch it so the drain persists. + try: + if self.metadata is not None and self.metadata.get("series_finished", False): + self._series_finished = True + except Exception: + pass + # Find common IDs between images and positions if self.images is not None and self.image_ids.size > 0 and self.position_ids.size > 0: common_ids = np.intersect1d(self.image_ids, self.position_ids).astype(int) - + + drain = self._series_finished and int(common_ids.size) > 0 #if common_ids.size > 0: - if int(common_ids.size) >= self.batch_size: + if int(common_ids.size) >= self.batch_size or drain: + blocked_common = int(common_ids.size) + if self._transition_blocked(): + if blocked_common != self._last_blocked_common: + self._last_blocked_common = blocked_common + self.logger.info( + "Transition blocked: caching %d matched frames in GatherOp", + blocked_common, + ) + return + + # Transition resumed / not blocked. + if self._last_blocked_common >= 0: + self.logger.info( + "Transition unblocked: resuming GatherOp emit with %d cached matched frames", + blocked_common, + ) + self._last_blocked_common = -1 + # Create vectorized masks for efficient filtering mask_positions = np.isin(self.position_ids, common_ids) mask_images = np.isin(self.image_ids, common_ids) diff --git a/pipeline/header_io.py b/pipeline/header_io.py new file mode 100644 index 0000000..72fd8a6 --- /dev/null +++ b/pipeline/header_io.py @@ -0,0 +1,180 @@ +""" +Header Input for the Holoscan Ptycho Pipeline + +Defines HeaderRxOp: a dedicated ZMQ SUB operator that listens for live +scan-geometry headers and reconfigures the pipeline on the fly (PR2). + +A header is a JSON object, e.g.:: + + {"npoints_h": 100, "npoints_v": 100, + "step_size_h": 0.25, "step_size_v": 0.25, + "num_projections": 1} + +On a valid header the operator: + 1. updates the always-present shared ``scan_state`` (projection count), + 2. stages the new geometry in ``ptycho_state`` and requests preemption of an + in-flight reconstruction (R-4 handshake; the recon op applies the geometry + once it has finished the current iteration, saved, and quiesced), and + 3. emits a ``"header"`` token to ControlOp so the STXM path flushes for the + new dataset (works even when ptychography is disabled). + +Malformed headers are logged and ignored without disturbing an in-flight scan. +""" + +import logging + +import zmq + +from holoscan.core import Operator, OperatorSpec, ConditionType + + +class HeaderRxOp(Operator): + """Receive live scan-geometry headers over a dedicated ZMQ SUB socket.""" + + def __init__( + self, + fragment, + *args, + zmq_endpoint: str = None, + receive_timeout_ms: int = 100, + scan_state: dict = None, + ptycho_state: dict = None, + **kwargs, + ): + """ + Args: + fragment: Holoscan fragment + zmq_endpoint: ZMQ endpoint to connect to (e.g. "tcp://host:5557") + receive_timeout_ms: recv timeout in ms (short so the blocking recv + does not hold a worker thread for long) + scan_state: always-present shared holder for projection/frame counts + ptycho_state: ptycho state (None when ptychography is disabled) + """ + self.logger = logging.getLogger(kwargs.get("name", "HeaderRxOp")) + logging.basicConfig(level=logging.INFO) + + self.endpoint = zmq_endpoint + context = zmq.Context() + self.socket = context.socket(zmq.SUB) + self.socket.setsockopt_string(zmq.SUBSCRIBE, "") + self.socket.setsockopt(zmq.RCVTIMEO, receive_timeout_ms) + + try: + self.socket.connect(self.endpoint) + except zmq.error.ZMQError: + self.logger.error("Failed to connect header socket to %s", self.endpoint) + + self.scan_state = scan_state + self.ptycho_state = ptycho_state + + super().__init__(fragment, *args, **kwargs) + + def setup(self, spec: OperatorSpec): + # Token to ControlOp; NONE condition so this source runs freely. + spec.output("header").condition(ConditionType.NONE) + + def _validate(self, msg): + """Validate a header dict; return the parsed tuple or None if malformed.""" + if not isinstance(msg, dict): + self.logger.warning("Header is not a JSON object: %r", msg) + return None + try: + npoints_h = int(msg["npoints_h"]) + npoints_v = int(msg["npoints_v"]) + step_size_h = float(msg["step_size_h"]) + step_size_v = float(msg["step_size_v"]) + num_projections = int(msg.get("num_projections", 1)) + except (KeyError, TypeError, ValueError) as exc: + self.logger.warning("Malformed header %r: %s", msg, exc) + return None + if ( + npoints_h <= 0 or npoints_v <= 0 + or step_size_h <= 0 or step_size_v <= 0 + or num_projections < 1 + ): + self.logger.warning("Header has non-positive values: %r", msg) + return None + return npoints_h, npoints_v, step_size_h, step_size_v, num_projections + + def compute(self, op_input, op_output, context): + try: + msg = self.socket.recv_json() + except zmq.error.Again: + return # recv timed out — nothing to do this tick + except Exception as exc: # noqa: BLE001 - don't let a bad message kill the op + self.logger.warning("Header receive error: %s", exc) + return + + parsed = self._validate(msg) + if parsed is None: + return # malformed — ignore, leave any in-flight scan untouched + + npoints_h, npoints_v, step_size_h, step_size_v, num_projections = parsed + self.logger.info( + "Received header: %d x %d points, step %.4g x %.4g µm, " + "num_projections=%d", + npoints_h, npoints_v, step_size_h, step_size_v, num_projections, + ) + + # Reject a grid that exceeds the pre-allocated GPU capacity BEFORE staging + # it (buffers are never realloced, R-6). Rejecting here keeps the bad + # header off the recon's apply path, which would otherwise raise inside + # compute() and take down the pipeline. + if self.ptycho_state is not None: + capacity = self.ptycho_state.get("capacity") + if capacity is not None and npoints_h * npoints_v > capacity: + self.logger.warning( + "Header grid %dx%d = %d frames exceeds capacity %d — rejected " + "(increase max_npoints_h/max_npoints_v). In-flight scan " + "untouched.", + npoints_h, npoints_v, npoints_h * npoints_v, capacity, + ) + return + + # 1. Update the always-present shared holder (S11). Set no_frames here too + # (not only via configure_scan_geometry) so the STXM sink can segment + # per projection even when ptychography is disabled. + if self.scan_state is not None: + prev_phase = self.scan_state.get("transition_phase", "idle") + prev_blocked = False + blocked_event = self.scan_state.get("transition_blocked_event") + if blocked_event is not None: + prev_blocked = blocked_event.is_set() + self.scan_state["num_projections"] = num_projections + self.scan_state["current_projection"] = 0 + self.scan_state["no_frames"] = npoints_h * npoints_v + self.logger.info( + "Header accepted: no_frames=%d num_projections=%d (prev_phase=%s, prev_blocked=%s)", + self.scan_state["no_frames"], + self.scan_state["num_projections"], + prev_phase, + prev_blocked, + ) + + # 2. Ptycho path: stage geometry + request preemption (R-4). The recon op + # applies configure_scan_geometry once it has quiesced, so no buffer + # view is swapped under a live PIE iteration. + if self.ptycho_state is not None: + with self.ptycho_state["lock"]: + self.ptycho_state["pending_geometry"] = { + "npoints_h": npoints_h, + "npoints_v": npoints_v, + "step_size_h": step_size_h, + "step_size_v": step_size_v, + } + self.ptycho_state["preempt_requested"].set() + if self.scan_state is not None: + self.scan_state["transition_blocked_event"].set() + self.scan_state["transition_phase"] = "waiting_quiesce" + self.scan_state["ptycho_accum_flushed"] = False + self.scan_state["ptycho_recon_flushed"] = False + self.scan_state["transition_error"] = None + self.logger.info( + "Transition state -> waiting_quiesce (blocked=True, acks reset)" + ) + self.logger.info("Header received — transition blocked, waiting for recon quiesce") + + # 3. Notify ControlOp so the STXM path flushes for the new dataset + # (works even when ptychography is disabled). + self.logger.info("Emitting header control token") + op_output.emit("header", "header") diff --git a/pipeline/pipeline.py b/pipeline/pipeline.py index 7a2316a..7b39d6d 100644 --- a/pipeline/pipeline.py +++ b/pipeline/pipeline.py @@ -14,6 +14,7 @@ """ import logging +import threading from argparse import ArgumentParser from holoscan.core import Application @@ -28,8 +29,9 @@ GatherOp ) from processing import MaskingOp -from publish import SinkAndPublishOp, PublishToCloudOp +from publish import SinkAndPublishOp from control import ControlOp +from header_io import HeaderRxOp # Try to import NATS (optional for testing) try: @@ -58,6 +60,20 @@ def __init__(self, *args, **kwargs): self.num_decompress_ops = 4 self.ptychography_enabled = False self.ptycho_state = None + # Always-present shared holder (S11) for projection/frame counts, written + # by the header op and read by both the STXM and ptycho paths. Populated + # in main(); the frame count is filled in by configure_scan_geometry. + self.scan_state = { + "no_frames": 0, + "num_projections": 1, + "current_projection": 0, + "transition_blocked_event": threading.Event(), + "transition_phase": "idle", + "ptycho_accum_flushed": False, + "ptycho_recon_flushed": False, + "transition_error": None, + "max_blocked_frames": None, + } super().__init__(*args, **kwargs) self.enable_metadata(True) @@ -130,20 +146,15 @@ def compose(self): sink_and_publish_op = SinkAndPublishOp(self, tensor2subject=tensor2subject, publish_backend=publish_backend, + scan_state=self.scan_state, **self.kwargs('sink_and_publish_op'), name="sink_and_publish_op") - # ===== Cloud Publishing Operator ===== - publish_folder = self.kwargs('sink_and_publish_op')['publish_folder'] - temp_folder = self.kwargs('sink_and_publish_op')['temp_folder'] - - publish_to_cloud_op = PublishToCloudOp(self, - publish_folder=publish_folder, - temp_folder=temp_folder, - name="publish_to_cloud_op") - # ===== Control Operator ===== - flushable_ops = [gather_op, position_src, sink_and_publish_op] + stxm_flush_ops = [gather_op, position_src, sink_and_publish_op] + ptycho_flush_ops = [] + ptycho_accum = None # set below when ptychography is enabled + ptycho_recon = None # ===== Ptychography Branch (conditional) ===== if self.ptychography_enabled: @@ -170,6 +181,7 @@ def compose(self): housekeeping_interval=ptycho_cfg["housekeeping_interval"], publish_interval=ptycho_cfg["publish_interval"], reset_probe=ptycho_cfg.get("reset_probe", False), + publish_folder=sink_config.get("publish_folder"), name="ptycho_reconstruction", ) @@ -179,14 +191,32 @@ def compose(self): name="ptycho_publish", ) - flushable_ops.append(ptycho_accum) - flushable_ops.append(ptycho_recon) + ptycho_flush_ops.append(ptycho_accum) + ptycho_flush_ops.append(ptycho_recon) control_op = ControlOp(self, - flushable_ops=flushable_ops, + stxm_flush_ops=stxm_flush_ops, + ptycho_flush_ops=ptycho_flush_ops, publish_backend=publish_backend, + ptycho_accum=ptycho_accum, + ptycho_recon=ptycho_recon, + scan_state=self.scan_state, name="control_op") + # ===== Header Source (live scan geometry, optional) ===== + # A dedicated ZMQ SUB socket for JSON scan-geometry headers. Reconfigures + # geometry on the fly and preempts an in-flight recon (R-4 handshake). + header_src = None + header_cfg = self.kwargs('header_src') + if header_cfg: + header_src = HeaderRxOp(self, + name="header_src", + scan_state=self.scan_state, + ptycho_state=self.ptycho_state, + **header_cfg) + else: + logger.warning("No header_src config — live scan headers disabled") + # ===== Connect Operators ===== # I/O: Image reception and decompression -> gather for i in range(self.num_decompress_ops): @@ -204,11 +234,17 @@ def compose(self): if self.ptychography_enabled: self.add_flow(gather_op, ptycho_accum, {("output", "input")}) self.add_flow(ptycho_recon, ptycho_publish, {("output", "input")}) + # Completion signal → control (logged in PR1; drives flush in PR3) + self.add_flow(ptycho_recon, control_op, {("complete", "input")}) - # Control path: flush and completion signals + # Control path: flush signal (start message → unconditional idempotent flush) self.add_flow(img_src, control_op, {("flush", "input")}) - self.add_flow(sink_and_publish_op, control_op, {("processing_end", "input")}) - self.add_flow(control_op, publish_to_cloud_op, {("output", "trigger")}) + + # Header path: live geometry header → control (flush for new dataset). + # The ptycho geometry reconfigure is driven separately via the R-4 + # handshake in ptycho_state (header op sets preempt_requested). + if header_src is not None: + self.add_flow(header_src, control_op, {("header", "input")}) def main(): @@ -231,21 +267,30 @@ def main(): # Load config to make kwargs available app.config(args.config) + image_src_config = app.kwargs('image_src') + # Get scheduler parameters from config via kwargs scheduler_config = app.kwargs('scheduler') num_decompress_ops = scheduler_config.get('num_decompress_ops', 4) worker_threads = scheduler_config.get('worker_threads', 6) + ptycho_cfg = app.kwargs("ptychography") + default_blocked_frames = int(image_src_config.get("batch_size", 100)) * 10 + app.scan_state["max_blocked_frames"] = int( + ptycho_cfg.get("max_blocked_frames", default_blocked_frames) + ) if ptycho_cfg else default_blocked_frames + # Set num_decompress_ops - will be used in compose() when run() is called app.num_decompress_ops = num_decompress_ops # Ptychography setup (before compose) - ptycho_cfg = app.kwargs("ptychography") if ptycho_cfg and ptycho_cfg.get("enabled", False): from ptychography_setup import init_ptycho_state logger.info("Initialising ptychography state…") - app.ptycho_state = init_ptycho_state(ptycho_cfg) + # Pass the shared scan_state so configure_scan_geometry can mirror the + # frame count for the STXM path and header op (S11). + app.ptycho_state = init_ptycho_state(ptycho_cfg, app.scan_state) app.ptychography_enabled = True worker_threads = max(worker_threads, 8) logger.info("Ptychography enabled (worker_threads=%d)", worker_threads) diff --git a/pipeline/ptychography_ops.py b/pipeline/ptychography_ops.py index 12f40ba..b345374 100644 --- a/pipeline/ptychography_ops.py +++ b/pipeline/ptychography_ops.py @@ -40,6 +40,41 @@ ) +def _ack_flush_and_maybe_release(scan_state, ack_key, logger): + """Mark one ptycho flush-execution ack and release the transition barrier + only when both ptycho acks are present in waiting_flush_exec.""" + if not scan_state: + return + + blocked_event = scan_state.get("transition_blocked_event") + if blocked_event is None or not blocked_event.is_set(): + return + + if scan_state.get("transition_phase", "idle") != "waiting_flush_exec": + return + + scan_state[ack_key] = True + logger.info( + "Transition flush ack set: %s=True (phase=%s accum=%s recon=%s)", + ack_key, + scan_state.get("transition_phase", "idle"), + scan_state.get("ptycho_accum_flushed", False), + scan_state.get("ptycho_recon_flushed", False), + ) + + if ( + scan_state.get("ptycho_accum_flushed", False) + and scan_state.get("ptycho_recon_flushed", False) + ): + blocked_event.clear() + scan_state["transition_phase"] = "idle" + scan_state["ptycho_accum_flushed"] = False + scan_state["ptycho_recon_flushed"] = False + logger.info( + "Transition barrier released after both ptycho flush executions (phase=idle, blocked=False)" + ) + + class PtychoAccumulatorOp(Operator): """Fast batch accumulator for ptychography. @@ -51,6 +86,22 @@ class PtychoAccumulatorOp(Operator): def __init__(self, fragment, *args, ptycho_state, **kwargs): self.ptycho_state = ptycho_state self.lock = ptycho_state["lock"] + # Tracks whether any frames have been accumulated since the last flush, + # so a redundant flush (e.g. the unconditional start-flush, R-2) is free. + self._dirty = False + # Deferred flush (see GatherOp): flush() sets this flag; the actual reset + # happens at the top of the next compute(), so it never zeros the GPU + # buffers while compute() is mid-write — race-safe even for a mid-stream + # flush (PR2 header preemption). + self._flush_requested = False + # PR3 tomography: frames that straddle a projection boundary are split — + # the head fills the current projection, the tail is carried here and + # written into the next projection once the boundary advance (flush) has + # reset filled_until to 0. None when there is no pending carry. + self._carry = None + # One-shot flag so the "recon lagging" backpressure warning fires once per + # stall, not every 10 ms tick. + self._lapped = False self.logger = logging.getLogger(kwargs.get("name", "PtychoAccumulatorOp")) super().__init__(fragment, *args, **kwargs) @@ -60,19 +111,109 @@ def setup(self, spec: OperatorSpec): ).condition(ConditionType.NONE) def flush(self): - """Reset fill level and zero out GPU buffers for a new series.""" + """Request a FULL reset (scan start/end/preempt); performed at the top of + the next compute() (deferred, so it never races the buffer writes).""" + self._flush_requested = True + + def _reset_pingpong(self): + """Reset the PR4 double-buffer ping-pong to its canonical scan-start state + under the lock. Idempotent, so the accumulator's and recon's deferred + flushes can both call it in any order without disagreeing.""" + nbuf = self.ptycho_state["num_buffers"] with self.lock: - self.ptycho_state["filled_until"] = 0 - self.ptycho_state["raw_gpu"][:] = 0 - self.ptycho_state["positions_full"][:] = 0 - self.ptycho_state["tilts_full"][:] = 0 + self.ptycho_state["filled_until"] = [0] * nbuf + self.ptycho_state["write_idx"] = 0 + self.ptycho_state["read_idx"] = 0 + self.ptycho_state["buf_free"] = [i != 0 for i in range(nbuf)] + + def _perform_flush(self): + """Reset fill levels + ping-pong and zero the GPU buffers. Only ever called + from compute(), so it is single-threaded w.r.t. the buffer writes. No-ops + when nothing has been accumulated since the last flush (free redundant + flush).""" + self._flush_requested = False + scan_state = self.ptycho_state.get("scan_state") or {} + # A full flush is a scan boundary — any carried straddle-tail is stale. + self._carry = None + if not self._dirty: + _ack_flush_and_maybe_release( + scan_state, "ptycho_accum_flushed", self.logger + ) + return + self._reset_pingpong() + for b in range(self.ptycho_state["num_buffers"]): + self.ptycho_state["raw_gpu"][b][:] = 0 + self.ptycho_state["positions_full"][b][:] = 0 + self.ptycho_state["tilts_full"][b][:] = 0 # Clear auto-centre so the new scan re-derives its own scan centre # from the first batch rather than reusing the previous scan's. self.ptycho_state["scan_center_py"] = None self.ptycho_state["scan_center_px"] = None - self.logger.info("Flushed ptychography accumulator buffers") + self._dirty = False + self.logger.info("Flushed ptychography accumulator buffers (both)") + _ack_flush_and_maybe_release(scan_state, "ptycho_accum_flushed", self.logger) + + def _try_flip(self): + """PR4: move the write cursor to the next buffer for the next projection, + iff that buffer is free (the recon has released it). Returns False when the + other buffer is still owned by the recon (we've lapped it by a full + projection) so the caller backpressures instead of overwriting live data.""" + nbuf = self.ptycho_state["num_buffers"] + with self.lock: + other = (self.ptycho_state["write_idx"] + 1) % nbuf + if not self.ptycho_state["buf_free"][other]: + return False + self.ptycho_state["buf_free"][other] = False + self.ptycho_state["write_idx"] = other + self.ptycho_state["filled_until"][other] = 0 + self.ptycho_state["raw_gpu"][other][:] = 0 + self.ptycho_state["positions_full"][other][:] = 0 + self.ptycho_state["tilts_full"][other][:] = 0 + # Per-projection auto-centering: the first batch written into the + # new buffer must derive a fresh center for that projection. + self.ptycho_state["scan_center_py"] = None + self.ptycho_state["scan_center_px"] = None + self.logger.info("Accumulator flipped to buffer %d for next projection", other) + return True def compute(self, op_input, op_output, context): + # Perform any requested flush here — single-threaded w.r.t. the buffers. + if self._flush_requested: + self._perform_flush() + + scan_state = self.ptycho_state.get("scan_state") or {} + num_projections = int(scan_state.get("num_projections", 1)) + no_frames = self.ptycho_state["no_frames"] + w = self.ptycho_state["write_idx"] + + # PR4: the current write buffer is full. For tomography, flip to the next + # buffer so we keep draining projection N+1 while the recon finalizes N on + # the read buffer. If the other buffer isn't free yet (recon hasn't + # released it — we've lapped it by a full projection), fall back to + # backpressure so frames queue upstream rather than being dropped. Single + # projection never flips. + if num_projections > 1 and self.ptycho_state["filled_until"][w] >= no_frames: + if not self._try_flip(): + if not self._lapped: # warn once per stall, not every tick + self._lapped = True + self.logger.warning( + "Recon lagging: write buffer %d full but the recon still " + "holds the other buffer — backpressuring upstream. Both " + "double-buffer slots are in use; the detector is outrunning " + "reconstruction.", w, + ) + return # lapped → clean backpressure fallback + self._lapped = False + w = self.ptycho_state["write_idx"] + + # Write a carried straddle-tail into the (freshly-flipped, empty) write + # buffer first. + if self._carry is not None and self.ptycho_state["filled_until"][w] == 0: + carry = self._carry + self._carry = None + self._accumulate(carry["images"], carry["positions"]) + w = self.ptycho_state["write_idx"] + data = op_input.receive("input") if data is None: return @@ -81,10 +222,35 @@ def compute(self, op_input, op_output, context): positions = np.asarray(data["positions"]) # (N, 4) [x, y, z, theta] batch_size = images.shape[0] - filled = self.ptycho_state["filled_until"] - if filled + batch_size > self.ptycho_state["no_frames"]: - return # buffer full, drop batch + filled = self.ptycho_state["filled_until"][w] + if filled + batch_size <= no_frames: + # Batch fits within the current projection. + self._accumulate(images, positions) + elif num_projections > 1: + # Batch straddles the projection boundary — write the head to fill this + # projection; carry the tail. Next compute sees the buffer full, flips, + # and writes the carry into the new buffer. + head = no_frames - filled + self._accumulate(images[:head], positions[:head]) + self._carry = {"images": images[head:], "positions": positions[head:]} + self.logger.info( + "Projection boundary: wrote %d frames, carrying %d to next projection", + head, batch_size - head, + ) + else: + return # single projection, buffer full → drop batch + + def _accumulate(self, images, positions): + """Preprocess a batch and write it into the pre-allocated buffers. + Extracted so the projection-boundary dispatch in compute() can call it + for a whole batch, a straddle-head, or a carried straddle-tail. + """ + batch_size = images.shape[0] + if batch_size == 0: + return + w = self.ptycho_state["write_idx"] + filled = self.ptycho_state["filled_until"][w] pty_data = self.ptycho_state["pty_data"] # H2D + crop @@ -130,11 +296,11 @@ def compute(self, op_input, op_output, context): positions_txyz = positions[:, [3, 0, 1, 2]] pos_y, pos_x = self._transform_positions(positions_txyz) - # Write into pre-allocated buffers + # Write into pre-allocated buffers (the current write buffer, PR4) new_end = filled + batch_size - self.ptycho_state["raw_gpu"][filled:new_end] = images_gpu - self.ptycho_state["positions_full"][0, 0, filled:new_end] = cp.asarray(pos_y) - self.ptycho_state["positions_full"][0, 1, filled:new_end] = cp.asarray(pos_x) + self.ptycho_state["raw_gpu"][w][filled:new_end] = images_gpu + self.ptycho_state["positions_full"][w][0, 0, filled:new_end] = cp.asarray(pos_y) + self.ptycho_state["positions_full"][w][0, 1, filled:new_end] = cp.asarray(pos_x) # Diagnostics on first batch if filled == 0: @@ -152,14 +318,17 @@ def compute(self, op_input, op_output, context): pos_y.min(), pos_y.max(), pos_x.min(), pos_x.max(), ) - # Atomically update fill counter + # Atomically update fill counter (for the current write buffer) with self.lock: - self.ptycho_state["filled_until"] = new_end + self.ptycho_state["filled_until"][w] = new_end + self._dirty = True - # Summary when buffer is full + # Summary when buffer is full. Buffers are allocated at capacity (R-6), + # so slice to the logical no_frames rather than the full buffer extent. if new_end >= self.ptycho_state["no_frames"]: - all_py = cp.asnumpy(self.ptycho_state["positions_full"][0, 0, :]) - all_px = cp.asnumpy(self.ptycho_state["positions_full"][0, 1, :]) + no_frames = self.ptycho_state["no_frames"] + all_py = cp.asnumpy(self.ptycho_state["positions_full"][w][0, 0, :no_frames]) + all_px = cp.asnumpy(self.ptycho_state["positions_full"][w][0, 1, :no_frames]) pty_model = self.ptycho_state["pty_model"] obj_h = int(pty_model.obj.sz_glo[-2]) obj_w = int(pty_model.obj.sz_glo[-1]) @@ -232,12 +401,23 @@ def _transform_positions(self, positions_txyz): # Auto-centre: capture scan centre from first batch if self.ptycho_state["scan_center_py"] is None: - halfview = self.ptycho_state["N"][0]/1.2/2 * 1e-6 * pty_model.scan.scale[0] - self.ptycho_state["scan_center_py"] = float(cp.mean(py)) + halfview - self.ptycho_state["scan_center_px"] = float(cp.mean(px)) + scan_state = self.ptycho_state.get("scan_state") or {} + projection = int(scan_state.get("current_projection", 0)) + halfview = self.ptycho_state["N"][0]/2/2 * 1e-6 * pty_model.scan.scale[0] + + pyi = cp.mean(py[0:5]) + pyf = cp.mean(py[-5:-1]) + #sign = 1 if projection % 2 == 0 else -1 + sign = 1 if pyi < pyf else -1 + + batch_size_here = py.shape[0] + self.ptycho_state["scan_center_py"] = -0.0 #float(cp.mean(py)) + sign * halfview + self.ptycho_state["scan_center_px"] = float(cp.mean(px)) #float(cp.mean(px[(batch_size_here//2):])) + self.logger.info( - "Auto-centring scan: center_py=%.6e m, center_px=%.6e m " + "Auto-centring projection %d: center_py=%.6e m, center_px=%.6e m " "(theta=%.2f°)", + projection, self.ptycho_state["scan_center_py"], self.ptycho_state["scan_center_px"], theta, @@ -279,10 +459,13 @@ def __init__( housekeeping_interval=10, publish_interval=5, reset_probe=False, + publish_folder=None, **kwargs, ): self.ptycho_state = ptycho_state self.lock = ptycho_state["lock"] + # PR3: folder for per-projection reconstruction HDF5 files (tomography). + self.publish_folder = publish_folder self.total_iterations = int(total_iterations) self.post_stream_iterations = int(post_stream_iterations) self.housekeeping_interval = int(housekeeping_interval) @@ -292,6 +475,13 @@ def __init__( self.all_data_arrived = False self.post_stream_count = 0 self.initialized_gpu = False + # Emitted exactly once per scan when the final iteration is reached; + # reset on flush so the next scan/projection can signal again. + self._completed = False + # Deferred flush (see GatherOp/accumulator): flush() sets this flag; the + # object/counter reset happens at the top of the next compute(), so it + # never races the PIE object update — race-safe for a mid-stream flush. + self._flush_requested = False # Pristine reconstruction state, snapshotted on first GPU init and # used to reset the object (and optionally the probe) on flush. self._obj_initial = None @@ -304,31 +494,91 @@ def __init__( def setup(self, spec: OperatorSpec): spec.output("output").condition(ConditionType.NONE) + # Completion signal to ControlOp (plumbing for PR2/PR3). + spec.output("complete").condition(ConditionType.NONE) def flush(self): - """Reset reconstruction state for a new scan. + """Request a FULL reset (scan end/header/start); performed at the top of + the next compute() (deferred, so it never races the PIE update).""" + self._flush_requested = True + + def _perform_advance(self): + """Per-projection reset: fresh object + iteration counters for the next + projection, probe carried over (warm start). Does NOT touch the buffers — + the read buffer was just released and the recon now reads the other one. + Called from _flip_read (PR4) when the recon moves to the next projection's + buffer, so it never re-completes the projection it just finished.""" + self.current_iteration = 0 + self.all_data_arrived = False + self.post_stream_count = 0 + self._completed = False + if self.initialized_gpu: + pty_model = self.ptycho_state["pty_model"] + pty_model.obj.array_global[:] = self._obj_initial + pty_model.obj.array_global_old[:] = self._obj_initial + pty_model.obj.array_global_kernel[:] = self._obj_kernel_initial + pty_model.obj.array_global_kernel_old[:] = self._obj_kernel_initial + if self.reset_probe: + pty_model.probe.array_states[:] = self._probe_initial + pty_model.source.flux = self._flux_initial + self.logger.info("Recon advanced to next projection (object reset, probe carried)") + + def _flip_read(self, scan_state): + """PR4: finished the current read buffer — release it back to the + accumulator, advance current_projection (the recon owns it, so the file + save can't race a ControlOp bump), move the read cursor to the next + buffer, and reset the object/counters for the next projection.""" + nbuf = self.ptycho_state["num_buffers"] + with self.lock: + r = self.ptycho_state["read_idx"] + self.ptycho_state["buf_free"][r] = True # accumulator may reuse it + self.ptycho_state["raw_gpu"][r][:] = 0 + self.ptycho_state["positions_full"][r][:] = 0 + self.ptycho_state["tilts_full"][r][:] = 0 + self.ptycho_state["filled_until"][r] = 0 + self.ptycho_state["read_idx"] = (r + 1) % nbuf + scan_state["current_projection"] = int( + scan_state.get("current_projection", 0) + ) + 1 + self._perform_advance() # fresh object/counters (probe carried) + + def _perform_flush(self): + """Reset reconstruction state for a new scan. Only ever called from + compute(), so the object reset is single-threaded w.r.t. the PIE update. Resets iteration counters and the object to its initial guess. By default the probe (and its flux) are CARRIED OVER from the previous scan as a warm start, since consecutive scans usually share illumination. Set ``reset_probe=True`` to fully reset the probe too. """ + self._flush_requested = False + scan_state = self.ptycho_state.get("scan_state") or {} + self.current_iteration = 0 + self.all_data_arrived = False + self.post_stream_count = 0 + self._completed = False + # Reset the shared ping-pong (PR4) too, so this op immediately sees "no + # data" and won't re-process the just-finished scan before the + # accumulator's own (deferred) flush zeros the buffers. Idempotent with the + # accumulator's identical reset — order-independent. + nbuf = self.ptycho_state["num_buffers"] with self.lock: - self.current_iteration = 0 - self.all_data_arrived = False - self.post_stream_count = 0 - if self.initialized_gpu: - pty_model = self.ptycho_state["pty_model"] - pty_model.obj.array_global[:] = self._obj_initial - pty_model.obj.array_global_old[:] = self._obj_initial - if self.reset_probe: - # Full reset: restore initial probe + flux so iteration 0 - # recomputes flux and re-normalises the probe. - pty_model.probe.array_states[:] = self._probe_initial - pty_model.source.flux = self._flux_initial - # else: leave the previous scan's probe and flux untouched. - # flux stays >= 0, so the iter-0 re-normalisation branch in - # compute() does not fire and the carried probe is preserved. + self.ptycho_state["filled_until"] = [0] * nbuf + self.ptycho_state["write_idx"] = 0 + self.ptycho_state["read_idx"] = 0 + self.ptycho_state["buf_free"] = [i != 0 for i in range(nbuf)] + if self.initialized_gpu: + pty_model = self.ptycho_state["pty_model"] + pty_model.obj.array_global[:] = self._obj_initial + pty_model.obj.array_global_old[:] = self._obj_initial + if self.reset_probe: + # Full reset: restore initial probe + flux so iteration 0 + # recomputes flux and re-normalises the probe. + pty_model.probe.array_states[:] = self._probe_initial + pty_model.source.flux = self._flux_initial + # else: leave the previous scan's probe and flux untouched. + # flux stays >= 0, so the iter-0 re-normalisation branch in + # compute() does not fire and the carried probe is preserved. if self.reset_probe: self.logger.info( @@ -342,26 +592,149 @@ def flush(self): "ptychography.reset_probe: true in the config." ) + _ack_flush_and_maybe_release(scan_state, "ptycho_recon_flushed", self.logger) + + # ------------------------------------------------------------------ + # Header preemption handshake (R-4) + + def _save_and_signal_complete(self, op_output): + """Persist the in-flight partial result, then emit recon_complete. + + Called at the top of compute() when a header preemption is requested, so + the current reconstruction is saved (finish-current-iteration → save) + before the geometry is reconfigured. No-ops the save when the recon has + not produced anything yet (idle / pre-first-iteration). + """ + if self.initialized_gpu and self.current_iteration > 0: + pty_data = self.ptycho_state["pty_data"] + pty_model = self.ptycho_state["pty_model"] + pty_params = self.ptycho_state["pty_params"] + # Bring the object/probe to host and run the end-of-iteration save + # (writes pty_out). We are reconfiguring next, so we do not push the + # arrays back to device (no to_device). + from_device(pty_model, pty_params) + setup.after_iteration(pty_data, pty_model, pty_params, pty_plot=None) + obj_2d = np.squeeze(cp.asnumpy(pty_model.obj.array_global)) + probe_2d = np.squeeze(cp.asnumpy(pty_model.probe.array_states)) + out = { + "object_phase": np.angle(obj_2d).astype(np.float32), + "object_amp": np.abs(obj_2d).astype(np.float32), + "probe_phase": np.angle(probe_2d).astype(np.float32), + "probe_amp": np.abs(probe_2d).astype(np.float32), + "iteration": self.current_iteration, + } + op_output.emit(out, "output") + self.logger.info( + "Saved partial result before preemption (iter %d)", + self.current_iteration, + ) + # Signal completion → ControlOp flushes all ops for the next dataset. + if not self._completed: + self._completed = True + op_output.emit("recon_complete", "complete") + + def _apply_pending_geometry(self): + """Apply the header's staged geometry while quiesced, then clear the + handshake. Safe because no PIE iteration is in flight (Phase 2).""" + from ptychography_setup import configure_scan_geometry + + with self.lock: + pending = self.ptycho_state.get("pending_geometry") + self.ptycho_state["pending_geometry"] = None + if pending is not None: + try: + configure_scan_geometry(self.ptycho_state, **pending) + self.logger.info("Applied new scan geometry from header") + except Exception: + # Never let a bad reconfigure take down the pipeline; keep the + # previous geometry and resume. (HeaderRxOp already rejects + # over-capacity grids; this guards anything unexpected.) + self.logger.exception( + "Failed to apply new scan geometry — keeping previous geometry" + ) + self.ptycho_state["needs_gpu_reinit"] = False + # A full geometry reconfigure subsumes any pending flush (the object is + # rebuilt from scratch), so drop a stale deferred flush that would + # otherwise reference the previous geometry's initial-object snapshot. + self._flush_requested = False + # Clear the handshake — next compute re-inits GPU for the new object. + self.ptycho_state["quiesced"].clear() + self.ptycho_state["preempt_requested"].clear() + + def _reset_for_new_geometry(self): + """Re-init reconstruction state after a geometry change. + + configure_scan_geometry rebuilt the object arrays on the host, so the + one-time GPU transfer and the pristine-object snapshot must be redone. + """ + self.ptycho_state["needs_gpu_reinit"] = False + self.initialized_gpu = False + self._obj_initial = None + self._probe_initial = None + self._flux_initial = None + self.current_iteration = 0 + self.all_data_arrived = False + self.post_stream_count = 0 + self._completed = False + self.logger.info("Recon reset for new scan geometry") + def compute(self, op_input, op_output, context): - # Snapshot fill level + # Header preemption handshake (R-4) takes priority over any flush so the + # in-flight partial is saved with the object still intact. + if self.ptycho_state["preempt_requested"].is_set(): + if not self.ptycho_state["quiesced"].is_set(): + # Phase 1: the current iteration is already finished (compute is + # atomic). Save the partial result + signal completion, then + # quiesce so the geometry can be re-pointed without racing a + # live PIE view. + self._save_and_signal_complete(op_output) + self.ptycho_state["quiesced"].set() + self.logger.info("Recon quiesced for header preemption") + else: + # Phase 2: still preempted and quiesced — apply the staged + # geometry now (nothing is touching the buffers) and clear the + # handshake. Next compute re-inits GPU for the new object. + self._apply_pending_geometry() + return + + # Perform any requested flush here — single-threaded w.r.t. the PIE update. + if self._flush_requested: + self._perform_flush() + + # Geometry changed while quiesced — re-init GPU state for the new object. + if self.ptycho_state.get("needs_gpu_reinit"): + self._reset_for_new_geometry() + + scan_state = self.ptycho_state.get("scan_state") or {} + num_projections = int(scan_state.get("num_projections", 1)) + no_frames = self.ptycho_state["no_frames"] + + # PR4: after the FINAL projection completes we idle until the scan-end + # flush (ControlOp flushes on recon_complete). Non-final projections reset + # _completed in _flip_read on the same tick, so this only catches the + # final-idle case — the recon never re-processes a finished projection. + if self._completed: + return + + # Snapshot fill level of the buffer we're reading (PR4 double-buffer). with self.lock: - n_filled = self.ptycho_state["filled_until"] + r = self.ptycho_state["read_idx"] + n_filled = self.ptycho_state["filled_until"][r] if n_filled == 0: return - # ITER_TIMING instrumentation (diagnostic, uncommitted) + # ITER_TIMING instrumentation (per-iteration timing diagnostic; INFO level, + # ~sub-ms against a ~450ms PIE iteration) t_start = time.perf_counter() - no_frames = self.ptycho_state["no_frames"] - # Detect when all data has arrived if n_filled >= no_frames and not self.all_data_arrived: self.all_data_arrived = True self.post_stream_count = 0 pty_model = self.ptycho_state["pty_model"] pty_model.scan.original = cp.copy( - self.ptycho_state["positions_full"] + self.ptycho_state["positions_full"][r] ) self.logger.info("All %d frames arrived", no_frames) @@ -385,30 +758,36 @@ def compute(self, op_input, op_output, context): pty_model = self.ptycho_state["pty_model"] pty_params = self.ptycho_state["pty_params"] - pty_model.scan.positions = self.ptycho_state["positions_full"][ + pty_model.scan.positions = self.ptycho_state["positions_full"][r][ :, :, :n_filled ] - pty_model.scan.tilts = self.ptycho_state["tilts_full"][ + pty_model.scan.tilts = self.ptycho_state["tilts_full"][r][ :, :, :n_filled ] - pty_data.raw_expanded = self.ptycho_state["raw_gpu"][:n_filled][ + pty_data.raw_expanded = self.ptycho_state["raw_gpu"][r][:n_filled][ cp.newaxis, :, :, : ] # Flux normalization — compute once on first iteration - if self.current_iteration == 0 and pty_model.source.flux < 0: - raw_cpu = cp.asnumpy(self.ptycho_state["raw_gpu"][:n_filled]) + if self.current_iteration == 0: + raw_cpu = cp.asnumpy(self.ptycho_state["raw_gpu"][r][:n_filled]) dp = pty_data.dp - pty_model.source.flux = float(np.sum( - np.sum(raw_cpu, 0)[dp == 1] - ) / raw_cpu.shape[0]) - self.logger.info("Computed flux = %.2f from %d frames", pty_model.source.flux, n_filled) - for trial_idx in range(pty_model.scan.tris_n): - pty_model.probe.array_states[:, :, :, :, trial_idx, :, :] = setPower( - pty_model.probe.array_states[:, :, :, :, trial_idx, :, :], - pty_model.source.flux, - ) - self.logger.info("Probe power normalized to flux") + + if not pty_model.source.flux == -2: + + if pty_model.source.flux < 0: + pty_model.source.flux = float(np.sum( + np.sum(raw_cpu, 0)[dp == 1] + ) / raw_cpu.shape[0]) + self.logger.info("Computed flux = %.2f from %d frames", pty_model.source.flux, n_filled) + + + for trial_idx in range(pty_model.scan.tris_n): + pty_model.probe.array_states[:, :, :, :, trial_idx, :, :] = setPower( + pty_model.probe.array_states[:, :, :, :, trial_idx, :, :], + pty_model.source.flux, + ) + self.logger.info("Probe power normalized to flux") pty_params.current_iteration = cp.asarray(min( self.current_iteration, self.total_iterations - 1 @@ -452,8 +831,8 @@ def compute(self, op_input, op_output, context): "No valid positions (0/%d in object bounds), skipping iteration", n_filled, ) - pty_model.scan.positions = self.ptycho_state["positions_full"] - pty_model.scan.tilts = self.ptycho_state["tilts_full"] + pty_model.scan.positions = self.ptycho_state["positions_full"][r] + pty_model.scan.tilts = self.ptycho_state["tilts_full"][r] return pty_params.frame_IDs = cp.asnumpy(valid_ids) @@ -478,15 +857,21 @@ def compute(self, op_input, op_output, context): combine_subsets_stream(pty_model, pty_params, recon_data) t_pie = time.perf_counter() - # Restore full buffer references for next accumulator writes - pty_model.scan.positions = self.ptycho_state["positions_full"] - pty_model.scan.tilts = self.ptycho_state["tilts_full"] + # Restore full (read-buffer) references after the sliced PIE view + pty_model.scan.positions = self.ptycho_state["positions_full"][r] + pty_model.scan.tilts = self.ptycho_state["tilts_full"][r] - # Housekeeping (every N iterations or on last) - is_last = self.current_iteration >= self.total_iterations - 1 and ( - not self.all_data_arrived - or self.post_stream_count >= self.post_stream_iterations - ) + # Housekeeping (every N iterations or on last). For tomography + # (num_projections > 1) a projection completes as soon as all its frames + # are in — finish this in-flight iteration, then advance (plan PR3). For a + # single projection, run to total_iterations + post_stream (as before). + if num_projections > 1: + is_last = self.all_data_arrived + else: + is_last = self.current_iteration >= self.total_iterations - 1 and ( + not self.all_data_arrived + or self.post_stream_count >= self.post_stream_iterations + ) if ( self.current_iteration % self.housekeeping_interval == 0 or is_last @@ -509,6 +894,51 @@ def compute(self, op_input, op_output, context): } op_output.emit(out, "output") + # Completion signal — emitted once, only when the scan is GENUINELY + # complete: all frames have arrived AND the final (post-stream) iteration + # is done. Gating on all_data_arrived is essential — without it, a low + # total_iterations exhausts mid-stream and fires "complete" while frames + # are still arriving, which would flush GatherOp mid-compute (race) and + # reset the object before the full scan is reconstructed. ControlOp + # flushes on this signal (Task 3: flush after the last iteration). + projection_advanced = False + if is_last and self.all_data_arrived and not self._completed: + self._completed = True + if num_projections > 1: + # Tomography: save this projection using the CURRENT index, then + # either end the scan (final) or flip to the next projection's + # buffer. PR4: the recon owns current_projection and the read-buffer + # flip itself — no ControlOp round-trip — so it can never save the + # next projection under a stale index (the old race). + self._save_projection_file() + current_proj = int(scan_state.get("current_projection", 0)) + if current_proj >= num_projections - 1: + # Final projection → full scan end. ControlOp flushes on this. + op_output.emit("recon_complete", "complete") + self.logger.info( + "Projection %d/%d complete (final) at iteration %d", + current_proj, num_projections, self.current_iteration, + ) + # _completed stays True → idle until the scan-end flush. + else: + # Release the finished read buffer to the accumulator and move + # to the next projection's buffer (which the accumulator has + # been filling meanwhile). Resets _completed → resume next tick. + iter_done = self.current_iteration # _flip_read resets it to 0 + self._flip_read(scan_state) + projection_advanced = True + self.logger.info( + "Projection %d/%d complete at iteration %d — flipped read " + "buffer to %d for next projection", + current_proj, num_projections, iter_done, + self.ptycho_state["read_idx"], + ) + else: + op_output.emit("recon_complete", "complete") + self.logger.info( + "Reconstruction complete at iteration %d", self.current_iteration + ) + t_end = time.perf_counter() self.logger.info( "ITER_TIMING iter=%d n_filled=%d valid=%d total_ms=%.1f " @@ -529,10 +959,42 @@ def compute(self, op_input, op_output, context): n_filled, no_frames, ) - self.current_iteration += 1 + + if not projection_advanced: + self.current_iteration += 1 # ------------------------------------------------------------------ + def _save_projection_file(self): + """PR3: write this projection's reconstruction to its own HDF5 file, + named with the shared series_id + projection index (M2 — all projections + of a tomography scan share one series_id, so the index is mandatory).""" + if self.publish_folder is None: + return + scan_state = self.ptycho_state.get("scan_state") or {} + series_id = scan_state.get("series_id", "unknown") + proj = int(scan_state.get("current_projection", 0)) + pty_model = self.ptycho_state["pty_model"] + obj_2d = np.squeeze(cp.asnumpy(pty_model.obj.array_global)) + probe_2d = np.squeeze(cp.asnumpy(pty_model.probe.array_states)) + import h5py + try: + os.makedirs(self.publish_folder, exist_ok=True) + path = os.path.join( + self.publish_folder, f"{series_id}_proj{proj:02d}_recon.h5" + ) + with h5py.File(path, "w") as f: + f.create_dataset("object_phase", data=np.angle(obj_2d).astype(np.float32)) + f.create_dataset("object_amp", data=np.abs(obj_2d).astype(np.float32)) + f.create_dataset("probe_phase", data=np.angle(probe_2d).astype(np.float32)) + f.create_dataset("probe_amp", data=np.abs(probe_2d).astype(np.float32)) + f.attrs["projection"] = proj + f.attrs["series_id"] = str(series_id) + f.attrs["iteration"] = int(self.current_iteration) + self.logger.info("Saved projection recon → %s", path) + except Exception: + self.logger.exception("Failed to save projection recon file") + def _init_gpu(self): """One-time transfer of static model data to GPU.""" pty_data = self.ptycho_state["pty_data"] @@ -550,6 +1012,7 @@ def _init_gpu(self): # can reset the object back to its initial guess on flush. The probe # snapshot is only used when reset_probe is enabled. self._obj_initial = pty_model.obj.array_global.copy() + self._obj_kernel_initial = pty_model.obj.array_global_kernel.copy() self._probe_initial = pty_model.probe.array_states.copy() self._flux_initial = pty_model.source.flux # may be < 0 (auto) diff --git a/pipeline/ptychography_setup.py b/pipeline/ptychography_setup.py index 77e6947..6dd6e37 100644 --- a/pipeline/ptychography_setup.py +++ b/pipeline/ptychography_setup.py @@ -1,11 +1,15 @@ """ Ptychography State Initialization -Builds the shared ptycho_state dict at application launch by: -1. Loading PtyREX model from JSON config -2. Computing scan extent from npoints and step_size (matching PtyREX streaming) -3. Running one-time PtyREX setup (mirrors pre_process_reconstruct_stream) -4. Pre-allocating GPU buffers +Builds the shared ptycho_state dict at application launch, split into a +grid-INDEPENDENT one-time model load and a grid-DEPENDENT geometry +configuration so scan geometry can be reconfigured on the fly from a live +header (PR2) without reallocating GPU buffers (R-6): + +1. ``load_ptycho_model`` — load PtyREX model, jsplitter, detector pre-load +2. GPU buffers allocated once at the configured MAX capacity (never realloced) +3. ``configure_scan_geometry`` — grid-dependent object sizing + view re-pointing +4. ``init_ptycho_state`` — load + allocate-at-max + one default configure """ import threading @@ -54,12 +58,24 @@ def update(self, *args, **kwargs): pass -def init_ptycho_state(ptycho_cfg: dict) -> dict: - """Build ptycho_state from PtyREX JSON config + pipeline YAML overrides. +def _ensure_ptyrex_on_path(): + """Make the vendored PtyREX package importable (mirrors the original + lazy sys.path insertion so the host can import this module without PtyREX).""" + import importlib.util, os, sys + + _spec = importlib.util.find_spec("ptyrex") + if _spec and _spec.origin: + _ptyrex_root = os.path.dirname(os.path.dirname(_spec.origin)) + if _ptyrex_root not in sys.path: + sys.path.insert(0, _ptyrex_root) + + +def load_ptycho_model(ptycho_cfg: dict): + """Grid-INDEPENDENT one-time load. - The only scan parameters required are npoints_h, npoints_v, step_size_h, - and step_size_v — the scan extent in pixels (R) and object size are - derived automatically, matching PtyREX's streaming workflow. + Loads the PtyREX model from JSON, installs the single-rank jsplitter stub, + and runs the detector pre-load (crop geometry). Nothing here depends on the + scan grid, so it runs exactly once at startup. Parameters ---------- @@ -68,23 +84,11 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: Returns ------- - dict - Shared state containing PtyREX model objects and pre-allocated - GPU buffers. + (pty_data, pty_model, pty_params, H, W) + PtyREX objects plus the cropped detector frame size (H, W). """ - import importlib.util, os, sys - import cupy as cp - - _spec = importlib.util.find_spec("ptyrex") - if _spec and _spec.origin: - _ptyrex_root = os.path.dirname(os.path.dirname(_spec.origin)) - if _ptyrex_root not in sys.path: - sys.path.insert(0, _ptyrex_root) - + _ensure_ptyrex_on_path() from ptyrex.core.io import json_read - from ptyrex.reconstruct.core import setup - from ptyrex.reconstruct.iterator.process_pty_model import generate_grow_scan_params - from ptyrex.reconstruct.utils import numpy as utils_np # ── 1. Load PtyREX model from JSON config ────────────────────────── ptyrex_config_path = ptycho_cfg["ptyrex_config"] @@ -92,25 +96,6 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: ID = ptycho_cfg.get("ID", [1, 1, 1]) pty_data, pty_model, pty_params = json_read.load(ptyrex_config_path, scan_ID, ID) - # ── 2. Compute streaming parameters from npoints / step_size ─────── - npoints_h = ptycho_cfg["npoints_h"] - npoints_v = ptycho_cfg["npoints_v"] - step_size_h = ptycho_cfg["step_size_h"] - step_size_v = ptycho_cfg["step_size_v"] - - no_frames = npoints_h * npoints_v - # Scan extent in microns with 20% padding (same formula as PtyREX streaming) - N = [ - ((npoints_v - 1) * step_size_v) * 1.2, - ((npoints_h - 1) * step_size_h) * 1.2, - ] - logger.info( - "Scan: %d x %d points, step %.3f x %.3f µm → " - "N = [%.2f, %.2f] µm, %d frames", - npoints_h, npoints_v, step_size_h, step_size_v, - N[1], N[0], no_frames, - ) - pty_params.total_iterations = ptycho_cfg["total_iterations"] # Ensure string attributes expected by PtyREX save/config routines @@ -118,13 +103,76 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: pty_data.ID = str(pty_data.ID[0]) if isinstance(pty_data.ID, list) else str(pty_data.ID) pty_data.scan_ID = str(pty_data.scan_ID[0]) if isinstance(pty_data.scan_ID, list) else str(pty_data.scan_ID) - # ── 3. Dummy jsplitter for single-rank pipeline ──────────────────── + # ── Dummy jsplitter for single-rank pipeline ─────────────────────── pty_params.jsplitter = DummyJSplitter() - # ── 4. Initialise scan arrays (mirrors pre_process_reconstruct_stream) ─ + # ── Detector pre-load (crop geometry — grid-independent) ─────────── pty_data.pre_load(pty_model.detector, pty_params) H = int(pty_data.crop_bottom - pty_data.crop_top) W = int(pty_data.crop_right - pty_data.crop_left) + + return pty_data, pty_model, pty_params, H, W + + +def configure_scan_geometry( + ptycho_state: dict, + npoints_h: int, + npoints_v: int, + step_size_h: float, + step_size_v: float, +): + """Grid-DEPENDENT (re)configuration of scan geometry. + + Recomputes ``no_frames`` and the scan extent ``N``, re-runs the PtyREX + object-sizing setup, refreshes the detector pixel mask / dp transforms, and + re-points the scan arrays at the pre-allocated GPU buffers. **No GPU + reallocation** (R-6): a grid whose ``no_frames`` exceeds the configured + capacity is rejected with ``ValueError``. + + Safe to call at startup (from ``init_ptycho_state``) and again on a live + header, but ONLY while the reconstruction is quiesced (R-4 handshake) — it + re-points views the recon reads during a PIE iteration. + + Sets ``ptycho_state["needs_gpu_reinit"] = True`` so the reconstruction op + re-runs its one-time GPU transfer and re-snapshots the pristine object for + the new geometry. + """ + _ensure_ptyrex_on_path() + from ptyrex.reconstruct.core import setup + from ptyrex.reconstruct.iterator.process_pty_model import generate_grow_scan_params + from ptyrex.reconstruct.utils import numpy as utils_np + import cupy as cp + + pty_data = ptycho_state["pty_data"] + pty_model = ptycho_state["pty_model"] + pty_params = ptycho_state["pty_params"] + H = ptycho_state["H"] + W = ptycho_state["W"] + capacity = ptycho_state["capacity"] + + # ── Compute streaming parameters from npoints / step_size ────────── + no_frames = int(npoints_h) * int(npoints_v) + if no_frames > capacity: + raise ValueError( + f"Requested grid {npoints_h}x{npoints_v} = {no_frames} frames exceeds " + f"the pre-allocated capacity of {capacity} frames " + f"(increase max_npoints_h/max_npoints_v in the config). " + f"Buffers are never reallocated at runtime (R-6)." + ) + + # Scan extent in microns with 20% padding (same formula as PtyREX streaming) ## increased it on 11/09/26 to 100% padding for cases with position overshoots + N = [ + ((npoints_v - 1) * step_size_v) * 2, + ((npoints_h - 1) * step_size_h) * 2, + ] + logger.info( + "Configuring scan: %d x %d points, step %.3f x %.3f µm → " + "N = [%.2f, %.2f] µm, %d frames (capacity %d)", + npoints_h, npoints_v, step_size_h, step_size_v, + N[1], N[0], no_frames, capacity, + ) + + # ── Initialise scan arrays (mirrors pre_process_reconstruct_stream) ─ pty_data.raw = np.zeros((no_frames, H, W), dtype=np.uint32) pty_model.scan.positions = np.ones([no_frames, 2], np.float32) @@ -144,13 +192,15 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: pty_model.scan.sz = [pty_model.scan.positions.shape[0], 1] pty_data.reg_ind = pty_model.scan.reg_ind - # ── 5. PtyREX one-time setup — pass N in microns (not R in pixels) ─ + # ── PtyREX one-time setup — pass N in microns (not R in pixels) ──── pty_plot = DummyPtyPlot() pty_data, pty_model, pty_params, pty_plot = setup.before_reconstruction_stream( pty_data, pty_model, pty_params, pty_plot, N ) - # ── 6. Pixel mask + dp transforms (mirrors post_process_stream) ── + # ── Pixel mask + dp transforms (mirrors post_process_stream) ─────── + # Re-run every reconfigure so ordering matches the validated single-scan + # path (before_reconstruction_stream then get_pixel_mask); cheap. df, ff, dp = pty_data.get_pixel_mask(pty_model.detector, pty_params) pty_data.dp = dp[ pty_data.crop_top : pty_data.crop_bottom, @@ -164,10 +214,9 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: pty_model.detector.mask[pty_data.dp > 0] = 0 pty_model.detector.mask_inv[pty_data.dp > 0] = 1 - # Now apply the same dp transforms that post_process_stream does on - # its first iteration so that dp matches the preprocessed data layout - # and uses the inverted convention expected by the flux computation - # (dp == 1 → good pixel). + # Apply the same dp transforms that post_process_stream does on its first + # iteration so dp matches the preprocessed data layout and uses the inverted + # convention expected by the flux computation (dp == 1 → good pixel). det = pty_model.detector if det.orientation == "01": pty_data.dp = pty_data.dp[:, ::-1] @@ -187,36 +236,152 @@ def init_ptycho_state(ptycho_cfg: dict) -> dict: pty_model.obj.array_global_old[:] = pty_model.obj.array_global[:] generate_grow_scan_params(pty_params) - # ── 7. Pre-allocate GPU buffers ──────────────────────────────────── - raw_gpu = cp.zeros((no_frames, H, W), dtype=cp.uint32) - positions_full = cp.zeros((1, 2, no_frames), dtype=cp.float32) - tilts_full = cp.zeros((1, 2, no_frames), dtype=cp.float32) - - pty_model.scan.positions = positions_full - pty_model.scan.tilts = tilts_full - pty_model.scan.original = cp.zeros_like(positions_full) - pty_model.scan.previous = cp.zeros_like(positions_full) + # ── Re-point scan arrays at the pre-allocated GPU buffers ────────── + # The buffers are allocated once at capacity in init_ptycho_state and never + # realloced; the accumulator/recon index them via ptycho_state directly and + # bound their reads to [:n_filled] (n_filled <= no_frames <= capacity). + # PR4: point the model at buffer 0 to start; the recon re-points these views + # to positions_full[read_idx] / raw_gpu[read_idx] each compute, so this is + # just the initial binding + the shape source for scan.original/previous. + positions_full = ptycho_state["positions_full"] + tilts_full = ptycho_state["tilts_full"] + pty_model.scan.positions = positions_full[0] + pty_model.scan.tilts = tilts_full[0] + pty_model.scan.original = cp.zeros_like(positions_full[0]) + pty_model.scan.previous = cp.zeros_like(positions_full[0]) + + # ── Update shared state ──────────────────────────────────────────── + # A (re)configure is a clean scan boundary: reset the ping-pong to buffer 0, + # both fill levels empty, buffer 0 claimed for the first projection. + with ptycho_state["lock"]: + ptycho_state["no_frames"] = no_frames + nbuf = ptycho_state["num_buffers"] + ptycho_state["filled_until"] = [0] * nbuf + ptycho_state["write_idx"] = 0 + ptycho_state["read_idx"] = 0 + ptycho_state["buf_free"] = [i != 0 for i in range(nbuf)] + ptycho_state["N"] = N + # Clear auto-centre so the new geometry re-derives its own scan centre. + ptycho_state["scan_center_py"] = None + ptycho_state["scan_center_px"] = None + # Signal the reconstruction op to re-init GPU state for the new object size. + ptycho_state["needs_gpu_reinit"] = True + + # Mirror the frame count into the always-present scan_state holder (S11) so + # the STXM path and header op can read it even when ptycho is disabled. + scan_state = ptycho_state.get("scan_state") + if scan_state is not None: + scan_state["no_frames"] = no_frames + scan_state["npoints_h"] = int(npoints_h) + scan_state["npoints_v"] = int(npoints_v) + scan_state["step_size_h"] = float(step_size_h) + scan_state["step_size_v"] = float(step_size_v) logger.info( - "ptycho_state initialized: %d frames, image size %dx%d, " - "object size %s, %d total iterations", + "Scan geometry configured: %d frames, image size %dx%d, object size %s", no_frames, H, W, tuple(int(x) for x in pty_model.obj.sz_glo), - pty_params.total_iterations, ) - # ── 8. Assemble ptycho_state ─────────────────────────────────────── - return { + +def init_ptycho_state(ptycho_cfg: dict, scan_state: dict = None) -> dict: + """Build ptycho_state: load the model, allocate GPU buffers at the configured + MAX capacity, then configure a default scan geometry so the pipeline can run + before any header arrives. + + Parameters + ---------- + ptycho_cfg : dict + The ``ptychography`` section of the pipeline YAML config. Must provide + ``max_npoints_h``/``max_npoints_v`` (buffer capacity + default grid) and + ``default_step_size_h``/``default_step_size_v`` (startup step sizes). + scan_state : dict, optional + The always-present shared holder for projection/frame counts (S11). Its + ``no_frames`` is populated by the default configure below. + + Returns + ------- + dict + Shared state containing PtyREX model objects, pre-allocated GPU buffers + (at max capacity), geometry, and the preemption handshake primitives. + """ + import cupy as cp + + # ── 1. Grid-independent model load ───────────────────────────────── + pty_data, pty_model, pty_params, H, W = load_ptycho_model(ptycho_cfg) + + # ── 2. Capacity + default startup grid ───────────────────────────── + max_npoints_h = int(ptycho_cfg["max_npoints_h"]) + max_npoints_v = int(ptycho_cfg["max_npoints_v"]) + capacity = max_npoints_h * max_npoints_v + default_step_h = float(ptycho_cfg["default_step_size_h"]) + default_step_v = float(ptycho_cfg["default_step_size_v"]) + + # ── 3. Pre-allocate GPU buffers ONCE at max capacity (R-6) ───────── + # PR4 double-buffering: TWO buffer sets (ping-pong). While the recon + # finalizes projection N on the read buffer, the accumulator fills + # projection N+1 into the write buffer — so the accumulator never stops + # draining and a free-running detector is never backpressured across the + # ~one-iteration finalize window. Single-projection scans always use + # buffer 0 (no flip). Cost: 2× raw_gpu (~tens of MB). + NUM_BUFFERS = 2 + raw_gpu = [cp.zeros((capacity, H, W), dtype=cp.uint32) for _ in range(NUM_BUFFERS)] + positions_full = [cp.zeros((1, 2, capacity), dtype=cp.float32) for _ in range(NUM_BUFFERS)] + tilts_full = [cp.zeros((1, 2, capacity), dtype=cp.float32) for _ in range(NUM_BUFFERS)] + + logger.info( + "ptycho_state buffers allocated at capacity: %d frames × %d buffers " + "(double-buffered), image size %dx%d, %d total iterations", + capacity, NUM_BUFFERS, H, W, pty_params.total_iterations, + ) + + # ── 4. Assemble ptycho_state (geometry filled in by configure) ───── + ptycho_state = { "pty_data": pty_data, "pty_model": pty_model, "pty_params": pty_params, - "raw_gpu": raw_gpu, - "positions_full": positions_full, - "tilts_full": tilts_full, - "filled_until": 0, - "no_frames": no_frames, + "raw_gpu": raw_gpu, # list[NUM_BUFFERS] of (capacity,H,W) + "positions_full": positions_full, # list[NUM_BUFFERS] + "tilts_full": tilts_full, # list[NUM_BUFFERS] + "num_buffers": NUM_BUFFERS, + # PR4 ping-pong state (all touched under "lock"): + # filled_until[i] — fill level of buffer i + # write_idx — buffer the accumulator writes (accumulator owns) + # read_idx — buffer the recon reads (recon owns) + # buf_free[i] — buffer i holds no data the recon still needs, so the + # accumulator may claim it for a new projection. Init: + # buffer 0 is claimed for the first projection. + "filled_until": [0] * NUM_BUFFERS, + "write_idx": 0, + "read_idx": 0, + "buf_free": [i != 0 for i in range(NUM_BUFFERS)], + "no_frames": 0, # set by configure_scan_geometry + "H": H, + "W": W, + "capacity": capacity, "scan_center_py": None, "scan_center_px": None, - "N": N, + "N": None, # set by configure_scan_geometry "lock": threading.Lock(), + "scan_state": scan_state, + # Preemption handshake (R-4): header stages pending_geometry + sets + # preempt_requested; recon saves the partial, sets quiesced, then applies + # the geometry while quiesced and clears the flags. + "preempt_requested": threading.Event(), + "quiesced": threading.Event(), + "pending_geometry": None, + "needs_gpu_reinit": False, } + + # ── 5. Configure the default (startup) scan geometry ─────────────── + configure_scan_geometry( + ptycho_state, + npoints_h=max_npoints_h, + npoints_v=max_npoints_v, + step_size_h=default_step_h, + step_size_v=default_step_v, + ) + # First configure just did startup init; no reconfigure has happened yet. + ptycho_state["needs_gpu_reinit"] = False + + return ptycho_state diff --git a/pipeline/publish.py b/pipeline/publish.py index ca64e09..5c5856e 100644 --- a/pipeline/publish.py +++ b/pipeline/publish.py @@ -80,6 +80,7 @@ def __init__(self, fragment, *args, publish_backend: PublishBackend = None, backend: str = "nats", backend_endpoint: str = None, + scan_state: dict = None, **kwargs): """ Initialize sink and publish operator. @@ -100,10 +101,19 @@ def __init__(self, fragment, *args, self.processed_frame_count = 0 self.processed_batch_count = 0 - # In-memory accumulator of per-batch arrays for the current scan. + # In-memory accumulator of per-batch arrays for the current scan/projection. # Avoids per-batch HDF5 open/close on the compute hot path; the file - # is written once at scan end (processing_end). + # is written once at scan end (or per projection for tomography). self.scan_buffer = [] + # Save-state tracking so flush never discards an unwritten scan buffer. + self._written = False + self._series_id = None + # PR3 tomography: shared holder (num_projections, no_frames) + a LOCAL + # per-projection index/counter so the STXM path segments itself (S2), + # independent of the ptycho path's current_projection. + self.scan_state = scan_state + self._projection = 0 + self._proj_frame_count = 0 self.publish_folder = publish_folder self.publish_tensors = publish_tensors if publish_tensors is not None else [] @@ -116,7 +126,6 @@ def __init__(self, fragment, *args, def setup(self, spec: OperatorSpec): spec.input("input").connector(IOSpec.ConnectorType.DOUBLE_BUFFER, capacity=128).condition(ConditionType.NONE) - spec.output("processing_end").condition(ConditionType.NONE) def write_scan_file(self, series_id): """Write the buffered scan to a single HDF5 file. @@ -135,13 +144,58 @@ def write_scan_file(self, series_id): with h5py.File(filepath, 'w') as f: f.create_dataset('stxm', data=data) self.logger.info(f"Wrote {data.shape[0]} frames to {filepath}") + # Self-contained: mark saved and clear the buffer so a later flush has + # nothing to discard. + self._written = True + self.scan_buffer = [] + + def _write_projection_file(self, series_id, projection): + """PR3: write the current projection's buffered STXM to its own file, + named with the shared series_id + projection index (M2), then clear the + buffer for the next projection.""" + if self.publish_folder is None or not self.scan_buffer: + return + os.makedirs(self.publish_folder, exist_ok=True) + data = np.concatenate(self.scan_buffer, axis=0) + filepath = os.path.join(self.publish_folder, f"{series_id}_proj{projection:02d}.h5") + with h5py.File(filepath, 'w') as f: + f.create_dataset('stxm', data=data) + f.attrs['projection'] = projection + f.attrs['series_id'] = str(series_id) + self.logger.info(f"Wrote projection {projection} ({data.shape[0]} frames) to {filepath}") + self._written = True + self.scan_buffer = [] def flush(self): - """Reset counters and discard any buffered scan data on flush.""" + """Reset counters. If the buffer still holds data that was never written + (a flush arrived before the end-of-scan write), write it out first so a + scan's STXM file is never lost.""" + if (self.scan_buffer and not self._written + and self.publish_folder is not None and self._series_id is not None): + self.logger.warning( + "Flush with %d unwritten STXM batch(es) — writing before clearing", + len(self.scan_buffer), + ) + # Name with the projection index if we're mid-tomography. + scan_state = self.scan_state or {} + if int(scan_state.get("num_projections", 1)) > 1: + self._write_projection_file(self._series_id, self._projection) + else: + self.write_scan_file(self._series_id) self.processed_frame_count = 0 self.processed_batch_count = 0 self.scan_buffer = [] - + self._written = False + self._projection = 0 + self._proj_frame_count = 0 + + def _publish_stxm_flush(self): + """Notify downstream consumers that the current STXM projection is done.""" + if self.backend is None: + return + import numpy as np + self.backend.publish("stxm_flush", np.array([1])) + def compute(self, op_input, op_output, context): """Receive, publish, and save processed data using metadata.""" # Initialize backend on first call @@ -200,23 +254,49 @@ def compute(self, op_input, op_output, context): if self.publish_folder is not None: if len(arrays_to_publish) > 0: self.scan_buffer.append(np.concatenate(arrays_to_publish, axis=1)) + self._written = False # new unsaved data + self._series_id = series_id # remembered for a defensive flush-save self.processed_batch_count += 1 self.processed_frame_count += tensor.shape[0] - # Check if processing is complete using metadata from upstream - if self.processed_frame_count == series_frame_count: - # Write the whole scan once, now that no more batches are arriving - if self.publish_folder is not None and series_id is not None: - self.write_scan_file(series_id) - - op_output.emit("processing_end", "processing_end") - - _n = self.processed_frame_count - _b = self.processed_batch_count - _elapsed = time.time() - series_start_time if series_start_time > 0 else 0 - _rate = _n/_elapsed if _elapsed > 0 else 0 - self.logger.info(f"{_n} processed in {_elapsed:.1f}s. speed: {_rate:.1f} Hz (in {_b} batches)") + # Share series_id so the ptycho per-projection recon files can be named + # with the same identifier (PR3). + if self.scan_state is not None and series_id is not None: + self.scan_state["series_id"] = series_id + + scan_state = self.scan_state or {} + num_projections = int(scan_state.get("num_projections", 1)) + proj_no_frames = int(scan_state.get("no_frames", 0)) + + if num_projections > 1 and proj_no_frames > 0: + # Tomography (S2): save one STXM file per projection, segmented by + # frame count — self-contained, no ControlOp round-trip. Uses >= with + # a count carry (I1); exact frame boundaries require no_frames to be a + # multiple of batch_size (a few overshoot frames otherwise land in the + # current projection's file). + self._proj_frame_count += tensor.shape[0] + if self._proj_frame_count >= proj_no_frames: + if self.publish_folder is not None and series_id is not None: + self._write_projection_file(series_id, self._projection) + + # Notify the visualizer / downstream consumers that the + # previous projection is complete and should be cleared. + self._publish_stxm_flush() + self._proj_frame_count -= proj_no_frames # carry overshoot count + self._projection += 1 + else: + # Single scan: write the whole series once (existing behaviour). Use + # >= (not ==) so a batch overshooting the exact count still triggers. + if series_frame_count > 0 and self.processed_frame_count >= series_frame_count: + if self.publish_folder is not None and series_id is not None: + self.write_scan_file(series_id) + + _n = self.processed_frame_count + _b = self.processed_batch_count + _elapsed = time.time() - series_start_time if series_start_time > 0 else 0 + _rate = _n/_elapsed if _elapsed > 0 else 0 + self.logger.info(f"{_n} processed in {_elapsed:.1f}s. speed: {_rate:.1f} Hz (in {_b} batches)") class PublishToCloudOp(Operator): @@ -228,17 +308,20 @@ class PublishToCloudOp(Operator): NOTE: STXM saving now happens in SinkAndPublishOp.write_scan_file (a single end-of-scan write), so the per-batch temp files this op consolidated are no - longer produced. It is retained for any external/temp-file workflow and - no-ops gracefully when no temp file is present. + longer produced. It is retained for a possible future external/temp-file + workflow and no-ops gracefully when no temp file is present. + + NOT CURRENTLY WIRED into the pipeline (see pipeline.py). To re-enable, feed a + completion trigger into its "trigger" input (e.g. from SinkAndPublishOp). """ - + def __init__(self, fragment, publish_folder: str = None, temp_folder: str = None, *args, **kwargs): """ Initialize cloud publishing operator. - + Args: fragment: Holoscan fragment publish_folder: Final destination folder @@ -256,14 +339,14 @@ def compute(self, op_input, op_output, context): """Consolidate and publish dataset on trigger using metadata.""" # Receive trigger - metadata is automatically merged trigger = op_input.receive("trigger") - + if trigger == "processing_end": if self.publish_folder is None or self.temp_folder is None: return # Get the series ID from metadata that flowed from upstream series_id = self.metadata.get("series_id") - + if series_id is None: self.logger.warning("No series_id found in metadata, cannot publish") return @@ -299,7 +382,7 @@ def compute(self, op_input, op_output, context): f.attrs[key] = value self.logger.info(f"Published concatenated data to {publish_file}") - + # Remove temp file os.remove(temp_file) diff --git a/pipeline/vis/vis_stxm.py b/pipeline/vis/vis_stxm.py index 41ddbde..19f2e70 100644 --- a/pipeline/vis/vis_stxm.py +++ b/pipeline/vis/vis_stxm.py @@ -225,19 +225,53 @@ def build_combined_figure(): placeholder = np.zeros((64, 64)) * np.nan ptycho_axes_info = [ - (ax_obj_phase, "twilight", "object_phase"), + (ax_obj_phase, "gray", "object_phase"), (ax_obj_amp, "gray", "object_amp"), - (ax_prb_phase, "twilight", "probe_phase"), + (ax_prb_phase, "gray", "probe_phase"), (ax_prb_amp, "gray", "probe_amp"), ] ptycho_ims = {} for ax, cmap, key in ptycho_axes_info: im = ax.imshow(placeholder, cmap=cmap, interpolation='nearest', aspect='equal') + ax.invert_xaxis() fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) ptycho_ims[key] = im return fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims +def build_combined_figure_reduced(): + """2-row x 2-col layout: STXM top, ptycho object phase + probe modulus bottom.""" + plt.style.use('dark_background') + matplotlib.rcParams.update({'font.size': 8}) + + fig = plt.figure(figsize=(10, 10)) + gs = fig.add_gridspec(2, 2, hspace=0.35, wspace=0.3) + + ax_stxm_outer = fig.add_subplot(gs[0, 0]) + ax_stxm_inner = fig.add_subplot(gs[0, 1]) + ax_obj_phase = fig.add_subplot(gs[1, 0]) + ax_prb_amp = fig.add_subplot(gs[1, 1]) + + ax_stxm_outer.set_title("STXM Outer") + ax_stxm_inner.set_title("STXM Inner") + ax_obj_phase.set_title("Object Phase") + ax_prb_amp.set_title("Probe Amplitude") + + placeholder = np.zeros((64, 64)) * np.nan + ptycho_axes_info = [ + (ax_obj_phase, "gray", "object_phase"), + (ax_prb_amp, "gray", "probe_amp"), + ] + ptycho_ims = {} + for ax, cmap, key in ptycho_axes_info: + im = ax.imshow(placeholder, cmap=cmap, interpolation='nearest', aspect='equal') + ax.invert_xaxis() + #fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + fig.colorbar(im, ax=ax) #, fraction=0.046, pad=0.04) + ptycho_ims[key] = im + + return fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims + # ===================== Combined animation ===================== @@ -316,7 +350,12 @@ def animate_combined(i): threading.Thread(target=receive_stxm_data, args=(sub_backend,), daemon=True).start() threading.Thread(target=receive_ptycho_data, args=(sub_backend,), daemon=True).start() - fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims = build_combined_figure() + reduced_flag = True + + if reduced_flag: + fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims = build_combined_figure_reduced() + else: + fig, ax_stxm_outer, ax_stxm_inner, ptycho_ims = build_combined_figure() if all(v is not None for v in (args.xmin, args.xmax, args.ymin, args.ymax)): ax_stxm_outer.set_xlim(-args.xmax, -args.xmin) diff --git a/selun_mask_unbinned.h5 b/selun_mask_unbinned.h5 new file mode 100644 index 0000000..684c518 Binary files /dev/null and b/selun_mask_unbinned.h5 differ