From 5d246a475cd1771a1175956ead0cab48bdb51a64 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Wed, 14 Oct 2026 14:03:23 -0400 Subject: [PATCH 1/4] ENH: FireANTs registrar, tuned heart/lung presets, Fireants-vs-Greedy benchmark - Add RegisterImagesFireants (GPU rigid/affine/deformable) with configure_for_heart_ct() and configure_for_lung_ct() using the tuned, identity-initialised settings from the benchmark REPORT.md - Add Fireants-vs-Greedy benchmark, tuning and report scripts in experiments/ and tests for the new registrar - Add the loop-kit skill - Physics-informed motion: clarify lambda_physics warmup docs and the collapse warning; update Duke heart parameters and tutorial 18 - Inference classes: model_tag becomes an __init__ argument - Sort __all__ and fix ruff findings; bump the ruff pre-commit hook to v0.16.6 to match the dev environment - Add a TotalSegmentator test; refresh tutorial 13 gifs --- .agents/skills/loop-kit/SKILL.md | 142 ++++ .pre-commit-config.yaml | 2 +- .../assets/tutorial_13_combined_motion_ct.gif | 4 +- .../tutorial_13_combined_motion_usd.gif | 4 +- docs/developer/migration_next.md | 32 + .../affine_diagnostic.py | 96 +++ .../benchmark.py | 582 ++++++++++++++ .../make_report.py | 715 ++++++++++++++++++ .../tuning_compare.py | 54 ++ src/monai_physio/__init__.py | 30 +- src/monai_physio/infer_physicsnemo_base.py | 13 +- src/monai_physio/infer_physicsnemo_mgn.py | 16 +- src/monai_physio/infer_physicsnemo_mlp.py | 14 +- src/monai_physio/register_images_fireants.py | 417 ++++++++++ ...ain_physicsnemo_physics_informed_motion.py | 10 +- tests/test_register_images_fireants.py | 385 ++++++++++ tests/test_segment_chest_total_segmentator.py | 38 + .../parameters_duke_heart_physics_informed.py | 20 +- .../tutorial_13_heart_and_lung_motion.py | 4 +- ...uke_heart_physics_informed_motion_infer.py | 17 +- 20 files changed, 2546 insertions(+), 49 deletions(-) create mode 100644 .agents/skills/loop-kit/SKILL.md create mode 100644 experiments/Registration-Fireants_vs_Greedy/affine_diagnostic.py create mode 100644 experiments/Registration-Fireants_vs_Greedy/benchmark.py create mode 100644 experiments/Registration-Fireants_vs_Greedy/make_report.py create mode 100644 experiments/Registration-Fireants_vs_Greedy/tuning_compare.py create mode 100644 src/monai_physio/register_images_fireants.py create mode 100644 tests/test_register_images_fireants.py diff --git a/.agents/skills/loop-kit/SKILL.md b/.agents/skills/loop-kit/SKILL.md new file mode 100644 index 00000000..2bd5aed9 --- /dev/null +++ b/.agents/skills/loop-kit/SKILL.md @@ -0,0 +1,142 @@ +# Claude Prompt: Generate a Verified Loop Kit for Any Task + +## How to use this prompt + +- **One-off:** paste everything below the line into Claude Code at the root of the target repository. Fill in the `` block, or point it at a file (e.g. `@task_request.md`). +- **Reusable:** save the part below the line as `.claude/skills/loop-kit/SKILL.md` (project) or `~/.claude/skills/loop-kit/SKILL.md` (personal). Put this frontmatter first: + ```yaml + --- + name: loop-kit + description: Generate a bounded, verified implement/verify/review loop (task contract, prompts, subagents, skills, verifier) for a given task. Use when asked to set up a loop, agent loop, or iterative refinement for a task. Does not implement the task. + --- + ``` + Then run it with `/loop-kit `. + +The design follows the earlier note, *MONAI Physio: verified initialization and code-review loops*, which adapts Youssef Hosni's [How to Create Loops with Claude Code](https://levelup.gitconnected.com/how-to-create-loops-with-claude-code-a-practical-guide-to-agentic-automation-6f422390a143). + +--- + +## Role + +You are a **loop architect**. Your job is to turn the task below into a small, file-based kit that lets Claude Code implement it in a **bounded, verified loop**. + +**You must not implement the task itself.** You build the scaffolding, and then you stop for human approval. + + +{{Describe the task here, or reference a file. Include the goal, the repository or working area, constraints, and any known acceptance criteria.}} + + +## Autonomy: minimize interrupts (applies to the kit and to the run) + +The loop is **automated**. A human halt is a cost; every halt must be justified. + +1. **Ask everything once, up front.** Before generating the kit, anticipate every question the run could raise and put them in ONE batch (use AskUserQuestion, or a single list): decisions and tolerances; data/mode choices; dependency installs and environment quirks (version pins, Windows/platform gaps, plugin conflicts); approval for expensive jobs; whether the supervisor may run the verifier itself; and **standing authorizations** for predictable amendments (verifier/environment fixes, baseline-relative lint rules, schema/path naming mismatches, API differences discovered in Phase 0). Do not trickle questions between steps. +2. **Standing authorization replaces per-step gates.** Record the answers in `TASK.md` and the supervisor skill. After that: the supervisor runs worker -> verifier -> reviewer -> decision and **auto-advances** to the next slice when the verifier is PASS and the reviewer has no blocking findings. [H] items are treated as confirmed under the standing authorization unless the human reserved them in the batch. Never end a turn with "waiting for acceptance". +3. **Self-resolve, then log.** Verifier/environment defects (outside the repo), schema/name mismatches, and pre-existing-lint baselines are fixed by the supervisor, re-verified under a new run id, and logged in `PROGRESS.md`. Contract-semantics changes that stay inside the pre-authorized classes are applied as amendments (re-freeze, refresh the control hashes) without asking. +4. **Halt only for true BLOCKED conditions:** repair limit exhausted; a protected-path change; a missing prerequisite the agent cannot create; an action outside the standing authorizations (commit/push/publish/messages, licensed downloads, new spend); or a decision that changes the task's intent. When halting, give one consolidated question list, not one at a time. +5. **Run to completion.** Continue through all slices and the final report in one go; report once at the end (plus a brief note on any self-resolved amendment). +6. **Make the generated kit say so:** `README.md`, `LOOP_INSTRUCTIONS.md`, and the supervisor skill must state these rules, include the auto-advance decision rule, and let the supervisor run the verifier when the human configured that. + +## Principles every generated file must follow + +1. **Roles are separate.** + - A **worker** proposes and implements. + - A **trusted verifier**, which the human controls, runs the checks and records evidence. + - A **fresh-context reviewer** judges the exact output without seeing the worker's reasoning. + - The **human** decides up front (decisions and standing authorizations); during the run the supervisor applies them and auto-advances. + + The worker never certifies its own success. +2. **Authority order:** + 1. the human-approved `TASK.md`, frozen by a hash; + 2. the verifier's `verification.json` for the current output hash; + 3. the reviewer's validated findings for the current output hash; + 4. `PROGRESS.md`, which is memory only. + + A line saying "done" or "approved" in a model-written file confers nothing. +3. **Small slices.** Split the work into ordered slices. Each slice gets its own acceptance criteria (ACs) and one worker run per attempt. Never start the next slice until the human accepts the current one. +4. **Every AC has an evidence type:** + - **[V]:** a deterministic check that the verifier executes; + - **[R]:** reviewer-confirmed from the output; + - **[H]:** human-confirmed. + + Prefer [V] wherever a check is possible. A required check that is missing, skipped, or not run counts as a failure, not a pass. +5. **Bounded repair.** + - Allow at most 3 repair cycles per slice, counted from the evidence folders rather than from conversation memory. + - Stop early on a repeated identical failure, alternating edits, a missing prerequisite, a protected-path change, or a decision that needs the human. + - An exhausted limit means `BLOCKED`, never success. +6. **Protected references.** + - The worker may never edit the contract, the verifier, the reference baselines, golden files, or tolerances. + - It may never add skips or delete assertions to make checks pass. + - Contract changes are proposed in a `contract_amendment.md` and decided by the human. +7. **No external side effects by default.** Agents never: + - commit, push, open PRs, resolve review threads, or publish anything; + - send messages, make purchases, or download licensed data; + - start expensive jobs without explicit approval. +8. **Evidence is immutable.** Each run writes `outputs//` (an output snapshot, raw logs, `verification.json`, `review.json`, `decision.md`). Closed runs are never edited. Any new edit invalidates earlier evidence. +9. **Reuse before invention.** Use the project's existing guidance (`CLAUDE.md`, `AGENTS.md`, existing skills, subagents, commands, CI, test markers). Reference them rather than duplicating them. Do not build a general-purpose orchestration framework. +10. **Prove value.** Keep a plain single-prompt version of the task as a baseline, and describe a paired comparison against it. + +## Step 1: Investigate (read-only) + +Before writing anything: +- **Find the current conventions.** Read `CLAUDE.md`, `AGENTS.md`, `README`, and the contributing docs. Look at existing `.claude/agents/`, `.claude/skills/`, `.claude/commands/`, and `.agents/skills/`, and at `.github/workflows/`, the test layout and markers, linters, type checkers, the docs build, and how the environment is set up (venvs, worktrees, GPU, data). +- **Find what the task touches.** Identify the files, APIs, and data, the existing tests, and the riskiest invariants. Examples: units, coordinate systems, numerical tolerances, security boundaries, public APIs, migrations. +- **Classify the task.** Is it code, documentation, data analysis, research or writing, infrastructure, or mixed? The classification decides what [V] checks are possible. For non-code tasks, still prefer deterministic checks: schema validation, link and citation checks, word or section limits, recomputing numbers from source data, spellcheck, build success. Where only judgment is possible, use [R] or [H] and say so explicitly in `TASK.md`. +- **Check the installed tools.** Record the versions of the Claude Code CLI, Python, and the test tools. **Check the installed Claude Code** (`claude --help`, and the docs for subagents, skills, hooks, `/goal`, headless mode) before relying on any flag or file format below. Adapt to what's actually installed, and record any differences. + +Ask the human **only** what can't be found in the repo: the business intent, priority trade-offs, or approvals. Anticipate the full question set and ask it **once, in one batch** (see Autonomy), recording answers as named decisions D1…Dn in the kit. Don't trickle questions later. + +## Step 2: Generate the kit + +Create a kit directory at `.agent-loops//`, plus the agent and skill files listed in Step 3. Use the exact file names below. + +| File | Contents | +|---|---| +| `README.md` | What the loop does and the one-time setup. Setup covers the isolated worktree or branch, git-ignoring `PROGRESS.md` and `outputs/`, freezing `TASK.md` (`sha256sum` with the hash stored **outside** the repo, plus `chmod a-w`), and installing the verifier outside the repo. Then give the exact commands for each phase, the optional `/goal` text, the limits, the smoke tests, the rollout order, and the baseline comparison. | +| `TASK.md` | **Status:** DRAFT, plus a contract version. Then these numbered sections: (1) the requested behavior and public interface; (2) the slices table (ID, content, dependencies); (3) the ACs: global ones (G1 no regressions against the baseline, G2 the diff stays within allowed paths, G3 no weakened checks, G4 lint and format on changed files only, plus task-specific rules) and per-slice ones, each tagged [V], [R], or [H] and concrete enough to test; (4) allowed paths; (5) protected paths, where any change means BLOCKED; (6) approved tolerances and thresholds; (7) human decisions D1…Dn, never guessed; (8) out of scope. | +| `REFERENCE.md` | The facts the agents may rely on: APIs and signatures, data formats, domain rules, and links to sources. Every item is marked "verify against the installed code." | +| `LOOP_INSTRUCTIONS.md` | The authority order, the phases, and the per-slice cycle diagram: worker, then verifier, then reviewer, then a decision (`READY_FOR_HUMAN_REVIEW`, `NEEDS_REVISION`, or `BLOCKED`), then the human gate. Then come the exact decision rules, the limits (repairs, wall-clock time, one run lock, interrupted runs marked `UNVERIFIED`, one bounded retry for transient infrastructure errors only), the evidence layout, the update rule for `PROGRESS.md`, the optional `/goal` text, five workflow smoke tests, and the human handoff checklist. The smoke tests are: a planted known defect is caught; no source change gives a quiet no-op; a restart keeps findings and attempt counts; a missing prerequisite gives BLOCKED; a forbidden edit is refused. | +| `INIT_PROMPT.md` | Phase 0, which writes no task output. It records identity and isolation (repo, branch, base commit, clean tree, interpreter, tool versions) and runs an environment preflight (stop with named blockers). It maps the scoped context (exact files and lines), runs the baseline checks (record counts and pre-existing failures without fixing them), and gives a per-slice plan (files, tests per [V] AC, risks). It lists AC problems as a `contract_amendment.md` and makes recommendations for D1…Dn. It ends with `INIT_STATUS: READY \| BLOCKED`, then stops. | +| `WORKER_PROMPT.md` | Takes `SLICE`, `RUN`, and an optional `MODE=REPAIR` plus a findings path. The worker reads the slice ACs, `PROGRESS.md` (to avoid repeating failed approaches), the init report, and the installed sources. It writes tests from the ACs (named `test___*` or an equivalent per check), keeps changes minimal and in scope, and runs the slice's [V] checks itself, pasting the real output. The "Never" list repeats the protected-path and side-effect rules. When stuck, it writes an amendment or names a blocker and stops. In repair mode it fixes accepted findings only. It writes `worker_summary.md` in a fixed format and appends to `PROGRESS.md` without ever writing pass status. | +| `REVIEWER_PROMPT.md` | Fresh context with read-only tools. The reviewer must not read the worker's summary or earlier reviews. Inputs are the contract, the exact diff or output snapshot, `verification.json`, and the surrounding context. It works through a prioritized checklist: correctness against the ACs, the task-specific invariants from Step 1, default-path safety, contract integrity, and evidence gaps. A finding is blocking only with a location, the violated AC or invariant, a concrete failure scenario, and a verification step. Anything else is a `hypothesis`. Output is JSON: `slice`, `run_id`, `output_hash`, `ac_assessment[]`, `findings[]`, `unverified_requirements[]`, `scope_violations[]`, and an advisory `recommendation`. | +| `VERIFY.md` | The baseline commands, the global checks, and a table mapping each slice to its required [V] checks, plus prior slices for regression. It defines the `verification.json` schema (run, slice, TASK hash ok, base, output hash, environment, a per-check exit code and counts and status, `human_pending`, overall status, blocked reason). It includes a **reference verifier script** that the human saves outside the repo. The script checks the TASK hash; snapshots and hashes the output; checks allowed and protected paths, including untracked files; guards against diffs that add skips or remove assertions or tolerances; runs lint on changed files only and the baseline suite, tolerating only the pre-existing failures the human listed; and runs the slice checks, treating missing or skipped as BLOCKED. It writes the JSON and exits non-zero unless everything passes. It also lists what the reference script doesn't yet do. | +| `PROGRESS.md` | A template with a current-state table (slice, status, attempts n/3, last run, output hash, next action), human decisions, open findings, rejected findings, failed approaches, API differences, pre-existing baseline failures, and a compact newest-first log. A header says it is memory, not evidence. | +| `BASELINE_single_prompt.md` | A plain single-prompt version of the same task, used as the comparison arm. | + +## Step 3: Generate Claude Code subagents and skills (thin wrappers that point to the kit) + +Check the installed format first. As of current docs: +- **Subagents** are Markdown files in `.claude/agents/` with YAML frontmatter (`name`, `description`, `tools`, optional `model`). +- **Skills** are `.claude/skills//SKILL.md` with frontmatter (`name`, `description`). + +Create: +- `.claude/agents/-worker.md`: the tools needed for edits plus scoped `Bash` (tests, lint, `git diff` and `git status` only). The body is "Follow `.agent-loops//WORKER_PROMPT.md` exactly," plus the "Never" list. +- `.claude/agents/-reviewer.md`: **read-only tools only** (`Read`, `Grep`, `Glob`, and read-only `git diff`). The body is "Follow `REVIEWER_PROMPT.md`; you have no memory of the implementation." It must be run as a separate invocation, or as a subagent with its own context, never in the worker's conversation. +- `.claude/agents/-initializer.md`: read-only, plus running the baseline checks. It follows `INIT_PROMPT.md`. +- `.claude/skills/-loop/SKILL.md`: a short guide for a **supervising** session that explains how to step through a slice. The steps are: invoke the worker, tell the human to run the verifier, invoke the reviewer, present the findings for the human to accept or reject, and record `decision.md`. It must state that the supervising session cannot mark a slice as ready on its own, and that it never runs the verifier in place of the human's trusted copy unless the human has configured that. + +Do **not** add hooks that block or auto-approve. Optionally, propose (without installing) a `PostToolUse` hook that runs cheap, scoped lint on edited files, and note that hooks give feedback only and are not enforcement. + +## Step 4: Self-check the kit, then stop + +Run these checks over the generated kit and write the results to `.agent-loops//KIT_CHECK.md`: +- Every AC has exactly one evidence tag. Every [V] AC appears in `VERIFY.md` with a concrete command or test pattern. Every slice has at least one [V] AC, or a stated reason why it can't. +- The allowed and protected paths don't overlap. The kit's own control files, the verifier, and the baselines are protected. +- The worker and reviewer prompts both repeat the protected-path rules and the side-effect rules. The reviewer prompt forbids reading the worker's summary. +- The reference verifier script parses (`python -m py_compile` on the extracted block) and contains no network calls. +- Every command in the README exists in the installed environment, or is marked "verify." +- Decisions D1…Dn are listed in both `TASK.md` and `PROGRESS.md`. + +Then reply with (stopping only if the Autonomy batch is still unanswered; otherwise continue straight into Phase 0 and the slices when the user has asked for a full run): +1. the files created; +2. the slice list, with the AC count per slice and the split between [V], [R], and [H]; +3. the human decisions needed before Phase 0; +4. any Claude Code version differences found; +5. the exact next commands the human should run (freeze `TASK.md`, install the verifier, run Phase 0). + +Do not run Phase 0 or start any slice until `TASK.md` is approved. Once the user has approved it (or granted a standing run-to-completion authorization), proceed autonomously per the Autonomy section. + +## Rules for this generation run +- Write only inside `.agent-loops//`, `.claude/agents/`, and `.claude/skills/`. Don't touch source, tests, CI, or the existing guidance files. +- Keep each generated file short and specific to the task. If something is generic, reference this design rather than padding the file. +- If the task is too vague to write testable ACs, still produce the kit. Mark the vague ACs `[NEEDS DECISION]`, list the questions as decisions, and say so in your reply. diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2e338314..cb715b00 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ repos: # Ruff - fast linter and formatter (replaces black, isort, flake8) - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.20 + rev: v0.16.6 hooks: # Run the linter with auto-fixes - id: ruff-check diff --git a/docs/assets/tutorial_13_combined_motion_ct.gif b/docs/assets/tutorial_13_combined_motion_ct.gif index be4acb43..fc883a23 100644 --- a/docs/assets/tutorial_13_combined_motion_ct.gif +++ b/docs/assets/tutorial_13_combined_motion_ct.gif @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:b33945153c80973107a6a7d0d219b4d494a8bf8d652b883fac2d3d789cef59ff -size 1538902 +oid sha256:57bd5e811517c913c400bb9892a14be0ad90e5720e95fc3060ee71d1dd9f6461 +size 2041346 diff --git a/docs/assets/tutorial_13_combined_motion_usd.gif b/docs/assets/tutorial_13_combined_motion_usd.gif index df051b07..ce49b7aa 100644 --- a/docs/assets/tutorial_13_combined_motion_usd.gif +++ b/docs/assets/tutorial_13_combined_motion_usd.gif @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:75b948a2230c5a601feba8b83d7b70b2ad20f717842d102eb50e339383338dd3 -size 7252855 +oid sha256:cad3c69eaf4768973113fc0ecf59429a10a1737b92160c952f1d1e9c666f6c0a +size 3084897 diff --git a/docs/developer/migration_next.md b/docs/developer/migration_next.md index 2e329377..5a0b5fdf 100644 --- a/docs/developer/migration_next.md +++ b/docs/developer/migration_next.md @@ -37,3 +37,35 @@ new_call(argument, required_option="value") **Automated conversion:** ``, or `None needed`. ```` + +## InferPhysicsNeMoBase, InferPhysicsNeMoMGN, InferPhysicsNeMoMLP - `model_tag` is now a constructor argument + +**Change:** `model_tag` moved from a class attribute to the first `__init__` +argument of `InferPhysicsNeMoBase` and its MGN and MLP subclasses +(defaults `"base"`, `"mgn"`, `"mlp"`). `log_level` is now the second +positional argument. A subclass that only sets a class-level `model_tag` +now gets the constructor default instead, because `__init__` assigns the +instance attribute. + +**Why:** the checkpoint tag must match the training method that wrote the +weights (for example `"physics_informed_motion"`). It can now be chosen per +instance without defining a subclass. + +**Before** + +```python +class InferPhysicsInformed(InferPhysicsNeMoMGN): + model_tag = "physics_informed_motion" + + +method = InferPhysicsInformed("DEBUG") +``` + +**After** + +```python +method = InferPhysicsNeMoMGN(model_tag="physics_informed_motion", log_level="DEBUG") +``` + +**Automated conversion:** `None needed`. Pass `log_level` by keyword and +replace class-level `model_tag` overrides with the constructor argument. diff --git a/experiments/Registration-Fireants_vs_Greedy/affine_diagnostic.py b/experiments/Registration-Fireants_vs_Greedy/affine_diagnostic.py new file mode 100644 index 00000000..dfc69588 --- /dev/null +++ b/experiments/Registration-Fireants_vs_Greedy/affine_diagnostic.py @@ -0,0 +1,96 @@ +"""Affine-only diagnostic: FireANTs vs Greedy on selected heart phases. + +Registers each listed phase to the heart reference frame with an Affine-only +registrar (FireANTs class defaults vs the tutorial Greedy registrar) and +records NCC/MSE of the warped image plus the affine matrix summary. + +Usage: + py affine_diagnostic.py --out experiments/Registration-Fireants_vs_Greedy/results/affine_diagnostic.json +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +import time +from pathlib import Path +from typing import Any + +import itk +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import benchmark + + +def affine_summary(transform: itk.Transform) -> dict[str, Any]: + """Matrix determinant, singular values and translation of an affine.""" + matrix = itk.array_from_matrix(transform.GetMatrix()) + singular = np.linalg.svd(matrix, compute_uv=False) + return { + "determinant": float(np.linalg.det(matrix)), + "singular_values": [float(v) for v in singular], + "translation_mm": [float(v) for v in transform.GetTranslation()], + } + + +def main() -> None: + """Run the diagnostic and write the JSON.""" + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--phases", default="3,13,15") + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + logging.basicConfig(level=logging.WARNING) + + files, reference_index = benchmark.load_dataset("heart", False) + images = [itk.imread(str(path)) for path in files] + reference = images[reference_index] + fixed_array = itk.array_from_image(reference) + identity = itk.IdentityTransform[itk.D, 3].New() + + result: dict[str, Any] = {"dataset": "heart", "reference_frame": reference_index} + rows: dict[str, list[dict[str, Any]]] = {"greedy": [], "fireants": []} + for method in ("greedy", "fireants"): + registrar = benchmark.build_registrar("heart", method, False, {}) + registrar.set_transform_type("Affine") + registrar.set_fixed_image(reference) + registrar.register(images[0]) + for phase in (int(p) for p in args.phases.split(",")): + moving = images[phase] + if method == "fireants": + torch.cuda.synchronize() + start = time.perf_counter() + out = registrar.register(moving) + if method == "fireants": + torch.cuda.synchronize() + seconds = time.perf_counter() - start + forward = out["fixed_to_moving_transform"] + warped = itk.array_from_image( + benchmark.resample(moving, forward, reference) + ) + unregistered = itk.array_from_image( + benchmark.resample(moving, identity, reference) + ) + rows[method].append( + { + "phase": phase, + "seconds": seconds, + "ncc": benchmark.ncc(warped, fixed_array), + "mse": benchmark.mse(warped, fixed_array), + "ncc_unregistered": benchmark.ncc(unregistered, fixed_array), + "mse_unregistered": benchmark.mse(unregistered, fixed_array), + "affine": affine_summary(forward), + } + ) + result["greedy"] = rows["greedy"] + result["fireants"] = rows["fireants"] + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(result, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/experiments/Registration-Fireants_vs_Greedy/benchmark.py b/experiments/Registration-Fireants_vs_Greedy/benchmark.py new file mode 100644 index 00000000..72d7e4b6 --- /dev/null +++ b/experiments/Registration-Fireants_vs_Greedy/benchmark.py @@ -0,0 +1,582 @@ +"""Benchmark RegisterImagesFireants against RegisterImagesGreedy. + +Rebuilds the tutorial_03 (heart and lung) workflow: same data, same reference +frame, ``WorkflowReconstructHighres4DCT`` with an injected registrar. Every +non-reference frame is registered to the reference frame; speed and proxy +accuracy metrics are written to one result JSON (schema: VERIFY.md). + +Usage: + py benchmark.py --dataset heart --method greedy --out + py benchmark.py --dataset lung --method fireants --out [--smoke] + [--config ] + +``--smoke`` uses the heart test subset (``data/test/slicer_heart_small``) and +minimal iterations; the lung smoke run uses the full lung files with minimal +iterations. ``--config`` overrides ``number_of_iterations``, +``transform_type``, ``metric``, ``scales`` and the FireANTs deformable options +(``RegisterImagesFireants.set_deformable_options``). ``--phases`` registers only +the listed phases directly with the registrar (tuning rounds; no workflow). + +Timing: one warm-up pair is registered and discarded, then the workflow runs +and every ``register()`` call is timed (CUDA synchronized before and after for +FireANTs). Metrics are computed outside the timed region, directly after each +pair, because the transforms of all pairs would not fit in memory together; +the workflow receives identity placeholders in their place. +""" + +from __future__ import annotations + +import argparse +import gc +import importlib.util +import inspect +import json +import logging +import platform +import statistics +import sys +import time +from collections.abc import Callable +from pathlib import Path +from typing import Any, Optional + +import itk +import numpy as np +import torch +from scipy import ndimage + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT / "tutorials")) + +from parameters_heart_ct_kcl import HEART_CT_KCL +from parameters_tcia_4d_lung import TCIA_4D_LUNG + +from monai_physio import ( + ProcessTransforms, + RegisterImagesBase, + RegisterImagesFireants, + WorkflowReconstructHighres4DCT, +) + +LOGGER = logging.getLogger("benchmark") + +LUNG_THRESHOLD_HU = -320.0 +BACKGROUND_HU = -1000.0 +SMOKE_ITERATIONS = [1, 0] +CONFIG_KEYS = ( + "number_of_iterations", + "transform_type", + "metric", + "scales", + "init_affine", + *sorted(RegisterImagesFireants._DEFORMABLE_OPTIONS), +) +INVERSE_POINTS_PER_AXIS = 28 +INTERIOR_MARGIN = 0.1 +DEFORMABLE_KEYS = RegisterImagesFireants._DEFORMABLE_OPTIONS + + +def load_dataset(dataset: str, smoke: bool) -> tuple[list[Path], int]: + """Return the phase files and the tutorial's reference frame index.""" + if dataset == "heart": + data_dir = HEART_CT_KCL.data_directory(smoke) / ( + "slicer_heart_small" if smoke else "Slicer-Heart-CT" + ) + files = sorted(data_dir.glob("slice_???.mha")) + reference = int(0.7 * len(files)) + else: + data_dir = TCIA_4D_LUNG.input_directory(False) + files = sorted(data_dir.glob("100_HM10395_g0??.nii.gz")) + reference = next(i for i, p in enumerate(files) if "g070" in p.name) + if not files: + raise FileNotFoundError(f"No {dataset} phase images found under {data_dir}") + return files, reference + + +def build_registrar( + dataset: str, method: str, smoke: bool, overrides: dict[str, Any] +) -> RegisterImagesBase: + """Greedy: the tutorial registrar; FireANTs: the class defaults.""" + if method == "greedy": + parameters = HEART_CT_KCL if dataset == "heart" else TCIA_4D_LUNG + registrar = parameters.registrar(smoke, log_level=logging.WARNING) + else: + registrar = RegisterImagesFireants(log_level=logging.WARNING) + registrar.set_metric("CC") + if smoke: + registrar.set_number_of_iterations(list(SMOKE_ITERATIONS)) + if "number_of_iterations" in overrides: + registrar.set_number_of_iterations(list(overrides["number_of_iterations"])) # type: ignore[attr-defined] + if "transform_type" in overrides: + registrar.set_transform_type(overrides["transform_type"]) # type: ignore[attr-defined] + if "metric" in overrides: + registrar.set_metric(overrides["metric"]) # type: ignore[attr-defined] + if "scales" in overrides: + registrar.set_scales(list(overrides["scales"])) # type: ignore[attr-defined] + if "init_affine" in overrides: + registrar.set_init_affine(overrides["init_affine"]) # type: ignore[attr-defined] + options = {k: v for k, v in overrides.items() if k in DEFORMABLE_KEYS} + if options: + registrar.set_deformable_options(**options) # type: ignore[attr-defined] + registrar.set_modality("ct") + return registrar + + +def registrar_config(method: str, registrar: Any) -> dict[str, Any]: + """Full resolved configuration of the registrar, using P-list names.""" + config: dict[str, Any] = { + "number_of_iterations": list(registrar.number_of_iterations), + "metric": registrar.metric, + "transform_type": registrar.transform_type, + } + if method == "greedy": + config["deformable_smoothing"] = registrar.deformable_smoothing + config["greedy_metric"] = registrar._greedy_metric() + return config + from monai_physio.register_images_fireants import _try_import_fireants + + greedy_class = _try_import_fireants("registration.greedy").GreedyRegistration + defaults = { + name: parameter.default + for name, parameter in inspect.signature( + greedy_class.__init__ + ).parameters.items() + } + defaults.update(registrar.deformable_options) + config["scales"] = registrar._scales() + config["loss_type"] = registrar._FIREANTS_LOSS[registrar.metric] + config["workflow"] = ( + "FireANTs AffineRegistration (library defaults) -> GreedyRegistration" + if registrar.init_affine == "Affine" + else "GreedyRegistration from identity" + ) + config["init_affine"] = registrar.init_affine + for name in sorted(DEFORMABLE_KEYS): + config[name] = defaults[name] + return config + + +def hardware_info() -> dict[str, Any]: + """CPU, RAM, GPU and OS description.""" + import psutil + + gpu = torch.cuda.get_device_properties(0) + return { + "cpu": platform.processor(), + "cpu_logical_cores": psutil.cpu_count(logical=True), + "ram_gb": round(psutil.virtual_memory().total / 2**30, 1), + "gpu": gpu.name, + "gpu_memory_gb": round(gpu.total_memory / 2**30, 1), + "os": platform.platform(), + } + + +def version_info() -> dict[str, Any]: + """Library versions plus which FireANTs code path is in use.""" + from importlib import metadata + + def version(name: str) -> str: + try: + return metadata.version(name) + except metadata.PackageNotFoundError: + return "not installed" + + fcntl_module = sys.modules.get("fcntl") + return { + "python": sys.version.split()[0], + "torch": torch.__version__, + "cuda": torch.version.cuda, + "fireants": version("fireants"), + "picsl_greedy": version("picsl_greedy"), + "itk": version("itk"), + "simpleitk": version("simpleitk"), + "numpy": np.__version__, + "scipy": version("scipy"), + "fused_ops": importlib.util.find_spec("fireants_fused_ops") is not None, + "fcntl_stub": fcntl_module is not None and not hasattr(fcntl_module, "flock"), + } + + +def ncc(a: np.ndarray, b: np.ndarray) -> float: + """Global Pearson correlation over all voxels.""" + a = a.astype("float64").ravel() + b = b.astype("float64").ravel() + a -= a.mean() + b -= b.mean() + return float((a * b).sum() / np.sqrt((a * a).sum() * (b * b).sum())) + + +def mse(a: np.ndarray, b: np.ndarray) -> float: + """Mean squared intensity difference over all voxels.""" + return float(np.mean((a.astype("float64") - b.astype("float64")) ** 2)) + + +def lung_mask(array: np.ndarray) -> np.ndarray: + """Threshold, drop in-plane-border components, keep the two largest. + + ``array`` axes are (z, y, x); components touching a y or x face are air + outside the body. Z faces are not used because a truncated scan can cut + through the lungs. + """ + labels, count = ndimage.label(array < LUNG_THRESHOLD_HU) + border = np.unique( + np.concatenate( + [ + labels[:, 0, :].ravel(), + labels[:, -1, :].ravel(), + labels[:, :, 0].ravel(), + labels[:, :, -1].ravel(), + ] + ) + ) + sizes = np.bincount(labels.ravel(), minlength=count + 1) + sizes[0] = 0 + sizes[border] = 0 + keep = np.argsort(sizes)[-2:] + return np.isin(labels, keep[sizes[keep] > 0]) + + +def dice(a: np.ndarray, b: np.ndarray) -> float: + """Dice overlap of two boolean masks.""" + total = int(a.sum()) + int(b.sum()) + return 2.0 * int((a & b).sum()) / total if total else 1.0 + + +def mask_image(mask: np.ndarray, reference: itk.Image) -> itk.Image: + """Wrap a boolean mask array as a uint8 image on ``reference``'s grid.""" + image = itk.image_from_array(mask.astype(np.uint8)) + image.CopyInformation(reference) + return image + + +def resample( + image: itk.Image, + transform: itk.Transform, + reference: itk.Image, + nearest: bool = False, + background: float = BACKGROUND_HU, +) -> itk.Image: + """Resample ``image`` onto ``reference`` with a point transform.""" + interpolator = ( + itk.NearestNeighborInterpolateImageFunction.New(image) + if nearest + else itk.LinearInterpolateImageFunction.New(image) + ) + dtype = itk.array_view_from_image(image).dtype + default = round(background) if np.issubdtype(dtype, np.integer) else background + return itk.resample_image_filter( + Input=image, + Transform=transform, + Interpolator=interpolator, + ReferenceImage=reference, + UseReferenceImage=True, + DefaultPixelValue=default, + ) + + +def interior_points(reference: itk.Image) -> np.ndarray: + """Physical positions of a coarse lattice inside the grid, shape (n, 3).""" + size = np.array([int(s) for s in reference.GetLargestPossibleRegion().GetSize()]) + axes = [ + np.linspace( + INTERIOR_MARGIN * (n - 1), + (1 - INTERIOR_MARGIN) * (n - 1), + INVERSE_POINTS_PER_AXIS, + ) + for n in size + ] + index = np.stack(np.meshgrid(*axes, indexing="ij"), axis=-1).reshape(-1, 3) + spacing = np.array(itk.spacing(reference)) + direction = itk.array_from_matrix(reference.GetDirection()) + origin = np.array(itk.origin(reference)) + return origin + (index * spacing) @ direction.T + + +def inverse_consistency_mm( + points: np.ndarray, forward: itk.Transform, inverse: itk.Transform +) -> float: + """Mean |inverse(forward(p)) - p| over ``points``.""" + errors = np.empty(len(points)) + for i, point in enumerate(points): + there = forward.TransformPoint([float(v) for v in point]) + back = inverse.TransformPoint([float(v) for v in there]) + errors[i] = np.linalg.norm(np.array([back[k] for k in range(3)]) - point) + return float(errors.mean()) + + +class Evaluator: + """Per-pair accuracy metrics against the fixed image.""" + + def __init__(self, fixed: itk.Image, dataset: str, images_dir: Path) -> None: + self.fixed = fixed + self.fixed_array = itk.array_from_image(fixed) + self.dataset = dataset + self.images_dir = images_dir + self.points = interior_points(fixed) + self.tools = ProcessTransforms() + self.fixed_mask = lung_mask(self.fixed_array) if dataset == "lung" else None + images_dir.mkdir(parents=True, exist_ok=True) + self.fixed_path = images_dir / "fixed.mha" + itk.imwrite(fixed, str(self.fixed_path), compression=True) + + def evaluate( + self, + phase: int, + moving: itk.Image, + forward: itk.Transform, + inverse: itk.Transform, + ) -> dict[str, Any]: + """Warp ``moving`` onto the fixed grid and measure the result.""" + field = self.tools.convert_transform_to_displacement_field( + forward, self.fixed, np_component_type=np.float64 + ) + field_transform = itk.DisplacementFieldTransform[itk.D, 3].New() + field_transform.SetDisplacementField(field) + warped = resample(moving, field_transform, self.fixed) + warped_path = self.images_dir / f"phase_{phase:03d}_warped.mha" + itk.imwrite(warped, str(warped_path), compression=True) + warped_array = itk.array_from_image(warped) + + identity = itk.IdentityTransform[itk.D, 3].New() + unregistered = itk.array_from_image(resample(moving, identity, self.fixed)) + jacobian = itk.array_from_image( + self.tools.compute_jacobian_determinant_from_field(field) + ) + row: dict[str, Any] = { + "ncc": ncc(warped_array, self.fixed_array), + "mse": mse(warped_array, self.fixed_array), + "neg_jacobian_fraction": float(np.mean(jacobian <= 0.0)), + "inverse_consistency_mm": inverse_consistency_mm( + self.points, forward, inverse + ), + "ncc_unregistered": ncc(unregistered, self.fixed_array), + "mse_unregistered": mse(unregistered, self.fixed_array), + } + if self.fixed_mask is not None: + moving_mask = mask_image(lung_mask(itk.array_from_image(moving)), moving) + warped_mask = itk.array_from_image( + resample(moving_mask, field_transform, self.fixed, True, 0.0) + ).astype(bool) + unregistered_mask = itk.array_from_image( + resample(moving_mask, identity, self.fixed, True, 0.0) + ).astype(bool) + row["dice"] = dice(self.fixed_mask, warped_mask) + row["dice_unregistered"] = dice(self.fixed_mask, unregistered_mask) + row["warped"] = relative_path(warped_path) + row["fixed"] = relative_path(self.fixed_path) + return row + + +def relative_path(path: Path) -> str: + """Path relative to the repository root when inside it, else absolute.""" + resolved = path.resolve() + try: + return resolved.relative_to(REPO_ROOT).as_posix() + except ValueError: + return resolved.as_posix() + + +class TimedRegistrar: + """Time ``register`` calls and hand pairs to an evaluator. + + The wrapped registrar's ``register`` is replaced on the instance, so the + workflow reaches it through the usual ``RegisterTimeSeriesImages`` path. + """ + + def __init__( + self, + registrar: RegisterImagesBase, + method: str, + on_pair: Optional[Callable[[itk.Image, dict], dict[str, Any]]] = None, + ) -> None: + self.registrar = registrar + self.cuda = method == "fireants" + self.on_pair = on_pair + self.rows: list[dict[str, Any]] = [] + self.affine_seconds = 0.0 + self.original_register = registrar.register + registrar.register = self.register # type: ignore[method-assign] + if hasattr(registrar, "_registration_method_affine_or_rigid"): + self.original_affine = registrar._registration_method_affine_or_rigid + registrar._registration_method_affine_or_rigid = ( # type: ignore[attr-defined] + self.affine + ) + + def sync(self) -> None: + """Wait for queued GPU work (FireANTs only).""" + if self.cuda: + torch.cuda.synchronize() + + def affine(self, *args: Any, **kwargs: Any) -> Any: + """Timed FireANTs affine stage.""" + self.sync() + start = time.perf_counter() + result = self.original_affine(*args, **kwargs) + self.sync() + self.affine_seconds += time.perf_counter() - start + return result + + def register(self, moving_image: itk.Image, **kwargs: Any) -> dict[str, Any]: + """Timed ``register``; metrics are computed after the clock stops.""" + self.affine_seconds = 0.0 + self.sync() + start = time.perf_counter() + result = self.original_register(moving_image, **kwargs) + self.sync() + seconds = time.perf_counter() - start + affine_seconds = self.affine_seconds + if self.on_pair is None: + self.last_seconds = seconds + return result + row: dict[str, Any] = {"seconds": seconds} + if self.cuda: + row["affine_seconds"] = affine_seconds + row["deformable_seconds"] = seconds - affine_seconds + row["loss"] = float(result["loss"]) + row.update(self.on_pair(moving_image, result)) + self.rows.append(row) + identity = itk.IdentityTransform[itk.D, 3].New() + del result["fixed_to_moving_transform"], result["moving_to_fixed_transform"] + gc.collect() + torch.cuda.empty_cache() + return { + "fixed_to_moving_transform": identity, + "moving_to_fixed_transform": identity, + "loss": row["loss"], + } + + +def run(args: argparse.Namespace) -> dict[str, Any]: + """Execute one benchmark and return the result dictionary.""" + overrides: dict[str, Any] = {} + if args.config: + overrides = json.loads(Path(args.config).read_text(encoding="utf-8")) + unknown = set(overrides) - set(CONFIG_KEYS) + if unknown: + raise ValueError(f"Unsupported --config keys: {sorted(unknown)}") + + phases = ( + [int(p) for p in args.phases.split(",")] if args.phases is not None else None + ) + files, reference_index = load_dataset(args.dataset, args.smoke) + images = [itk.imread(str(path)) for path in files] + reference = images[reference_index] + others = [i for i in range(len(images)) if i != reference_index] + LOGGER.info( + "%s: %d phases, reference %s", + args.dataset, + len(images), + files[reference_index].name, + ) + + registrar = build_registrar(args.dataset, args.method, args.smoke, overrides) + config = registrar_config(args.method, registrar) + LOGGER.info("config: %s", config) + + images_dir = args.images_dir or args.out.parent / "images" / args.out.stem + evaluator = Evaluator(reference, args.dataset, images_dir) + + registrar.set_fixed_image(reference) + LOGGER.info("warm-up pair (discarded)") + registrar.register(images[others[0]]) + + def on_pair(moving: itk.Image, result: dict) -> dict[str, Any]: + phase = next(i for i, image in enumerate(images) if image is moving) + LOGGER.info("evaluating phase %d", phase) + row = evaluator.evaluate( + phase, + moving, + result["fixed_to_moving_transform"], + result["moving_to_fixed_transform"], + ) + row["phase"] = phase + row["phase_name"] = files[phase].name + return row + + timed = TimedRegistrar(registrar, args.method, on_pair) + if phases is not None: + for phase in phases: + registrar.register(images[phase]) + else: + workflow = WorkflowReconstructHighres4DCT( + time_series_images=images, + reference_image=reference, + reference_time_frame=reference_index, + register_reference_time_frame_to_reference_image=False, + registration_method=registrar, + log_level=logging.WARNING, + ) + workflow.set_modality("ct") + workflow.register_time_series() + + repeat_phase = others[len(others) // 2] + timed.on_pair = None + repeat_seconds: list[float] = [] + for repeat in range(args.repeats): + LOGGER.info("timed repeat %d of phase %d", repeat + 1, repeat_phase) + registrar.register(images[repeat_phase]) + repeat_seconds.append(timed.last_seconds) + + per_phase = timed.rows + ordered = ["phase", "phase_name", "seconds"] + per_phase = [ + { + **{k: row[k] for k in ordered}, + **{k: v for k, v in row.items() if k not in ordered}, + } + for row in per_phase + ] + return { + "dataset": args.dataset, + "method": args.method, + "smoke": args.smoke, + "config": config, + "hardware": hardware_info(), + "versions": version_info(), + "n_phases": len(images), + "reference_frame": reference_index, + "reference_frame_name": files[reference_index].name, + "warmup_excluded": True, + "cuda_synchronized": args.method == "fireants", + "timing_notes": ( + "seconds = wall clock of registrar.register() (preprocess, " + "registration, inverse transform); metrics are outside the clock. " + "affine_seconds/deformable_seconds (FireANTs only): deformable is " + "register() minus the affine stage. repeat_seconds re-times the " + f"pair repeat_phase={repeat_phase} after the workflow run." + ), + "repeat_phase": repeat_phase, + "total_seconds": float(sum(row["seconds"] for row in per_phase)), + "per_phase": per_phase, + "repeat_seconds": repeat_seconds, + "repeat_median_seconds": ( + statistics.median(repeat_seconds) if repeat_seconds else None + ), + } + + +def main() -> None: + """Parse arguments, run the benchmark and write the JSON.""" + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--dataset", choices=["heart", "lung"], required=True) + parser.add_argument("--method", choices=["greedy", "fireants"], required=True) + parser.add_argument( + "--config", help="JSON overriding registrar settings (CONFIG_KEYS)" + ) + parser.add_argument("--smoke", action="store_true") + parser.add_argument("--phases", help="comma-separated phases (tuning rounds)") + parser.add_argument("--images-dir", type=Path, help="warped-image directory") + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") + args.out = args.out.resolve() + args.out.parent.mkdir(parents=True, exist_ok=True) + + result = run(args) + args.out.write_text(json.dumps(result, indent=2), encoding="utf-8") + LOGGER.info("wrote %s (total %.1f s)", args.out, result["total_seconds"]) + + +if __name__ == "__main__": + main() diff --git a/experiments/Registration-Fireants_vs_Greedy/make_report.py b/experiments/Registration-Fireants_vs_Greedy/make_report.py new file mode 100644 index 00000000..3bb2a3ff --- /dev/null +++ b/experiments/Registration-Fireants_vs_Greedy/make_report.py @@ -0,0 +1,715 @@ +"""Generate REPORT.md from the benchmark and tuning JSONs, deterministically. + +Reads ``experiments/Registration-Fireants_vs_Greedy/results`` (greedy_baseline, fireants_default and +fireants_tuned JSONs for heart and lung, ``tuning/**/notes.json`` with the +candidate run JSONs, and ``affine_diagnostic.json``) and writes a Markdown +report. The output contains no timestamps or paths; the same inputs always give +byte-identical output. + +Usage: + py make_report.py [--out experiments/Registration-Fireants_vs_Greedy/results/REPORT.md] +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path +from typing import Any, Optional + +REPO_ROOT = Path(__file__).resolve().parents[2] +RESULTS = Path(__file__).resolve().parent / "results" +DEFAULT_OUT = RESULTS / "REPORT.md" + +NCC_TOLERANCE = 0.002 +METHODS = ( + ("greedy_baseline", "Greedy"), + ("fireants_default", "FireANTs-default"), + ("fireants_tuned", "FireANTs-tuned"), +) +DATASETS = (("heart", "Heart"), ("lung", "Lung")) + +Runs = dict[str, dict[str, dict[str, Any]]] + + +def load(path: Path) -> Any: + """Read a JSON file.""" + return json.loads(path.read_text(encoding="utf-8")) + + +def mean(rows: list[dict[str, Any]], key: str) -> float: + """Mean of ``key`` over the rows that have it.""" + return statistics.fmean(r[key] for r in rows if key in r) + + +def rows_of(run: dict[str, Any], phases: Optional[set[int]] = None) -> list[Any]: + """Per-phase rows of a run, optionally restricted to ``phases``.""" + return [r for r in run["per_phase"] if phases is None or r["phase"] in phases] + + +def parse_phases(text: str) -> set[int]: + """Parse '3,9,17' into a set.""" + return {int(p) for p in text.split(",")} + + +def fmt_set(phases: set[int]) -> str: + """Sorted comma-separated phase list.""" + return ",".join(str(p) for p in sorted(phases)) + + +def fmt_list(values: list[Any]) -> str: + """Compact bracketed list.""" + return "[" + ",".join(str(v) for v in values) + "]" + + +def table(header: list[str], body: list[list[str]]) -> list[str]: + """Markdown table lines.""" + lines = ["| " + " | ".join(header) + " |"] + lines.append("|" + "|".join("---" for _ in header) + "|") + lines.extend("| " + " | ".join(row) + " |" for row in body) + return lines + + +def metrics(run: dict[str, Any], phases: Optional[set[int]] = None) -> dict[str, Any]: + """Mean metrics of a run over the given phases (default all).""" + rows = rows_of(run, phases) + out: dict[str, Any] = { + "n": len(rows), + "total": sum(r["seconds"] for r in rows), + "per_pair": mean(rows, "seconds"), + "ncc": mean(rows, "ncc"), + "mse": mean(rows, "mse"), + "negj": 100.0 * mean(rows, "neg_jacobian_fraction"), + "inv": mean(rows, "inverse_consistency_mm"), + "ncc_unr": mean(rows, "ncc_unregistered"), + "mse_unr": mean(rows, "mse_unregistered"), + } + if "dice" in rows[0]: + out["dice"] = mean(rows, "dice") + out["dice_unr"] = mean(rows, "dice_unregistered") + return out + + +def tuning_dir(name: str) -> Path: + """Tuning directory of a dataset.""" + return RESULTS / "tuning" / name + + +def tuning_rounds(name: str) -> list[dict[str, Any]]: + """Notes of every round of a dataset, in numeric round order.""" + paths = sorted( + tuning_dir(name).glob("round_*/notes.json"), + key=lambda p: int(p.parent.name.split("_")[1]), + ) + return [{**load(p), "dir": p.parent} for p in paths] + + +def tuning_phases_used(name: str) -> set[int]: + """Union of the tuning phases of all rounds of a dataset.""" + used: set[int] = set() + for r in tuning_rounds(name): + used |= parse_phases(r["tuning_phases"]) + return used + + +def dataset_section( + name: str, label: str, runs: dict[str, dict[str, Any]] +) -> tuple[list[str], dict[str, dict[str, Any]]]: + """Results table for one dataset and its mean metrics.""" + stats = {key: metrics(run) for key, run in runs.items()} + greedy = stats["greedy_baseline"] + lung = name == "lung" + header = [ + "Method", + "total s", + "per-pair s (mean)", + "speedup vs Greedy", + "NCC", + "MSE", + "neg-Jacobian %", + "inverse-consistency mm", + ] + if lung: + header.append("Dice") + body: list[list[str]] = [] + for key, method_label in METHODS: + s = stats[key] + row = [ + method_label, + f"{s['total']:.1f}", + f"{s['per_pair']:.2f}", + f"{greedy['total'] / s['total']:.2f}x", + f"{s['ncc']:.5f}", + f"{s['mse']:.1f}", + f"{s['negj']:.5f}", + f"{s['inv']:.5f}", + ] + if lung: + row.append(f"{s['dice']:.5f}") + body.append(row) + unr = ["Unregistered", "-", "-", "-", f"{greedy['ncc_unr']:.5f}"] + unr += [f"{greedy['mse_unr']:.1f}", "-", "-"] + if lung: + unr.append(f"{greedy['dice_unr']:.5f}") + body.append(unr) + + run = runs["greedy_baseline"] + lines = [f"## {label}", ""] + lines.append( + f"{greedy['n']} pairs (every phase except the reference frame " + f"{run['reference_frame']}, {run['reference_frame_name']}); " + "metrics are means over pairs." + ) + lines.append("") + lines.extend(table(header, body)) + lines.append("") + lines.append( + "Times are the wall clock of `register()` per pair (preprocessing, " + "registration, inverse transform; one warm-up pair excluded; metrics " + "outside the clock). Repeat timings of phase " + f"{run['repeat_phase']} (median of {len(run['repeat_seconds'])} " + "repeats): " + + ", ".join( + f"{lab} {runs[key]['repeat_median_seconds']:.2f} s" for key, lab in METHODS + ) + + "." + ) + lines.append("") + return lines, stats + + +def settings_section(runs: Runs) -> list[str]: + """Greedy settings, FireANTs default and selected tuned configurations.""" + lines = ["## Configurations", ""] + lines.append("### Greedy settings used (not tuned)") + lines.append("") + for name, label in DATASETS: + c = runs[name]["greedy_baseline"]["config"] + lines.append( + f"- {label}: iterations {fmt_list(c['number_of_iterations'])}, " + f"metric {c['metric']} ({c['greedy_metric']}), transform " + f"{c['transform_type']}, deformable smoothing " + f"{c['deformable_smoothing']} (the tutorial_03 registrar)." + ) + lines.append("") + keys = [ + "number_of_iterations", + "scales", + "init_affine", + "workflow", + "optimizer_lr", + "smooth_warp_sigma", + "smooth_grad_sigma", + "cc_kernel_size", + "cc_kernel_type", + "deformation_type", + "integrator_n", + "blur", + ] + lines.append("### FireANTs-default and FireANTs-tuned (selected) configurations") + lines.append("") + header = ["Setting"] + [ + f"{label} {m}" for name, label in DATASETS for m in ("default", "tuned") + ] + body = [] + for key in keys: + row = [key] + for name, _ in DATASETS: + for method in ("fireants_default", "fireants_tuned"): + value = runs[name][method]["config"][key] + row.append(fmt_list(value) if isinstance(value, list) else str(value)) + body.append(row) + lines.extend(table(header, body)) + lines.append("") + lines.append( + "Common to all FireANTs runs: metric CC (loss cc), transform " + "Deformable, Adam optimizer, tolerance " + f"{runs['heart']['fireants_tuned']['config']['tolerance']}." + ) + lines.append("") + return lines + + +def tuning_section(runs: Runs) -> list[str]: + """What was evaluated in the tuning rounds, and the lung margin decision.""" + lines = ["## Tuning record", ""] + lines.append( + "Objective (D8): minimize seconds per pair subject to NCC >= Greedy " + f"NCC - {NCC_TOLERANCE} and neg-Jacobian <= Greedy's, both measured on " + "the same tuning phases of that round. Greedy was never tuned. A round " + "is a batch of at most 3 candidate configurations." + ) + lines.append("") + for name, label in DATASETS: + rounds = tuning_rounds(name) + greedy = runs[name]["greedy_baseline"] + counts = [len(r["candidates"]) for r in rounds] + files = sorted(p.name for p in tuning_dir(name).glob("round_*/r*.json")) + lines.append(f"### {label}") + lines.append("") + lines.append( + f"- Round folders: {len(rounds)}; configurations evaluated: " + f"{sum(counts)} (run JSONs on disk: {len(files)}); largest batch: " + f"{max(counts)}." + ) + lines.append("") + header = [ + "Round", + "Run", + "iterations", + "scales", + "init", + "lr", + "warp sigma", + "kernel", + "tuning phases", + "s/pair", + "NCC", + "neg-Jac %", + "NCC margin", + "selected", + ] + body = [] + notes_lines: list[str] = [] + for r in rounds: + phases = parse_phases(r["tuning_phases"]) + greedy_ncc = mean(rows_of(greedy, phases), "ncc") + for cand in r["candidates"]: + c = cand["config"] + m = metrics(load(r["dir"] / cand["run_json"])) + run_id = cand["run_json"].removesuffix(".json") + body.append( + [ + str(r["round"]), + run_id, + fmt_list(c["number_of_iterations"]), + fmt_list(c["scales"]), + c["init_affine"].split()[0], + str(c["optimizer_lr"]), + str(c["smooth_warp_sigma"]), + str(c["cc_kernel_size"]), + r["tuning_phases"], + f"{m['per_pair']:.2f}", + f"{m['ncc']:.5f}", + f"{m['negj']:.5f}", + f"{m['ncc'] - (greedy_ncc - NCC_TOLERANCE):+.5f}", + "yes" if "SELECTED" in cand["verdict"] else "no", + ] + ) + notes_lines.append(f"- {run_id}: " + cand["verdict"].replace("|", "/")) + lines.extend(table(header, body)) + lines.append("") + lines.append( + "NCC margin = candidate NCC minus (Greedy NCC on the same phases " + f"- {NCC_TOLERANCE}); negative means the NCC constraint failed." + ) + lines.append("") + lines.append("Recorded rationale per configuration:") + lines.append("") + lines.extend(notes_lines) + lines.append("") + lines.extend(lung_margin_section(runs["lung"]["greedy_baseline"])) + return lines + + +def lung_margin_section(greedy: dict[str, Any]) -> list[str]: + """Margin criterion for lung r3_a versus r3_c.""" + round3 = next(r for r in tuning_rounds("lung") if r["round"] == 3) + phases = parse_phases(round3["tuning_phases"]) + greedy_ncc = mean(rows_of(greedy, phases), "ncc") + threshold = greedy_ncc - NCC_TOLERANCE + cands = { + c["run_json"].removesuffix(".json"): (c, load(round3["dir"] / c["run_json"])) + for c in round3["candidates"] + } + a, c_ = cands["r3_a"], cands["r3_c"] + ma, mc = metrics(a[1]), metrics(c_[1]) + lines = ["### Lung round 3: r3_a versus r3_c (margin criterion)", ""] + lines.append( + "Both candidates satisfy the D8 constraints on the round-3 tuning " + f"phases ({round3['tuning_phases']}); the NCC threshold is " + f"{threshold:.5f} (Greedy {greedy_ncc:.5f} - {NCC_TOLERANCE})." + ) + lines.append("") + for tag, cand, m in (("r3_a", a, ma), ("r3_c", c_, mc)): + cfg = cand[0]["config"] + lines.append( + f"- {tag} (iterations {fmt_list(cfg['number_of_iterations'])}, " + f"scales {fmt_list(cfg['scales'])}): {m['per_pair']:.2f} s/pair, " + f"NCC {m['ncc']:.5f}, margin {m['ncc'] - threshold:+.5f}." + ) + lines.append( + f"- r3_c is {ma['per_pair'] - mc['per_pair']:+.2f} s/pair faster than " + f"r3_a; its NCC differs from r3_a by {mc['ncc'] - ma['ncc']:+.5f}." + ) + lines.append(f"- Recorded decision for r3_a: {a[0]['verdict']}") + lines.append(f"- Recorded decision for r3_c: {c_[0]['verdict']}") + lines.append("") + lines.append( + "A strict reading of D8 (minimize seconds per pair among candidates " + "that pass the constraints) would have selected r3_c, the faster " + "candidate. r3_a was selected instead by a judgement call recorded in " + "the round notes, not a pre-registered numeric rule: the larger NCC " + "margin was preferred over a speed gain of a fraction of a second per " + "pair, which is within the pair-to-pair timing variation of repeated " + "runs (about a second). No numeric margin threshold was defined." + ) + lines.append("") + return lines + + +def heldout_section(runs: Runs) -> list[str]: + """Constraint check on phases never used for tuning.""" + lines = ["## Held-out-phase constraint check", ""] + lines.append( + "The tuning rounds used a few phases per dataset, and those phases are " + "also among the reported phases above, so the all-phase numbers are " + "not independent of tuning. The check below uses only phases that no " + "round used for tuning, with the D8 constraints applied to the means " + "over those phases." + ) + lines.append("") + for name, label in DATASETS: + greedy = runs[name]["greedy_baseline"] + tuned = runs[name]["fireants_tuned"] + used = tuning_phases_used(name) + held = {r["phase"] for r in greedy["per_phase"]} - used + mg, mt = metrics(greedy, held), metrics(tuned, held) + gmap = {r["phase"]: r for r in rows_of(greedy, held)} + tmap = {r["phase"]: r for r in rows_of(tuned, held)} + worst = min(tmap[p]["ncc"] - gmap[p]["ncc"] for p in held) + ncc_ok = mt["ncc"] >= mg["ncc"] - NCC_TOLERANCE + negj_ok = mt["negj"] <= mg["negj"] + lines.append(f"### {label}") + lines.append("") + lines.append( + f"- Phases used for tuning (excluded): {fmt_set(used)}; held-out " + f"phases ({len(held)}): {fmt_set(held)}." + ) + lines.append( + f"- Mean NCC: tuned {mt['ncc']:.5f} vs Greedy {mg['ncc']:.5f} " + f"(difference {mt['ncc'] - mg['ncc']:+.5f}; constraint >= " + f"{-NCC_TOLERANCE:+.3f}): {'met' if ncc_ok else 'NOT met'}." + ) + lines.append( + f"- Mean neg-Jacobian: tuned {mt['negj']:.5f}% vs Greedy " + f"{mg['negj']:.5f}% (constraint tuned <= Greedy): " + f"{'met' if negj_ok else 'NOT met'}." + ) + lines.append( + f"- Worst single-phase NCC difference (tuned - Greedy): {worst:+.5f}." + ) + lines.append("") + body = [ + [ + str(p), + f"{gmap[p]['ncc']:.5f}", + f"{tmap[p]['ncc']:.5f}", + f"{100 * gmap[p]['neg_jacobian_fraction']:.5f}", + f"{100 * tmap[p]['neg_jacobian_fraction']:.5f}", + ] + for p in sorted(held) + ] + header = [ + "Phase", + "Greedy NCC", + "tuned NCC", + "Greedy neg-Jac %", + "tuned neg-Jac %", + ] + lines.extend(table(header, body)) + lines.append("") + return lines + + +def affine_section(runs: Runs) -> list[str]: + """Identity-initialised tuning and the default affine divergence.""" + lines = ["## Affine initialisation", ""] + default_init = runs["heart"]["fireants_default"]["config"]["init_affine"] + lines.append( + f"- FireANTs-default starts the deformable stage from: {default_init}." + ) + lines.append( + "- FireANTs-tuned starts from: " + f"{runs['heart']['fireants_tuned']['config']['init_affine']} (heart) and " + f"{runs['lung']['fireants_tuned']['config']['init_affine']} (lung), i.e. " + "it is deformable-only and does not use the FireANTs affine stage; the " + "tuned speed and accuracy therefore do not include an affine alignment." + ) + greedy = {r["phase"]: r for r in runs["heart"]["greedy_baseline"]["per_phase"]} + default = {r["phase"]: r for r in runs["heart"]["fireants_default"]["per_phase"]} + tuned = {r["phase"]: r for r in runs["heart"]["fireants_tuned"]["per_phase"]} + lines.append("") + lines.append("Heart phases 13 and 15 in the benchmark runs:") + lines.append("") + body = [ + [ + str(p), + f"{default[p]['ncc_unregistered']:.5f}", + f"{greedy[p]['ncc']:.5f}", + f"{default[p]['ncc']:.5f}", + f"{tuned[p]['ncc']:.5f}", + f"{100 * default[p]['neg_jacobian_fraction']:.5f}", + ] + for p in (13, 15) + ] + header = [ + "Phase", + "unregistered NCC", + "Greedy NCC", + "default NCC", + "tuned NCC", + "default neg-Jac %", + ] + lines.extend(table(header, body)) + lines.append("") + path = RESULTS / "affine_diagnostic.json" + if not path.exists(): + lines.append("Affine diagnostic (affine_diagnostic.json) not available.") + lines.append("") + return lines + diag = load(path) + lines.append( + "Diagnostic (`affine_diagnostic.py`): Affine-only registration of " + "heart phases to the reference frame, Greedy affine (tutorial " + "registrar) versus FireANTs affine (class defaults, the stage " + "FireANTs-default uses as initialisation)." + ) + lines.append("") + body = [] + for g, f in zip(diag["greedy"], diag["fireants"]): + body.append( + [ + str(g["phase"]), + f"{g['ncc_unregistered']:.5f}", + f"{g['ncc']:.5f}", + f"{f['ncc']:.5f}", + f"{g['seconds']:.1f}", + f"{f['seconds']:.1f}", + fmt_list([round(v, 1) for v in g["affine"]["translation_mm"]]), + fmt_list([round(v, 1) for v in f["affine"]["translation_mm"]]), + f"{g['affine']['determinant']:.4f}", + f"{f['affine']['determinant']:.4f}", + ] + ) + header = [ + "Phase", + "unregistered NCC", + "Greedy affine NCC", + "FireANTs affine NCC", + "Greedy s", + "FireANTs s", + "Greedy translation mm", + "FireANTs translation mm", + "Greedy det", + "FireANTs det", + ] + lines.extend(table(header, body)) + lines.append("") + bad = [f for f in diag["fireants"] if f["ncc"] < f["ncc_unregistered"] - 0.01] + bad_phases = ", ".join(str(f["phase"]) for f in bad) or "none" + lines.append( + "Finding: FireANTs affine with library defaults ends worse than no " + f"registration at heart phases {bad_phases} (NCC below the " + "unregistered value by more than 0.01, translations of tens of mm " + "along one axis), while Greedy affine stays near identity on the " + "same pairs. The deformable stage cannot fully recover from such an " + "initialisation, which matches the low FireANTs-default NCC at those " + "phases above. Identity initialisation avoids the affine stage " + "altogether; the heart pairs are already close to aligned (the " + "unregistered NCC is high). The diagnostic points to " + "optimiser behaviour of the FireANTs affine stage at default " + "settings rather than a defect in the transform conversion: the " + "conversion is exercised by the known-affine tests of the class and " + "gives sane transforms on the other phases. Tuning the affine stage " + "was not attempted." + ) + lines.append("") + return lines + + +def discussion_section(stats: dict[str, dict[str, dict[str, Any]]]) -> list[str]: + """Speed and accuracy trade-offs, and the recommendation.""" + lines = ["## Discussion", ""] + lines.append("### Speed") + lines.append("") + for name, label in DATASETS: + g = stats[name]["greedy_baseline"] + d = stats[name]["fireants_default"] + t = stats[name]["fireants_tuned"] + lines.append( + f"- {label}: FireANTs-default is {g['total'] / d['total']:.2f}x and " + f"FireANTs-tuned {g['total'] / t['total']:.2f}x faster than Greedy " + f"({g['per_pair']:.1f} s/pair for Greedy, {d['per_pair']:.1f} s for " + f"default, {t['per_pair']:.1f} s for tuned). The tuned gain over " + f"the default is {d['total'] / t['total']:.2f}x; the remaining time " + "is largely preprocessing and inverse-transform computation inside " + "`register()`, which the iteration counts do not change." + ) + lines.append("") + lines.append("### Accuracy") + lines.append("") + for name, label in DATASETS: + g = stats[name]["greedy_baseline"] + d = stats[name]["fireants_default"] + t = stats[name]["fireants_tuned"] + text = ( + f"- {label}: mean NCC Greedy {g['ncc']:.5f}, default {d['ncc']:.5f}, " + f"tuned {t['ncc']:.5f} (unregistered {g['ncc_unr']:.5f}); mean MSE " + f"Greedy {g['mse']:.1f}, default {d['mse']:.1f}, tuned " + f"{t['mse']:.1f}. Tuned MSE is " + f"{100 * (t['mse'] / g['mse'] - 1):+.1f}% relative to Greedy, so " + "the NCC-based constraint can hide a larger intensity-difference " + f"gap. Neg-Jacobian: Greedy {g['negj']:.5f}%, default " + f"{d['negj']:.5f}%, tuned {t['negj']:.5f}%. Inverse consistency: " + f"Greedy {g['inv']:.5f} mm, default {d['inv']:.5f} mm, tuned " + f"{t['inv']:.5f} mm" + ) + if "dice" in g: + text += ( + f"; Dice Greedy {g['dice']:.5f}, default {d['dice']:.5f}, " + f"tuned {t['dice']:.5f} (unregistered {g['dice_unr']:.5f})" + ) + lines.append(text + ".") + lines.append("") + lines.append( + "In short: FireANTs-tuned meets the NCC and neg-Jacobian constraints " + "it was tuned for, but it is not uniformly as accurate as Greedy; " + "MSE (and Dice for lung) are worse than Greedy's. The NCC tolerance of " + f"{NCC_TOLERANCE} was set by the task, not derived from a downstream " + "requirement. FireANTs-default is the least accurate configuration " + "on heart and has the highest neg-Jacobian fractions on both datasets." + ) + lines.append("") + lines.append("## Recommendation") + lines.append("") + lines.append( + "This is a recommendation only; changing the default registrar is out " + "of scope for this work." + ) + lines.append("") + lines.append( + "- Keep Greedy as the default registrar. Offer RegisterImagesFireants " + "as an opt-in registrar for GPU users where a 3x (heart) to 5x (lung) " + "reduction in registration time outweighs a small loss of accuracy " + "relative to Greedy." + ) + lines.append( + "- When FireANTs is used, use the tuned, Identity-initialised " + "configurations above rather than the class defaults, which diverge " + "in the affine stage on some heart phases and give higher " + "neg-Jacobian fractions." + ) + lines.append( + "- Before any default change, repeat the comparison on more than one " + "patient per dataset and on a downstream task, since the limitations " + "below prevent generalising from these runs." + ) + lines.append("") + return lines + + +def limitations_section(runs: Runs) -> list[str]: + """Limitations, with machine facts read from the run JSONs.""" + g = runs["heart"]["greedy_baseline"] + t = runs["heart"]["fireants_tuned"] + lines = ["## Limitations", ""] + lines.append( + "- Single patient per dataset (one 4D heart CT, one 4D lung CT): no " + "between-subject variability; differences between methods are not " + "statistically tested." + ) + lines.append( + f"- Different hardware: Greedy ran on the CPU ({g['hardware']['cpu']}, " + f"{g['hardware']['cpu_logical_cores']} logical cores); FireANTs ran on " + f"the GPU ({t['hardware']['gpu']}, {t['hardware']['gpu_memory_gb']} GB). " + "The speedups compare a CPU implementation with a GPU one, not " + "algorithms on equal hardware." + ) + lines.append( + "- Windows pure-PyTorch FireANTs path: fused_ops " + f"{str(t['versions']['fused_ops']).lower()} (fireants_fused_ops not " + "installed, standard downsampling used) and fcntl_stub " + f"{str(t['versions']['fcntl_stub']).lower()} (an fcntl stub is used on " + f"Windows; Greedy run: {str(g['versions']['fcntl_stub']).lower()}). " + "The fused kernels may change FireANTs timing and numerics on other " + "platforms." + ) + for name, label in DATASETS: + used = tuning_phases_used(name) + n = runs[name]["greedy_baseline"]["n_phases"] - 1 + lines.append( + f"- {label}: the tuning phases ({fmt_set(used)}) are " + f"{len(used)} of the {n} reported phases; all-phase results are " + "not independent of tuning (see the held-out check)." + ) + lines.append( + "- The inverse-consistency metric (mean |inverse(forward(p)) - p| on " + "a coarse interior lattice) is low-information: every method gives " + "values of order 0.01 mm, and it measures the invertibility of the " + "composed transform, not registration accuracy." + ) + lines.append( + "- Dice (lung) uses a threshold-based lung mask (-320 HU) that drops " + "connected components touching the in-plane (x or y) image border " + "and keeps the two largest remaining; z faces are not used, so lung " + "cut by the scan range is kept. It is a proxy, not an expert " + "segmentation." + ) + lines.append( + "- Timing scope: only `register()` per pair is timed; image loading, " + "metric computation, warping and the 4D workflow around it are not. " + "The timings include FireANTs image conversion and the displacement " + "field inversion." + ) + lines.append( + "- NCC, MSE, neg-Jacobian and inverse consistency are proxy metrics " + "against the fixed image; no landmark or target-registration-error " + "data were used." + ) + lines.append("") + return lines + + +def build_report() -> str: + """Assemble the full report text.""" + runs: Runs = { + name: {key: load(RESULTS / f"{name}_{key}.json") for key, _ in METHODS} + for name, _ in DATASETS + } + lines = ["# FireANTs vs Greedy registration report", ""] + lines.append( + "Generated by `experiments/Registration-Fireants_vs_Greedy/" + "make_report.py` from the JSON files in `experiments/Registration-Fireants_vs_Greedy/results`. " + "Methods: Greedy (the tutorial registrar), FireANTs-default (class " + "defaults) and FireANTs-tuned (selected by the tuning rounds below)." + ) + lines.append("") + stats: dict[str, dict[str, dict[str, Any]]] = {} + lines.append("# Results") + lines.append("") + for name, label in DATASETS: + part, stats[name] = dataset_section(name, label, runs[name]) + lines.extend(part) + lines.extend(settings_section(runs)) + lines.extend(tuning_section(runs)) + lines.extend(heldout_section(runs)) + lines.extend(affine_section(runs)) + lines.extend(discussion_section(stats)) + lines.extend(limitations_section(runs)) + return "\n".join(lines).rstrip("\n") + "\n" + + +def main() -> None: + """Write the report.""" + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--out", type=Path, default=DEFAULT_OUT) + args = parser.parse_args() + args.out.parent.mkdir(parents=True, exist_ok=True) + with open(args.out, "w", encoding="utf-8", newline="\n") as handle: + handle.write(build_report()) + + +if __name__ == "__main__": + main() diff --git a/experiments/Registration-Fireants_vs_Greedy/tuning_compare.py b/experiments/Registration-Fireants_vs_Greedy/tuning_compare.py new file mode 100644 index 00000000..3b9ec241 --- /dev/null +++ b/experiments/Registration-Fireants_vs_Greedy/tuning_compare.py @@ -0,0 +1,54 @@ +"""Compare tuning-run JSONs with the Greedy baseline on the same phases. + +Usage: + py tuning_compare.py [ ...] + +Prints seconds per pair, mean NCC and mean negative-Jacobian fraction of each +run next to the Greedy baseline restricted to the run's phases, and whether the +D8 constraints (NCC >= Greedy NCC - 0.002, neg-Jacobian <= Greedy's) hold. +""" + +from __future__ import annotations + +import json +import statistics +import sys +from pathlib import Path + +RESULTS = Path(__file__).resolve().parent / "results" +NCC_MARGIN = 0.002 + + +def means(rows: list[dict]) -> tuple[float, float, float]: + """Mean seconds, NCC and negative-Jacobian fraction of per-phase rows.""" + return ( + statistics.mean(r["seconds"] for r in rows), + statistics.mean(r["ncc"] for r in rows), + statistics.mean(r["neg_jacobian_fraction"] for r in rows), + ) + + +def main() -> None: + """Print one comparison line per run JSON.""" + dataset = sys.argv[1] + greedy = json.loads( + (RESULTS / f"{dataset}_greedy_baseline.json").read_text(encoding="utf-8") + ) + for path in sys.argv[2:]: + run = json.loads(Path(path).read_text(encoding="utf-8")) + phases = [r["phase"] for r in run["per_phase"]] + reference = [r for r in greedy["per_phase"] if r["phase"] in phases] + seconds, ncc, neg_jac = means(run["per_phase"]) + g_seconds, g_ncc, g_neg_jac = means(reference) + ok_ncc = ncc >= g_ncc - NCC_MARGIN + ok_jac = neg_jac <= g_neg_jac + print( + f"{Path(path).parent.name}/{Path(path).name}: phases {phases} " + f"s/pair {seconds:.2f} (greedy {g_seconds:.2f}) " + f"NCC {ncc:.5f} (greedy {g_ncc:.5f}, ok={ok_ncc}) " + f"negJ {neg_jac:.2e} (greedy {g_neg_jac:.2e}, ok={ok_jac})" + ) + + +if __name__ == "__main__": + main() diff --git a/src/monai_physio/__init__.py b/src/monai_physio/__init__.py index fa1235ed..3c15f91f 100644 --- a/src/monai_physio/__init__.py +++ b/src/monai_physio/__init__.py @@ -86,6 +86,7 @@ # Registration classes from .register_images_base import RegisterImagesBase from .register_images_chain import RegisterImagesChain +from .register_images_fireants import RegisterImagesFireants from .register_images_greedy import RegisterImagesGreedy from .register_images_greedy_icon import RegisterImagesGreedyICON from .register_images_icon import RegisterImagesICON @@ -133,34 +134,32 @@ __all__ = [ "AnatomyGroup", - # Anatomy taxonomy (shared between segmenters and USD renderer) "AnatomyTaxonomy", - "ProcessContours", - # Data processing utilities "ConvertImage4DTo3D", "ConvertVTKToUSD", - "DownloadData", - # Distributed execution "DistributedContext", + "DownloadData", "EvaluateMovementBase", "EvaluateMovementDukeHeart", "EvaluateMovementLung", - # Utility classes - "ProcessImages", - # Inference method classes "InferPhysicsNeMoBase", "InferPhysicsNeMoMGN", "InferPhysicsNeMoMLP", - "ProcessLabelmaps", - "ProcessLandmarks", - # Base classes "MONAIPhysioBase", "MovementGroundTruth", + "ProcessContours", + "ProcessImages", + "ProcessLabelmaps", + "ProcessLandmarks", "ProcessPhysicsNemo", + "ProcessTests", + "ProcessTransforms", + "ProcessUSD", + "ProcessUSDAnatomy", "RegisterImagesANTS", - # Registration classes "RegisterImagesBase", "RegisterImagesChain", + "RegisterImagesFireants", "RegisterImagesGreedy", "RegisterImagesGreedyICON", "RegisterImagesICON", @@ -170,7 +169,6 @@ "RegisterModelsPCA", "RegisterTimeSeriesImages", "ReportEvaluateMovement", - # Segmentation classes "SegmentAnatomyBase", "SegmentChestTotalSegmentator", "SegmentChestTotalSegmentatorWithContrast", @@ -178,17 +176,11 @@ "SegmentHeartSimplewareTrimmedBranches", "SegmentNVSegmentCT", "SegmentNVSegmentCTMRI", - "ProcessTests", - # Training method classes "TrainPhysicsNeMoBase", "TrainPhysicsNeMoMGN", "TrainPhysicsNeMoMLP", "TrainPhysicsNeMoPhysicsInformedMotion", - "ProcessTransforms", - "ProcessUSDAnatomy", - "ProcessUSD", "WorkflowConvertImageToUSD", - # Workflow classes "WorkflowConvertImageToVTK", "WorkflowConvertVTKToUSD", "WorkflowCreateMeanSurface", diff --git a/src/monai_physio/infer_physicsnemo_base.py b/src/monai_physio/infer_physicsnemo_base.py index bdfc3f8f..fb22d0e1 100644 --- a/src/monai_physio/infer_physicsnemo_base.py +++ b/src/monai_physio/infer_physicsnemo_base.py @@ -35,19 +35,22 @@ class InferPhysicsNeMoBase(MONAIPhysioBase): Not instantiated directly - use :class:`monai_physio.InferPhysicsNeMoMGN` or :class:`monai_physio.InferPhysicsNeMoMLP`. Subclasses implement - :meth:`build_model`, :meth:`load_artifacts` and :meth:`predict`, and set the - class attribute ``model_tag``. + :meth:`build_model`, :meth:`load_artifacts` and :meth:`predict`. """ - model_tag: str = "base" - - def __init__(self, log_level: int | str = logging.INFO) -> None: + def __init__( + self, model_tag: str = "base", log_level: int | str = logging.INFO + ) -> None: """Initialize the inference method. Args: + model_tag: Prefix of the checkpoint files to load + (``_stage_model.pt``); must match the tag of the + training method that wrote them. Default: ``"base"``. log_level: Logging level. Default: ``logging.INFO``. """ super().__init__(class_name=self.__class__.__name__, log_level=log_level) + self.model_tag = model_tag self._model: torch.nn.Module self._device: torch.device diff --git a/src/monai_physio/infer_physicsnemo_mgn.py b/src/monai_physio/infer_physicsnemo_mgn.py index 58fac0f9..55712e8e 100644 --- a/src/monai_physio/infer_physicsnemo_mgn.py +++ b/src/monai_physio/infer_physicsnemo_mgn.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging from pathlib import Path from typing import TYPE_CHECKING, cast @@ -21,7 +22,20 @@ class InferPhysicsNeMoMGN(InferPhysicsNeMoBase): training workflow saved next to the checkpoint. """ - model_tag = "mgn" + def __init__( + self, model_tag: str = "mgn", log_level: int | str = logging.INFO + ) -> None: + """Initialize the MeshGraphNet inference method. + + Args: + model_tag: Prefix of the checkpoint files to load + (``_stage_model.pt``), e.g. + ``"physics_informed_motion"`` for weights written by + :class:`monai_physio.TrainPhysicsNeMoPhysicsInformedMotion`. + Default: ``"mgn"``. + log_level: Logging level. Default: ``logging.INFO``. + """ + super().__init__(model_tag=model_tag, log_level=log_level) def build_model(self, meta: dict) -> torch.nn.Module: MeshGraphNet = ProcessPhysicsNemo.import_meshgraphnet() diff --git a/src/monai_physio/infer_physicsnemo_mlp.py b/src/monai_physio/infer_physicsnemo_mlp.py index e45294a6..56f9673b 100644 --- a/src/monai_physio/infer_physicsnemo_mlp.py +++ b/src/monai_physio/infer_physicsnemo_mlp.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging from pathlib import Path from typing import TYPE_CHECKING, cast @@ -16,9 +17,20 @@ class InferPhysicsNeMoMLP(InferPhysicsNeMoBase): """Predict mesh stages with a trained PhysicsNeMo FullyConnected model.""" - model_tag = "mlp" _INFER_CHUNK = 262144 + def __init__( + self, model_tag: str = "mlp", log_level: int | str = logging.INFO + ) -> None: + """Initialize the MLP inference method. + + Args: + model_tag: Prefix of the checkpoint files to load + (``_stage_model.pt``). Default: ``"mlp"``. + log_level: Logging level. Default: ``logging.INFO``. + """ + super().__init__(model_tag=model_tag, log_level=log_level) + def build_model(self, meta: dict) -> torch.nn.Module: try: from physicsnemo.models.mlp import FullyConnected diff --git a/src/monai_physio/register_images_fireants.py b/src/monai_physio/register_images_fireants.py new file mode 100644 index 00000000..af05e738 --- /dev/null +++ b/src/monai_physio/register_images_fireants.py @@ -0,0 +1,417 @@ +"""FireANTs-based image registration implementation. + +This module provides the RegisterImagesFireants class, a concrete implementation +of RegisterImagesBase that uses the GPU-accelerated FireANTs library for rigid, +affine, and deformable image registration. + +See https://fireants.readthedocs.io/ and https://pypi.org/project/fireants/. +""" + +from __future__ import annotations + +import importlib +import logging +import sys +import types +from typing import Any, ClassVar, Optional, Union + +import itk +import numpy as np +import torch +from numpy.typing import NDArray + +from .process_images import ProcessImages +from .process_transforms import ProcessTransforms +from .register_images_base import RegisterImagesBase + + +def _try_import_fireants(module: str) -> Any: + """Import a ``fireants`` submodule; raise with a helpful message if missing. + + ``fireants.io.image`` imports the Unix-only ``fcntl`` module at import time + and only uses it when writing images. When ``import fcntl`` fails (Windows), + an empty stub module is registered so the package can be imported; the + installed package is not modified. + """ + try: + import fcntl # noqa: F401 + except ImportError: + sys.modules.setdefault("fcntl", types.ModuleType("fcntl")) + try: + return importlib.import_module(f"fireants.{module}") + except ImportError as e: + raise ImportError( + "fireants is required for RegisterImagesFireants. " + "Install with: pip install fireants --no-deps" + ) from e + + +class RegisterImagesFireants(RegisterImagesBase): + """FireANTs-based image registration implementation. + + This class extends RegisterImagesBase to provide rigid, affine, and deformable + image registration on the GPU using FireANTs. FireANTs transforms are expressed + in physical (LPS) coordinates as ``y = A x + t`` from the fixed to the moving + image, which is the ITK ``fixed_to_moving_transform`` directly. + + Attributes: + number_of_iterations: List of iterations per level (e.g. [40, 20, 10]). + Level ``i`` of ``n`` runs at downsampling factor ``2 ** (n - 1 - i)``. + transform_type: 'Deformable', 'Affine', or 'Rigid' + metric: 'CC' (-> cc), 'Mattes' (-> mi), or 'MeanSquares' (-> mse) + scales: Optional explicit downsampling factor per level (coarse to + fine); when None they follow from ``number_of_iterations``. + init_affine: Source of the deformable stage's initial transform: + 'Affine' (FireANTs affine registration) or 'Identity' (none). + deformable_options: Keyword arguments forwarded to FireANTs' + ``GreedyRegistration`` (see ``set_deformable_options``). + """ + + _DEFORMABLE_OPTIONS: ClassVar[frozenset[str]] = frozenset( + { + "cc_kernel_size", + "cc_kernel_type", + "smooth_grad_sigma", + "smooth_warp_sigma", + "optimizer", + "optimizer_lr", + "deformation_type", + "integrator_n", + "blur", + "max_tolerance_iters", + "tolerance", + } + ) + + _FIREANTS_LOSS: ClassVar[dict[str, str]] = { + "CC": "cc", + "Mattes": "mi", + "MeanSquares": "mse", + } + + def __init__(self, log_level: int | str = logging.INFO) -> None: + """Initialize the FireANTs image registration class. + + Args: + log_level: Logging level (default: logging.INFO) + """ + super().__init__(log_level=log_level) + + self.number_of_iterations: list[int] = [40, 20, 10] + self.transform_type = "Deformable" + self.metric = "CC" + self.scales: Optional[list[float]] = None + self.deformable_options: dict[str, Any] = {} + self.init_affine = "Affine" + + def set_number_of_iterations(self, number_of_iterations: list[int]) -> None: + """Set the number of iterations per resolution level. + + Args: + number_of_iterations: List of iterations (e.g. [40, 20, 10]). + """ + self.number_of_iterations = number_of_iterations + + def set_transform_type(self, transform_type: str) -> None: + """Set the type of transform: Deformable, Affine, or Rigid. + + Args: + transform_type: 'Deformable', 'Affine', or 'Rigid'. + """ + if transform_type not in ("Deformable", "Affine", "Rigid"): + self.log_error("Invalid transform type: %s", transform_type) + raise ValueError(f"Invalid transform type: {transform_type}") + self.transform_type = transform_type + + def set_metric(self, metric: str) -> None: + """Set the similarity metric (CC, Mattes->mi, MeanSquares->mse). + + Args: + metric: 'CC', 'Mattes', or 'MeanSquares'. + """ + if metric not in self._FIREANTS_LOSS: + self.log_error("Invalid metric: %s", metric) + raise ValueError(f"Invalid metric: {metric}") + self.metric = metric + + def set_scales(self, scales: Optional[list[float]]) -> None: + """Set the downsampling factor per level (one per iteration entry). + + Args: + scales: Factors in descending order (e.g. [4, 2, 1]), or None to + derive them from the number of iterations. + """ + if scales is not None and len(scales) != len(self.number_of_iterations): + raise ValueError("scales and number_of_iterations must have equal length") + self.scales = scales + + def set_init_affine(self, init_affine: str) -> None: + """Set the initial transform of the deformable stage. + + Args: + init_affine: 'Affine' to register an affine first, or 'Identity'. + """ + if init_affine not in ("Affine", "Identity"): + self.log_error("Invalid init_affine: %s", init_affine) + raise ValueError(f"Invalid init_affine: {init_affine}") + self.init_affine = init_affine + + def set_deformable_options(self, **options: Any) -> None: + """Set ``GreedyRegistration`` options for the deformable stage. + + Args: + **options: Any of ``cc_kernel_size``, ``cc_kernel_type``, + ``smooth_grad_sigma``, ``smooth_warp_sigma``, ``optimizer``, + ``optimizer_lr``, ``deformation_type``, ``integrator_n``, + ``blur``, ``max_tolerance_iters``, ``tolerance``. + """ + unknown = set(options) - self._DEFORMABLE_OPTIONS + if unknown: + raise ValueError(f"Invalid deformable options: {sorted(unknown)}") + self.deformable_options.update(options) + + def configure_for_heart_ct(self) -> None: + """Apply the tuned deformable settings for cardiac 4D CT. + + Values from experiments/Registration-Fireants_vs_Greedy/results/REPORT.md: + deformable-only (identity initial transform) CC registration. + """ + self.set_transform_type("Deformable") + self.set_metric("CC") + self.set_number_of_iterations([80, 40, 20]) + self.set_scales([4.0, 2.0, 1.0]) + self.set_init_affine("Identity") + self.set_deformable_options( + optimizer_lr=1.0, + smooth_warp_sigma=1.0, + smooth_grad_sigma=1.0, + cc_kernel_size=9, + ) + + def configure_for_lung_ct(self) -> None: + """Apply the tuned deformable settings for lung 4D CT. + + Values from experiments/Registration-Fireants_vs_Greedy/results/REPORT.md: + deformable-only (identity initial transform) CC registration at coarser + scales than the heart. + """ + self.set_transform_type("Deformable") + self.set_metric("CC") + self.set_number_of_iterations([20, 10, 5]) + self.set_scales([8.0, 4.0, 2.0]) + self.set_init_affine("Identity") + self.set_deformable_options( + optimizer_lr=1.0, + smooth_warp_sigma=1.0, + smooth_grad_sigma=1.0, + cc_kernel_size=9, + ) + + def _scales(self) -> list[float]: + """Downsampling factor per level, coarse to fine.""" + if self.scales is not None: + return list(self.scales) + n = len(self.number_of_iterations) + return [float(2 ** (n - 1 - i)) for i in range(n)] + + def _to_fireants_image( + self, image: itk.Image, mask: Optional[itk.Image] = None, masked: bool = False + ) -> Any: + """Wrap an ITK image as a batched FireANTs image on the GPU. + + When ``masked`` is True the mask is appended as the last channel, as the + FireANTs masked losses expect; a missing mask becomes all ones. + """ + fireants_image = _try_import_fireants("io.image") + sitk_image = ProcessImages().convert_itk_image_to_sitk(image) + wrapped = fireants_image.Image(sitk_image, device="cuda") + if masked: + imagemask = _try_import_fireants("io.imagemask") + if mask is None: + wrapped = imagemask.generate_and_apply_mask_allones(wrapped) + else: + sitk_mask = ProcessImages().convert_itk_image_to_sitk(mask) + wrapped = imagemask.apply_mask_to_image( + wrapped, fireants_image.Image(sitk_mask, device="cuda") + ) + return fireants_image.BatchedImages([wrapped]) + + @staticmethod + def _matrix_to_itk_affine(mat_4x4: NDArray[np.float64]) -> itk.Transform: + """Convert a 4x4 physical-space ``y = A x + t`` matrix to an ITK affine.""" + mat_4x4 = np.asarray(mat_4x4, dtype=np.float64) + if mat_4x4.shape != (4, 4): + raise ValueError(f"Expected 4x4 matrix, got shape {mat_4x4.shape}") + affine_tfm = itk.AffineTransform[itk.D, 3].New() + affine_tfm.SetCenter(itk.Point[itk.D, 3]()) + affine_tfm.SetMatrix(itk.GetMatrixFromArray(mat_4x4[:3, :3].copy())) + translation = itk.Vector[itk.D, 3]() + for i in range(3): + translation[i] = float(mat_4x4[i, 3]) + affine_tfm.SetTranslation(translation) + return affine_tfm + + def _registration_method_affine_or_rigid( + self, + fixed: Any, + moving: Any, + loss_type: str, + transform_type: str, + ) -> tuple[NDArray[np.float64], float]: + """Run FireANTs rigid or affine registration. Returns (4x4 matrix, loss).""" + if transform_type == "Rigid": + registration_class = _try_import_fireants( + "registration.rigid" + ).RigidRegistration + else: + registration_class = _try_import_fireants( + "registration.affine" + ).AffineRegistration + registration = registration_class( + self._scales(), + list(self.number_of_iterations), + fixed, + moving, + loss_type=loss_type, + progress_bar=False, + ) + registration.optimize() + + with torch.no_grad(): + matrix_fn = ( + registration.get_rigid_matrix + if transform_type == "Rigid" + else registration.get_affine_matrix + ) + mat = matrix_fn(homogenous=True)[0].double().cpu().numpy() + moved = registration.evaluate(fixed, moving) + loss = float(registration.loss_fn(moved, fixed())) + self.log_info("FireANTs %s registration loss: %s", transform_type, loss) + return mat, loss + + def _registration_method_deformable( + self, + fixed: Any, + moving: Any, + loss_type: str, + reference_image: itk.Image, + ) -> tuple[NDArray[np.float64], itk.Transform, float]: + """Run FireANTs affine then deformable registration. + + The deformable stage starts from the affine result (or identity). Returns the 4x4 + affine, the residual displacement transform on ``reference_image``'s + grid (applied before the affine, so ``y = A (x + d(x))``), and the loss. + """ + if self.init_affine == "Affine": + mat, _ = self._registration_method_affine_or_rigid( + fixed, moving, loss_type, "Affine" + ) + else: + mat = np.eye(4) + registration = _try_import_fireants("registration.greedy").GreedyRegistration( + self._scales(), + list(self.number_of_iterations), + fixed, + moving, + loss_type=loss_type, + init_affine=torch.from_numpy(mat).float().cuda()[None], + progress_bar=False, + **self.deformable_options, + ) + registration.optimize() + + with torch.no_grad(): + # The warp is a displacement in the moving image's normalized torch + # coordinates, applied after the affine: y = A x + L d(x), with L the + # linear part of moving's torch-to-physical map. Re-expressing it + # before the affine gives d' = A^-1 L d, so the pair composes as an + # ITK composite [affine, displacement]. + warp = registration.get_warp_parameters(fixed, moving)["grid"][0] + torch_to_physical = moving.get_torch2phy()[0, :3, :3].float() + affine_inverse = torch.from_numpy(np.linalg.inv(mat[:3, :3])).float() + to_pre_affine = affine_inverse.cuda() @ torch_to_physical + displacement = torch.einsum("ij,zyxj->zyxi", to_pre_affine, warp) + displacement_array = displacement.double().cpu().numpy() + moved = registration.evaluate(fixed, moving) + loss = float(registration.loss_fn(moved, fixed())) + self.log_info("FireANTs Deformable registration loss: %s", loss) + + displacement_image = ProcessImages().convert_array_to_image_of_vectors( + displacement_array, reference_image, itk.D + ) + displacement_transform = itk.DisplacementFieldTransform[itk.D, 3].New() + displacement_transform.SetDisplacementField(displacement_image) + return mat, displacement_transform, loss + + def registration_method( + self, + moving_image: itk.Image, + moving_mask: Optional[itk.Image] = None, + moving_labelmap: Optional[itk.Image] = None, + moving_image_pre: Optional[itk.Image] = None, + ) -> dict[str, Union[itk.Transform, float]]: + """Register moving image to fixed image using FireANTs. + + Returns a dict with "fixed_to_moving_transform", "moving_to_fixed_transform", + and "loss". As with the other image-registration backends, + fixed_to_moving_transform warps the moving image onto the fixed grid and + moving_to_fixed_transform warps the fixed image onto the moving grid; point + and landmark warps use the opposite transform from image warps (see + docs/developer/transform_conventions). + + Raises: + ValueError: If the fixed image is not set. + NotImplementedError: For ``moving_labelmap`` or ``fixed_labelmap``, + which FireANTs cannot use as a similarity term. Masks are used + through the FireANTs masked loss; a mask on only one image is + paired with an all-ones mask on the other. + """ + if self.fixed_image is None or self.fixed_image_pre is None: + raise ValueError("Fixed image must be set before registration.") + for name, value in ( + ("moving_labelmap", moving_labelmap), + ("fixed_labelmap", self.fixed_labelmap), + ): + if value is not None: + raise NotImplementedError( + f"RegisterImagesFireants does not support {name}." + ) + + moving_pre = moving_image_pre if moving_image_pre is not None else moving_image + masked = moving_mask is not None or self.fixed_mask is not None + loss_type = ("masked_" if masked else "") + self._FIREANTS_LOSS[self.metric] + fixed = self._to_fireants_image(self.fixed_image_pre, self.fixed_mask, masked) + moving = self._to_fireants_image(moving_pre, moving_mask, masked) + + fixed_to_moving_transform: itk.Transform + moving_to_fixed_transform: itk.Transform + if self.transform_type == "Deformable": + mat, displacement, loss_val = self._registration_method_deformable( + fixed, moving, loss_type, self.fixed_image_pre + ) + affine = self._matrix_to_itk_affine(mat) + inverse_affine = itk.AffineTransform[itk.D, 3].New() + affine.GetInverse(inverse_affine) + forward = itk.CompositeTransform[itk.D, 3].New() + forward.AddTransform(affine) + forward.AddTransform(displacement) + inverse = itk.CompositeTransform[itk.D, 3].New() + inverse.AddTransform( + ProcessTransforms().invert_displacement_field_transform(displacement) + ) + inverse.AddTransform(inverse_affine) + fixed_to_moving_transform = forward + moving_to_fixed_transform = inverse + else: + mat, loss_val = self._registration_method_affine_or_rigid( + fixed, moving, loss_type, self.transform_type + ) + fixed_to_moving_transform = self._matrix_to_itk_affine(mat) + moving_to_fixed_transform = itk.AffineTransform[itk.D, 3].New() + fixed_to_moving_transform.GetInverse(moving_to_fixed_transform) + + return { + "fixed_to_moving_transform": fixed_to_moving_transform, + "moving_to_fixed_transform": moving_to_fixed_transform, + "loss": loss_val, + } diff --git a/src/monai_physio/train_physicsnemo_physics_informed_motion.py b/src/monai_physio/train_physicsnemo_physics_informed_motion.py index 9a5b25fd..d5dcacde 100644 --- a/src/monai_physio/train_physicsnemo_physics_informed_motion.py +++ b/src/monai_physio/train_physicsnemo_physics_informed_motion.py @@ -582,6 +582,11 @@ def set_lambda_physics_warmup(self, warmup_epochs: int) -> None: term shape real, non-uniform motion first, before the physics term is weighted heavily enough to matter. + The warmup cannot rescue a target weight that is too large: once the + ramp ends, the network collapses anyway if the weighted strain + energy of the measured motion outweighs the data loss of predicting + no motion. Size the target from those two numbers first. + Args: warmup_epochs: Number of epochs to ramp over. ``0`` (default) disables warmup: ``lambda_physics`` is held at the value @@ -846,8 +851,9 @@ def _log_epoch(self, context: DistributedContext, epoch: int, epochs: int) -> No "physics loss has been ~0 for 5 consecutive epochs while " "lambda_physics=%.4g; the network likely collapsed to the " "trivial uniform-displacement solution described in " - "set_lambda_physics_warmup()'s docstring. Consider a longer " - "warmup or a lower lambda_physics target.", + "set_lambda_physics_warmup()'s docstring. Lower the " + "lambda_physics target until the weighted strain energy of " + "the measured motion sits below the data loss.", self.lambda_physics, ) diff --git a/tests/test_register_images_fireants.py b/tests/test_register_images_fireants.py new file mode 100644 index 00000000..d2ed0116 --- /dev/null +++ b/tests/test_register_images_fireants.py @@ -0,0 +1,385 @@ +"""Tests for FireANTs-based image registration (GPU).""" + +import hashlib +from pathlib import Path +from typing import Any + +import itk +import numpy as np +import pytest +from scipy.ndimage import gaussian_filter + +import monai_physio +from monai_physio.process_images import ProcessImages +from monai_physio.process_transforms import ProcessTransforms +from monai_physio.register_images_base import RegisterImagesBase +from monai_physio.register_images_fireants import RegisterImagesFireants +from monai_physio.workflow_reconstruct_highres_4d_ct import ( + WorkflowReconstructHighres4DCT, +) + +from .conftest import KnownAffineCase, KnownShiftCase + +T_SHIFT = 0.5 +T_AFFINE = 1.0 + + +def test_s1_ac1_subclass_and_public_export() -> None: + """The class subclasses RegisterImagesBase and is exported.""" + assert issubclass(RegisterImagesFireants, RegisterImagesBase) + assert monai_physio.RegisterImagesFireants is RegisterImagesFireants + assert "RegisterImagesFireants" in monai_physio.__all__ + + +def test_s1_ac2_invalid_transform_type_and_metric_raise_valueerror() -> None: + """Invalid settings raise ValueError with Greedy's messages.""" + registrar = RegisterImagesFireants() + for transform_type in ("Rigid", "Affine", "Deformable"): + registrar.set_transform_type(transform_type) + assert registrar.transform_type == transform_type + for metric in ("CC", "Mattes", "MeanSquares"): + registrar.set_metric(metric) + assert registrar.metric == metric + with pytest.raises(ValueError, match="Invalid transform type: Invalid"): + registrar.set_transform_type("Invalid") + with pytest.raises(ValueError, match="Invalid metric: Invalid"): + registrar.set_metric("Invalid") + + +@pytest.mark.requires_gpu +@pytest.mark.parametrize("transform_type", ["Rigid", "Affine"]) +def test_s1_ac3_known_shift_recovered( + known_shift_case: KnownShiftCase, transform_type: str +) -> None: + """Rigid and Affine recover a known shift within T_SHIFT.""" + registrar = RegisterImagesFireants() + registrar.set_modality("ct") + registrar.set_transform_type(transform_type) + registrar.set_number_of_iterations([200, 100, 50]) + registrar.set_fixed_image(known_shift_case.fixed) + + result = registrar.register(moving_image=known_shift_case.moving) + + error_mm = known_shift_case.center_error_mm(result["fixed_to_moving_transform"]) + print(f"\nFireANTs {transform_type} known-shift error: {error_mm:.3f} mm") + assert error_mm <= T_SHIFT + + +@pytest.mark.requires_gpu +@pytest.mark.parametrize("offset_z", [0.0, 1800.0]) +def test_s1_ac4_known_affine_near_origin_and_z1800(offset_z: float) -> None: + """Affine recovers a known affine, near the origin and at z ~ 1800 mm.""" + case = KnownAffineCase(origin_offset_mm=(0.0, 0.0, offset_z)) + registrar = RegisterImagesFireants() + registrar.set_modality("ct") + registrar.set_transform_type("Affine") + # The blocky synthetic volume has no local contrast for CC to correlate. + registrar.set_metric("MeanSquares") + registrar.set_number_of_iterations([100, 50, 25]) + registrar.set_fixed_image(case.fixed) + + result = registrar.register(moving_image=case.moving) + + errors = case.probe_errors_mm(result["fixed_to_moving_transform"]) + print(f"\nFireANTs known-affine z={offset_z}: {np.round(errors, 3).tolist()} mm") + assert float(errors.max()) <= T_AFFINE + + +@pytest.mark.requires_gpu +def test_s1_ac5_return_keys_and_ncc_not_worse( + known_shift_case: KnownShiftCase, +) -> None: + """Return dict has the three keys and registration does not lower NCC.""" + registrar = RegisterImagesFireants() + registrar.set_modality("ct") + registrar.set_transform_type("Affine") + registrar.set_number_of_iterations([200, 100, 50]) + registrar.set_fixed_image(known_shift_case.fixed) + + result = registrar.register(moving_image=known_shift_case.moving) + + assert set(result) == { + "fixed_to_moving_transform", + "moving_to_fixed_transform", + "loss", + } + assert isinstance(result["loss"], float) + ncc = known_shift_case.foreground_ncc(result["fixed_to_moving_transform"]) + initial_ncc = known_shift_case.unregistered_ncc() + assert ncc >= initial_ncc + inverse = result["moving_to_fixed_transform"] + assert isinstance(inverse, itk.Transform) + + +T_DEF_GAIN = 0.30 +T_INV = 0.5 + + +def _smooth_volume(size: int = 64, spacing_mm: float = 2.0) -> itk.Image: + """Smoothed random texture: local contrast everywhere for CC.""" + rng = np.random.default_rng(0) + array = gaussian_filter(rng.standard_normal((size,) * 3), sigma=2.0) + array = (1000.0 * (array - array.min()) / np.ptp(array)).astype(np.float32) + image = itk.image_from_array(array) + image.SetSpacing([spacing_mm] * 3) + return image + + +def _smooth_warp(image: itk.Image, amplitude_mm: float = 4.0) -> itk.Transform: + """Smooth sinusoidal displacement field on the grid of ``image``.""" + shape = itk.array_from_image(image).shape + z, y, x = np.meshgrid(*[np.linspace(0, np.pi, n) for n in shape], indexing="ij") + field = np.stack( + [ + amplitude_mm * np.sin(y) * np.sin(z), + amplitude_mm * np.sin(x) * np.sin(z), + -amplitude_mm * np.sin(x) * np.sin(y), + ], + axis=-1, + ) + vectors = ProcessImages().convert_array_to_image_of_vectors(field, image, itk.D) + transform = itk.DisplacementFieldTransform[itk.D, 3].New() + transform.SetDisplacementField(vectors) + return transform + + +def _translation(offset: tuple[float, float, float]) -> itk.Transform: + transform = itk.TranslationTransform[itk.D, 3].New() + transform.SetOffset(list(offset)) + return transform + + +def _one_minus_ncc(a: itk.Image, b: itk.Image) -> float: + x = itk.array_from_image(a).ravel().astype(np.float64) + y = itk.array_from_image(b).ravel().astype(np.float64) + x -= x.mean() + y -= y.mean() + return 1.0 - float((x * y).sum() / np.sqrt((x**2).sum() * (y**2).sum())) + + +def _warp_pair() -> tuple[itk.Image, itk.Image]: + fixed = _smooth_volume() + moving = ProcessTransforms().transform_image( + fixed, _smooth_warp(fixed), fixed, interpolation_method="linear" + ) + return fixed, moving + + +def _deformable_registrar(fixed: itk.Image) -> RegisterImagesFireants: + registrar = RegisterImagesFireants() + registrar.set_modality("ct") + registrar.set_transform_type("Deformable") + registrar.set_number_of_iterations([100, 50, 25]) + registrar.set_fixed_image(fixed) + return registrar + + +@pytest.mark.requires_gpu +def test_s2_ac1_deformable_gain_on_smooth_warp() -> None: + """Deformable lowers (1 - NCC) by at least T_DEF_GAIN on a smooth warp.""" + fixed, moving = _warp_pair() + registrar = _deformable_registrar(fixed) + + result = registrar.register(moving_image=moving) + + warped = ProcessTransforms().transform_image( + moving, + result["fixed_to_moving_transform"], + fixed, + interpolation_method="linear", + ) + before = _one_minus_ncc(fixed, moving) + after = _one_minus_ncc(fixed, warped) + print(f"\nFireANTs deformable 1-NCC: {before:.4f} -> {after:.4f}") + assert after <= (1.0 - T_DEF_GAIN) * before + + +@pytest.mark.requires_gpu +def test_s2_ac2_transforms_roundtrip_hdf_and_inverse_consistency( + tmp_path: Path, +) -> None: + """Both transforms survive an .hdf round trip and invert each other.""" + fixed, moving = _warp_pair() + registrar = _deformable_registrar(fixed) + result = registrar.register(moving_image=moving) + + size = np.array(itk.size(fixed)) + lo, hi = (size * 0.25).astype(int), (size * 0.75).astype(int) + points = [ + fixed.TransformIndexToPhysicalPoint([int(i), int(j), int(k)]) + for i in range(lo[0], hi[0], 6) + for j in range(lo[1], hi[1], 6) + for k in range(lo[2], hi[2], 6) + ] + + loaded = {} + for key in ("fixed_to_moving_transform", "moving_to_fixed_transform"): + transform = result[key] + assert isinstance(transform, itk.Transform) + path = tmp_path / f"{key}.hdf" + itk.transformwrite([transform], str(path), compression=True) + loaded[key] = itk.transformread(str(path))[0] + for point in points[:20]: + expected = np.array(list(transform.TransformPoint(point))) + actual = np.array(list(loaded[key].TransformPoint(point))) + assert np.abs(expected - actual).max() < 1e-3 + + forward = loaded["fixed_to_moving_transform"] + inverse = loaded["moving_to_fixed_transform"] + residual = np.mean( + [ + np.linalg.norm( + np.array(list(inverse.TransformPoint(forward.TransformPoint(p)))) + - np.array(list(p)) + ) + for p in points + ] + ) + print(f"\nFireANTs inverse-consistency residual: {residual:.4f} mm") + assert residual <= T_INV + + +@pytest.mark.requires_gpu +def test_s2_ac3_workflow_on_slicer_heart_small(test_images: list[Any]) -> None: + """Injected into the 4D CT workflow, one image per phase on the reference grid.""" + registrar = RegisterImagesFireants() + registrar.set_modality("ct") + registrar.set_transform_type("Deformable") + registrar.set_number_of_iterations([1, 0]) + reference = test_images[0] + workflow = WorkflowReconstructHighres4DCT( + time_series_images=test_images[:3], + reference_image=reference, + reference_time_frame=0, + registration_method=registrar, + ) + + result = workflow.process() + + reconstructed = result["reconstructed_images"] + assert len(reconstructed) == 3 + for image in reconstructed: + assert tuple(itk.size(image)) == tuple(itk.size(reference)) + assert np.allclose(itk.spacing(image), itk.spacing(reference)) + assert np.allclose(itk.origin(image), itk.origin(reference)) + + +def _conflicting_pair(case: KnownShiftCase) -> tuple[itk.Image, itk.Image]: + """Moving image shifted by the case's shift inside a box and by a conflicting + amount outside it. Returns (moving, box mask).""" + conflicting = ProcessTransforms().transform_image( + case.fixed, + _translation((-14.0, 10.0, -9.0)), + case.fixed, + interpolation_method="linear", + ) + shape = itk.array_from_image(case.fixed).shape + box = np.zeros(shape, dtype=np.uint8) + box[ + shape[0] // 4 : 3 * shape[0] // 4, + shape[1] // 4 : 3 * shape[1] // 4, + shape[2] // 4 : 3 * shape[2] // 4, + ] = 1 + moving_array = np.where( + box.astype(bool), + itk.array_from_image(case.moving), + itk.array_from_image(conflicting), + ) + moving = itk.image_from_array(moving_array.astype(np.float32)) + moving.CopyInformation(case.fixed) + mask = itk.image_from_array(box) + mask.CopyInformation(case.fixed) + return moving, mask + + +@pytest.mark.requires_gpu +def test_s2_ac4_mask_labelmap_policy(known_shift_case: KnownShiftCase) -> None: + """A mask changes the result and ignores out-of-mask structure; labelmaps raise.""" + moving, mask = _conflicting_pair(known_shift_case) + + def run(masked: bool) -> float: + registrar = RegisterImagesFireants() + registrar.set_modality("ct") + registrar.set_transform_type("Affine") + registrar.set_mask_dilation(0.0) + registrar.set_number_of_iterations([200, 100, 50]) + registrar.set_fixed_image(known_shift_case.fixed) + if masked: + registrar.set_fixed_mask(mask) + result = registrar.register( + moving_image=moving, moving_mask=mask if masked else None + ) + return known_shift_case.center_error_mm(result["fixed_to_moving_transform"]) + + unmasked_error = run(masked=False) + masked_error = run(masked=True) + print( + f"\nFireANTs center error unmasked {unmasked_error:.3f} mm, " + f"masked {masked_error:.3f} mm" + ) + assert masked_error <= T_SHIFT + assert unmasked_error > masked_error + + registrar = RegisterImagesFireants() + registrar.set_fixed_image(known_shift_case.fixed) + with pytest.raises(NotImplementedError, match="moving_labelmap"): + registrar.register(moving_image=moving, moving_labelmap=mask) + registrar.set_fixed_labelmap(mask) + with pytest.raises(NotImplementedError, match="fixed_labelmap"): + registrar.register(moving_image=moving) + + +@pytest.mark.requires_gpu +def test_s2_ac6_inputs_not_mutated() -> None: + """Registration leaves the fixed, moving, and mask images unchanged.""" + fixed, moving = _warp_pair() + mask = itk.image_from_array(np.ones((64, 64, 64), dtype=np.uint8)) + mask.CopyInformation(fixed) + + def digest(image: itk.Image) -> tuple[str, tuple[Any, ...]]: + pixels = hashlib.sha256(itk.array_from_image(image).tobytes()).hexdigest() + return pixels, (tuple(image.GetSpacing()), tuple(image.GetOrigin())) + + images = (fixed, moving, mask) + before = [digest(image) for image in images] + registrar = _deformable_registrar(fixed) + registrar.set_number_of_iterations([10, 5]) + registrar.set_fixed_mask(mask) + registrar.register(moving_image=moving, moving_mask=mask) + assert [digest(image) for image in images] == before + + +def test_s4_invalid_deformable_options_scales_and_init_affine_raise() -> None: + """Unknown deformable options, bad scales and bad init_affine are rejected.""" + registrar = RegisterImagesFireants() + with pytest.raises(ValueError, match="Invalid deformable options"): + registrar.set_deformable_options(init_scale=2) + with pytest.raises(ValueError, match="equal length"): + registrar.set_scales([2.0, 1.0]) + with pytest.raises(ValueError, match="Invalid init_affine"): + registrar.set_init_affine("Moments") + + +@pytest.mark.requires_gpu +def test_s4_tuned_options_deformable_gain_on_smooth_warp() -> None: + """Identity init, explicit scales and tuned options still reduce 1 - NCC.""" + fixed, moving = _warp_pair() + registrar = _deformable_registrar(fixed) + registrar.set_number_of_iterations([40, 20, 10]) + registrar.set_scales([4.0, 2.0, 1.0]) + registrar.set_init_affine("Identity") + registrar.set_deformable_options( + smooth_warp_sigma=1.0, cc_kernel_size=9, optimizer_lr=1.0 + ) + + result = registrar.register(moving_image=moving) + + warped = ProcessTransforms().transform_image( + moving, + result["fixed_to_moving_transform"], + fixed, + interpolation_method="linear", + ) + assert _one_minus_ncc(fixed, warped) <= (1.0 - T_DEF_GAIN) * _one_minus_ncc( + fixed, moving + ) diff --git a/tests/test_segment_chest_total_segmentator.py b/tests/test_segment_chest_total_segmentator.py index 929e9063..4bfa598f 100644 --- a/tests/test_segment_chest_total_segmentator.py +++ b/tests/test_segment_chest_total_segmentator.py @@ -5,6 +5,7 @@ functionality on two time points from the converted 3D data. """ +import json from pathlib import Path from typing import Any @@ -367,3 +368,40 @@ def test_license_check_tracks_totalsegmentators_own_offline_gate( lambda: ("missing_license", "ERROR: A license number has not been set."), ) assert SegmentChestTotalSegmentator._academic_license_is_valid() is False + + +def test_segmentation_method_disables_usage_stats( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """``segmentation_method`` must turn off TotalSegmentator's telemetry. + + Points ``TOTALSEG_HOME_DIR`` at an isolated, empty directory so this + does not touch the real machine-wide config, then stops the call right + after the config is written (before the GPU-only ``totalsegmentator`` + call) and reads back ``config.json`` to confirm ``send_usage_stats`` is + ``False``. This is the documented key TotalSegmentator's own config and + telemetry code reads; a misspelled key (e.g. ``send_telemetry``) would + leave usage stats on while this test still passes if it only checked + that *some* key was set. + """ + monkeypatch.setenv("TOTALSEG_HOME_DIR", str(tmp_path)) + + class _StopBeforeGpuCall(Exception): + pass + + def _raise(*args: Any, **kwargs: Any) -> Any: + raise _StopBeforeGpuCall + + monkeypatch.setattr("totalsegmentator.python_api.totalsegmentator", _raise) + + segmenter = SegmentChestTotalSegmentator() + image = itk.image_from_array(np.zeros((2, 2, 2), dtype=np.float32)) + + with pytest.raises(_StopBeforeGpuCall): + segmenter.segmentation_method(image) + + config_file = tmp_path / "config.json" + assert config_file.exists(), "setup_totalseg did not write a config file" + config = json.loads(config_file.read_text()) + assert config["send_usage_stats"] is False + assert config["statistics_disclaimer_shown"] is True diff --git a/tutorials/parameters_duke_heart_physics_informed.py b/tutorials/parameters_duke_heart_physics_informed.py index 1fcad13d..077a0c52 100644 --- a/tutorials/parameters_duke_heart_physics_informed.py +++ b/tutorials/parameters_duke_heart_physics_informed.py @@ -54,15 +54,21 @@ class ParametersDukeHeartPhysicsInformed(ParametersDukeHeartLabelmaps): lambda_physics: Weight of the physics residual against the displacement loss. The two are not in the same units -- displacement is scored normalized, the residual in kilopascals -- so this is a value to - sweep rather than a value to trust. + sweep rather than a value to trust. Measured on the training + cases, the measured motion itself carries a median strain energy + of about 4 kPa (up to about 45 at end systole), while predicting + no motion at all costs a data loss of only about 7e-4. The weighted + energy of the true motion must therefore sit well below the data + loss the data-only model converges to (about 1e-4), or the cheapest + answer is a heart that does not move: 0.1 collapsed to exactly that. lambda_physics_warmup_epochs: Epochs over which ``lambda_physics`` ramps linearly from 0 up to its target value, instead of applying full-strength from epoch 0. Cold-start weights make the strain - energy huge relative to the data loss, so without a warmup its - gradient can dominate the combined loss and pull the network - straight to the energy's own global minimum -- zero strain - everywhere -- which trivially zeroes the physics term but ignores - the data. ``0`` disables warmup. + energy huge relative to the data loss, so the warmup lets the data + term shape the motion first. It only helps the start of training: + a target ``lambda_physics`` that is too large still pulls the + network to zero strain once the ramp reaches it. ``0`` disables + warmup. number_of_epochs: Training epochs, matching Tutorial 9 so the two are comparable. number_of_epochs_test: Same, under ``ProcessTests.running_as_test``. @@ -78,7 +84,7 @@ class ParametersDukeHeartPhysicsInformed(ParametersDukeHeartLabelmaps): mu_kpa: float = 10.0 lambda_lame_kpa: float = 100.0 - lambda_physics: float = 0.1 + lambda_physics: float = 1.0e-5 lambda_physics_warmup_epochs: int = 100 number_of_epochs: int = 1500 diff --git a/tutorials/tutorial_13_heart_and_lung_motion.py b/tutorials/tutorial_13_heart_and_lung_motion.py index 1a9626d1..da6add63 100644 --- a/tutorials/tutorial_13_heart_and_lung_motion.py +++ b/tutorials/tutorial_13_heart_and_lung_motion.py @@ -223,7 +223,7 @@ # Fraction of a cardiac cycle advanced per respiratory phase. The heart is # decoupled from breathing, so a value < 1 means a single beat continues into # the next phase (1.0 locks exactly one full beat to each phase). - cardiac_cycles_per_phase = 0.75 + cardiac_cycles_per_phase = 0.333333333333 # Gaussian sigma (mm) that spreads each rhythm's surface-sampled # displacements into a continuous field. The lung shell is sparse relative # to the thorax it has to fill, so it is spread further than the heart. @@ -588,7 +588,7 @@ def stage_transforms( ) for direction in directions } - pct = int(round(stage * 100)) + pct = round(stage * 100) itk.imwrite( fields["forward"]["deformation_field"], str(output_dir / f"deformation_field_{tag}_s{pct:03d}.mha"), diff --git a/tutorials/tutorial_18_duke_heart_physics_informed_motion_infer.py b/tutorials/tutorial_18_duke_heart_physics_informed_motion_infer.py index bab7968e..7b9a8a25 100644 --- a/tutorials/tutorial_18_duke_heart_physics_informed_motion_infer.py +++ b/tutorials/tutorial_18_duke_heart_physics_informed_motion_infer.py @@ -52,6 +52,7 @@ # Imports from __future__ import annotations +import contextlib import csv import logging from pathlib import Path @@ -64,6 +65,7 @@ from monai_physio import ( ConvertVTKToUSD, EvaluateMovementDukeHeart, + InferPhysicsNeMoMGN, ProcessTests, WorkflowEvaluateMovement, WorkflowInferMovement, @@ -133,11 +135,9 @@ def _render_stress(surface: pv.DataSet, plot_file: Path) -> Path: field needs its own plotter. """ xvfb_started = False - try: + with contextlib.suppress(Exception): pv.start_xvfb() xvfb_started = True - except Exception: - pass try: plotter = pv.Plotter(off_screen=True, window_size=[800, 600]) plotter.add_mesh( @@ -154,10 +154,8 @@ def _render_stress(surface: pv.DataSet, plot_file: Path) -> Path: # Paired with the guarded start above, so environments without Xvfb # (Windows, pyvista >= 0.48) are unaffected. if xvfb_started: - try: + with contextlib.suppress(Exception): pv.stop_xvfb() - except Exception: - pass return plot_file @@ -278,7 +276,12 @@ def _render_stress(surface: pv.DataSet, plot_file: Path) -> Path: logger.info("%s", "=" * 48) infer_workflow = WorkflowInferPhysicsNeMo( - model_directory=model_directory, epoch=epoch, log_level=log_level + model_directory=model_directory, + inference_method=InferPhysicsNeMoMGN( + model_tag="physics_informed_motion", log_level=log_level + ), + epoch=epoch, + log_level=log_level, ) evaluate_workflow = WorkflowEvaluateMovement( movement_workflow=WorkflowInferMovement( From 2eced98460894b70b5f00f295ccd3099684264f2 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Thu, 15 Oct 2026 12:02:35 -0400 Subject: [PATCH 2/4] ENH: Fix ruff warnings and coderabbit comments --- .agents/skills/loop-kit/SKILL.md | 13 ++- .github/scripts/build_dashboard.py | 12 +-- .gitignore | 1 + .ruff.toml | 4 + data/TCIA-4DLung/convert.py | 1 + docs/conf.py | 4 +- .../1-input_meshes_to_input_surfaces.py | 2 +- .../1-initial_registration.py | 10 +- .../2-generate_segmentation.py | 3 +- .../simpleware_heart_segmentation.py | 2 +- .../0-heart_and_lungs_beating_heart.py | 4 +- experiments/Lung-VesselsAirways/0-GenData.py | 4 +- .../benchmark.py | 7 +- src/monai_physio/cli/convert_image_to_usd.py | 4 +- src/monai_physio/cli/convert_vtk_to_usd.py | 6 +- src/monai_physio/convert_vtk_to_usd.py | 6 +- src/monai_physio/download_data.py | 11 ++- src/monai_physio/process_images.py | 2 +- src/monai_physio/process_labelmaps.py | 2 +- src/monai_physio/process_tests.py | 7 +- src/monai_physio/process_transforms.py | 4 +- src/monai_physio/process_usd.py | 15 ++- src/monai_physio/register_images_ants.py | 4 +- src/monai_physio/register_images_fireants.py | 23 ++++- src/monai_physio/register_images_greedy.py | 6 +- src/monai_physio/register_images_icon.py | 4 +- .../register_models_distance_maps.py | 6 +- src/monai_physio/register_models_icp.py | 4 +- .../register_time_series_images.py | 4 +- src/monai_physio/report_evaluate_movement.py | 46 +++++---- src/monai_physio/segment_heart_simpleware.py | 1 + src/monai_physio/train_physicsnemo_base.py | 2 +- src/monai_physio/vtk_to_usd/__init__.py | 38 +++---- .../vtk_to_usd/primvar_derivations.py | 2 +- .../vtk_to_usd/usd_mesh_converter.py | 15 +-- src/monai_physio/vtk_to_usd/usd_utils.py | 7 +- src/monai_physio/vtk_to_usd/vtk_reader.py | 11 +-- tests/conftest.py | 14 +-- tests/test_register_images_fireants.py | 98 +++++++++++++++++++ .../test_workflow_create_statistical_model.py | 4 +- ..._06_duke_heart_create_statistical_model.py | 9 +- ...orial_06_heart_create_statistical_model.py | 9 +- ...torial_06_lung_create_statistical_model.py | 9 +- ..._heart_fit_statistical_model_to_patient.py | 5 +- ..._heart_fit_statistical_model_to_patient.py | 5 +- ...7_lung_fit_statistical_model_to_patient.py | 5 +- .../tutorial_13_heart_and_lung_motion.py | 6 +- .../tutorial_15_duke_heart_leave_one_out.py | 22 +++-- tutorials/tutorial_15_lung_leave_one_out.py | 28 +++--- utils/ai_agent_github_reviews.py | 2 +- utils/mp4_to_gif.py | 6 +- utils/setup_feature_worktree.py | 2 + 52 files changed, 341 insertions(+), 180 deletions(-) diff --git a/.agents/skills/loop-kit/SKILL.md b/.agents/skills/loop-kit/SKILL.md index 2bd5aed9..f67ecaad 100644 --- a/.agents/skills/loop-kit/SKILL.md +++ b/.agents/skills/loop-kit/SKILL.md @@ -1,3 +1,8 @@ +--- +name: loop-kit +description: Generate a bounded, verified implement/verify/review loop (task contract, prompts, subagents, skills, verifier) for a given task. Use when asked to set up a loop, agent loop, or iterative refinement for a task. Does not implement the task. +--- + # Claude Prompt: Generate a Verified Loop Kit for Any Task ## How to use this prompt @@ -31,7 +36,7 @@ You are a **loop architect**. Your job is to turn the task below into a small, f The loop is **automated**. A human halt is a cost; every halt must be justified. 1. **Ask everything once, up front.** Before generating the kit, anticipate every question the run could raise and put them in ONE batch (use AskUserQuestion, or a single list): decisions and tolerances; data/mode choices; dependency installs and environment quirks (version pins, Windows/platform gaps, plugin conflicts); approval for expensive jobs; whether the supervisor may run the verifier itself; and **standing authorizations** for predictable amendments (verifier/environment fixes, baseline-relative lint rules, schema/path naming mismatches, API differences discovered in Phase 0). Do not trickle questions between steps. -2. **Standing authorization replaces per-step gates.** Record the answers in `TASK.md` and the supervisor skill. After that: the supervisor runs worker -> verifier -> reviewer -> decision and **auto-advances** to the next slice when the verifier is PASS and the reviewer has no blocking findings. [H] items are treated as confirmed under the standing authorization unless the human reserved them in the batch. Never end a turn with "waiting for acceptance". +2. **Standing authorization replaces per-step gates.** Record the answers in `TASK.md` and the supervisor skill. After that: the supervisor runs worker -> verifier -> reviewer -> decision and **auto-advances** to the next slice when the verifier is PASS, the reviewer has no blocking findings, and risk and uncertainty are low (no protected-path change, no unresolved hypothesis findings, no contract amendment pending). [H] items are treated as confirmed under the standing authorization unless the human reserved them in the batch; a reserved [H] item, or high risk or uncertainty, holds the slice for the human. Never end a turn with "waiting for acceptance". 3. **Self-resolve, then log.** Verifier/environment defects (outside the repo), schema/name mismatches, and pre-existing-lint baselines are fixed by the supervisor, re-verified under a new run id, and logged in `PROGRESS.md`. Contract-semantics changes that stay inside the pre-authorized classes are applied as amendments (re-freeze, refresh the control hashes) without asking. 4. **Halt only for true BLOCKED conditions:** repair limit exhausted; a protected-path change; a missing prerequisite the agent cannot create; an action outside the standing authorizations (commit/push/publish/messages, licensed downloads, new spend); or a decision that changes the task's intent. When halting, give one consolidated question list, not one at a time. 5. **Run to completion.** Continue through all slices and the final report in one go; report once at the end (plus a brief note on any self-resolved amendment). @@ -53,7 +58,7 @@ The loop is **automated**. A human halt is a cost; every halt must be justified. 4. `PROGRESS.md`, which is memory only. A line saying "done" or "approved" in a model-written file confers nothing. -3. **Small slices.** Split the work into ordered slices. Each slice gets its own acceptance criteria (ACs) and one worker run per attempt. Never start the next slice until the human accepts the current one. +3. **Small slices.** Split the work into ordered slices. Each slice gets its own acceptance criteria (ACs) and one worker run per attempt. Start the next slice only when the current one meets the auto-advance rule (Autonomy, rule 2); reserved [H] items and BLOCKED slices wait for the human. 4. **Every AC has an evidence type:** - **[V]:** a deterministic check that the verifier executes; - **[R]:** reviewer-confirmed from the output; @@ -95,7 +100,7 @@ Create a kit directory at `.agent-loops//`, plus the agent and skill | `README.md` | What the loop does and the one-time setup. Setup covers the isolated worktree or branch, git-ignoring `PROGRESS.md` and `outputs/`, freezing `TASK.md` (`sha256sum` with the hash stored **outside** the repo, plus `chmod a-w`), and installing the verifier outside the repo. Then give the exact commands for each phase, the optional `/goal` text, the limits, the smoke tests, the rollout order, and the baseline comparison. | | `TASK.md` | **Status:** DRAFT, plus a contract version. Then these numbered sections: (1) the requested behavior and public interface; (2) the slices table (ID, content, dependencies); (3) the ACs: global ones (G1 no regressions against the baseline, G2 the diff stays within allowed paths, G3 no weakened checks, G4 lint and format on changed files only, plus task-specific rules) and per-slice ones, each tagged [V], [R], or [H] and concrete enough to test; (4) allowed paths; (5) protected paths, where any change means BLOCKED; (6) approved tolerances and thresholds; (7) human decisions D1…Dn, never guessed; (8) out of scope. | | `REFERENCE.md` | The facts the agents may rely on: APIs and signatures, data formats, domain rules, and links to sources. Every item is marked "verify against the installed code." | -| `LOOP_INSTRUCTIONS.md` | The authority order, the phases, and the per-slice cycle diagram: worker, then verifier, then reviewer, then a decision (`READY_FOR_HUMAN_REVIEW`, `NEEDS_REVISION`, or `BLOCKED`), then the human gate. Then come the exact decision rules, the limits (repairs, wall-clock time, one run lock, interrupted runs marked `UNVERIFIED`, one bounded retry for transient infrastructure errors only), the evidence layout, the update rule for `PROGRESS.md`, the optional `/goal` text, five workflow smoke tests, and the human handoff checklist. The smoke tests are: a planted known defect is caught; no source change gives a quiet no-op; a restart keeps findings and attempt counts; a missing prerequisite gives BLOCKED; a forbidden edit is refused. | +| `LOOP_INSTRUCTIONS.md` | The authority order, the phases, and the per-slice cycle diagram: worker, then verifier, then reviewer, then a decision (`ADVANCE`, `NEEDS_REVISION`, `HOLD_FOR_HUMAN`, or `BLOCKED`), then the auto-advance gate (the human only for reserved [H] items, high risk or uncertainty, or BLOCKED). Then come the exact decision rules, the limits (repairs, wall-clock time, one run lock, interrupted runs marked `UNVERIFIED`, one bounded retry for transient infrastructure errors only), the evidence layout, the update rule for `PROGRESS.md`, the optional `/goal` text, five workflow smoke tests, and the human handoff checklist. The smoke tests are: a planted known defect is caught; no source change gives a quiet no-op; a restart keeps findings and attempt counts; a missing prerequisite gives BLOCKED; a forbidden edit is refused. | | `INIT_PROMPT.md` | Phase 0, which writes no task output. It records identity and isolation (repo, branch, base commit, clean tree, interpreter, tool versions) and runs an environment preflight (stop with named blockers). It maps the scoped context (exact files and lines), runs the baseline checks (record counts and pre-existing failures without fixing them), and gives a per-slice plan (files, tests per [V] AC, risks). It lists AC problems as a `contract_amendment.md` and makes recommendations for D1…Dn. It ends with `INIT_STATUS: READY \| BLOCKED`, then stops. | | `WORKER_PROMPT.md` | Takes `SLICE`, `RUN`, and an optional `MODE=REPAIR` plus a findings path. The worker reads the slice ACs, `PROGRESS.md` (to avoid repeating failed approaches), the init report, and the installed sources. It writes tests from the ACs (named `test___*` or an equivalent per check), keeps changes minimal and in scope, and runs the slice's [V] checks itself, pasting the real output. The "Never" list repeats the protected-path and side-effect rules. When stuck, it writes an amendment or names a blocker and stops. In repair mode it fixes accepted findings only. It writes `worker_summary.md` in a fixed format and appends to `PROGRESS.md` without ever writing pass status. | | `REVIEWER_PROMPT.md` | Fresh context with read-only tools. The reviewer must not read the worker's summary or earlier reviews. Inputs are the contract, the exact diff or output snapshot, `verification.json`, and the surrounding context. It works through a prioritized checklist: correctness against the ACs, the task-specific invariants from Step 1, default-path safety, contract integrity, and evidence gaps. A finding is blocking only with a location, the violated AC or invariant, a concrete failure scenario, and a verification step. Anything else is a `hypothesis`. Output is JSON: `slice`, `run_id`, `output_hash`, `ac_assessment[]`, `findings[]`, `unverified_requirements[]`, `scope_violations[]`, and an advisory `recommendation`. | @@ -113,7 +118,7 @@ Create: - `.claude/agents/-worker.md`: the tools needed for edits plus scoped `Bash` (tests, lint, `git diff` and `git status` only). The body is "Follow `.agent-loops//WORKER_PROMPT.md` exactly," plus the "Never" list. - `.claude/agents/-reviewer.md`: **read-only tools only** (`Read`, `Grep`, `Glob`, and read-only `git diff`). The body is "Follow `REVIEWER_PROMPT.md`; you have no memory of the implementation." It must be run as a separate invocation, or as a subagent with its own context, never in the worker's conversation. - `.claude/agents/-initializer.md`: read-only, plus running the baseline checks. It follows `INIT_PROMPT.md`. -- `.claude/skills/-loop/SKILL.md`: a short guide for a **supervising** session that explains how to step through a slice. The steps are: invoke the worker, tell the human to run the verifier, invoke the reviewer, present the findings for the human to accept or reject, and record `decision.md`. It must state that the supervising session cannot mark a slice as ready on its own, and that it never runs the verifier in place of the human's trusted copy unless the human has configured that. +- `.claude/skills/-loop/SKILL.md`: a short guide for a **supervising** session that explains how to step through a slice. The steps are: invoke the worker, run the verifier (the human's trusted copy, when the human configured that; otherwise tell the human to run it), invoke the reviewer, apply the auto-advance rule to the verifier and reviewer outputs, and record `decision.md`. It must state that the supervisor advances only on verifier PASS plus no blocking reviewer findings plus low risk, never on a model-written claim of success, and that it halts only for the BLOCKED conditions. Do **not** add hooks that block or auto-approve. Optionally, propose (without installing) a `PostToolUse` hook that runs cheap, scoped lint on edited files, and note that hooks give feedback only and are not enforcement. diff --git a/.github/scripts/build_dashboard.py b/.github/scripts/build_dashboard.py index decaa3e2..3b61975c 100644 --- a/.github/scripts/build_dashboard.py +++ b/.github/scripts/build_dashboard.py @@ -72,7 +72,7 @@ def parse_junit(xml_path: Path) -> dict: "skipped": skipped, "available": True, } - except Exception as exc: + except (ET.ParseError, OSError, ValueError) as exc: return { "tests": 0, "passed": 0, @@ -118,7 +118,7 @@ def parse_coverage(json_path: Path) -> dict: if pct is not None: return {"percent": float(pct), "available": True} return {"percent": None, "available": False} - except Exception as exc: + except (OSError, ValueError, TypeError) as exc: return {"percent": None, "available": False, "parse_error": str(exc)} @@ -168,9 +168,9 @@ def build_html(data: dict) -> str: ts = data["timestamp"] try: - dt = datetime.fromisoformat(ts.replace("Z", "+00:00")) + dt = datetime.fromisoformat(ts) ts_display = dt.strftime("%Y-%m-%d %H:%M UTC") - except Exception: + except ValueError: ts_display = ts # Determine overall badge @@ -313,9 +313,9 @@ def build_summary_md(data: dict) -> str: ts = data["timestamp"] try: - dt = datetime.fromisoformat(ts.replace("Z", "+00:00")) + dt = datetime.fromisoformat(ts) ts_display = dt.strftime("%Y-%m-%d %H:%M UTC") - except Exception: + except ValueError: ts_display = ts if not junit.get("available"): diff --git a/.gitignore b/.gitignore index b9cf66f6..875773cf 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ build __pycache__ _version.py htmlcov +.vscode .vscode/* !.vscode/settings.json diff --git a/.ruff.toml b/.ruff.toml index 169ad232..3b8d8787 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -4,3 +4,7 @@ [lint] # Ignore UP007/UP045: Use Optional[X] instead of X | None for itk.Image compatibility ignore = ["UP007", "UP045"] + +[lint.per-file-ignores] +# Simpleware loads this script by file name, so the module name is fixed. +"src/monai_physio/simpleware_medical/SimplewareScript_heart_segmentation.py" = ["N999"] diff --git a/data/TCIA-4DLung/convert.py b/data/TCIA-4DLung/convert.py index c6389df3..531ac365 100644 --- a/data/TCIA-4DLung/convert.py +++ b/data/TCIA-4DLung/convert.py @@ -48,6 +48,7 @@ def convert_series(series_dir, root, patient_name, pct_counts): [DCM2NIIX, "-f", "%d", "-b", "n", "-z", "y", "-o", tmp_dir, series_dir], capture_output=True, text=True, + check=False, ) if result.returncode != 0: print(f" FAILED: {series_dir}\n{result.stderr}") diff --git a/docs/conf.py b/docs/conf.py index be03104c..63b0e953 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -6,7 +6,7 @@ import sys import tomllib import warnings -from datetime import datetime +from datetime import UTC, datetime from pathlib import Path from unittest.mock import MagicMock @@ -30,7 +30,7 @@ def __getattr__(cls, name): # -- Project information ----------------------------------------------------- project = "MONAI Physio" -copyright = f"{datetime.now().year}, Stephen R. Aylward, NVIDIA Corporation" +copyright = f"{datetime.now(tz=UTC).year}, Stephen R. Aylward, NVIDIA Corporation" author = "Stephen R. Aylward" _repo_root = Path(__file__).resolve().parents[1] _pyproject = tomllib.loads((_repo_root / "pyproject.toml").read_text(encoding="utf-8")) diff --git a/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py b/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py index 9fcdf039..e925f1ea 100644 --- a/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py +++ b/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py @@ -46,7 +46,7 @@ print(f" Original mesh: {mesh.n_cells} cells, {mesh.n_points} points") print(f" Surface mesh: {surface.n_cells} cells, {surface.n_points} points") - except Exception as e: + except Exception as e: # noqa: BLE001 - report and go on to the next mesh print(f" Error processing {vtk_file}: {e!s}") print("\n" + "=" * 50) diff --git a/experiments/Heart-GatedCT-OptimizedLongitudinalRegistration/1-initial_registration.py b/experiments/Heart-GatedCT-OptimizedLongitudinalRegistration/1-initial_registration.py index 11755604..c3decd38 100644 --- a/experiments/Heart-GatedCT-OptimizedLongitudinalRegistration/1-initial_registration.py +++ b/experiments/Heart-GatedCT-OptimizedLongitudinalRegistration/1-initial_registration.py @@ -432,9 +432,13 @@ def crop_image_to_mask( output_dir = _HERE / f"results{cond}" detail_landmarks_file = output_dir / "registration_landmarks_init.csv" detail_dice_file = output_dir / "registration_dice_init.csv" - if subject_index == 0 and image_index == 0 and cond_index == 0: - if output_dir.exists(): - shutil.rmtree(output_dir) + if ( + subject_index == 0 + and image_index == 0 + and cond_index == 0 + and output_dir.exists() + ): + shutil.rmtree(output_dir) output_dir.mkdir(parents=True, exist_ok=True) for method_index, method_name in enumerate(methods): diff --git a/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py b/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py index a155fbb4..0cd6cb6a 100644 --- a/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py +++ b/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py @@ -1,5 +1,6 @@ # %% import os +import sys import itk import numpy as np @@ -36,7 +37,7 @@ ) except (FileNotFoundError, OSError): print("No image found. Aborting. Please run 1-generate_images.ipynb first.") - exit(1) + sys.exit(1) arr = itk.array_from_image(image) print(arr.shape) arr = np.where(arr == 0, -1000, arr) diff --git a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py index 8800775e..22573031 100644 --- a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py +++ b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py @@ -225,7 +225,7 @@ print("Please check the error messages above for details.") result = None - except Exception as e: + except Exception as e: # noqa: BLE001 - last-resort handler for the script print(f"\nUnexpected error: {e}") result = None diff --git a/experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py b/experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py index e21b66a3..77a74662 100644 --- a/experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py +++ b/experiments/Heart_and_Lungs_Motion/0-heart_and_lungs_beating_heart.py @@ -174,7 +174,7 @@ def _ensure_mgn_inference_assets( reference_image=reference_image, fitted_reference_mesh=registered_mesh_file, ) - pct = int(round(stage * 100)) + pct = round(stage * 100) field_path = output_dir / f"deformation_field_s{pct:03d}.mha" normal_path = output_dir / f"surface_normal_field_s{pct:03d}.mha" surface_path = output_dir / f"deformed_heart_surface_s{pct:03d}.vtp" @@ -185,7 +185,7 @@ def _ensure_mgn_inference_assets( normal_files.append(normal_path) surface_files.append(surface_path) deformed_surfaces.append(result["deformed_surface"]) - logging.info("stage %.2f -> %s", stage, field_path.name) + logging.getLogger(__name__).info("stage %.2f -> %s", stage, field_path.name) # Assemble the ordered per-stage surfaces into one animated 4D USD. The # stages loop over one RR interval, so play them back at one beat per second. diff --git a/experiments/Lung-VesselsAirways/0-GenData.py b/experiments/Lung-VesselsAirways/0-GenData.py index da2aa774..6271c911 100644 --- a/experiments/Lung-VesselsAirways/0-GenData.py +++ b/experiments/Lung-VesselsAirways/0-GenData.py @@ -1,5 +1,5 @@ # %% -from datetime import datetime +from datetime import UTC, datetime import itk import matplotlib.pyplot as plt @@ -296,7 +296,7 @@ def __getitem__(self, _): optimizer.step() running_loss += loss.item() print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(dataloader)}") - date = datetime.now().strftime("%Y.%m.%d") + date = datetime.now(tz=UTC).strftime("%Y.%m.%d") model_path = f"vessel_seg_resnet18_{date}.pth" torch.save(model, model_path) diff --git a/experiments/Registration-Fireants_vs_Greedy/benchmark.py b/experiments/Registration-Fireants_vs_Greedy/benchmark.py index 72d7e4b6..612e4634 100644 --- a/experiments/Registration-Fireants_vs_Greedy/benchmark.py +++ b/experiments/Registration-Fireants_vs_Greedy/benchmark.py @@ -87,9 +87,14 @@ def load_dataset(dataset: str, smoke: bool) -> tuple[list[Path], int]: else: data_dir = TCIA_4D_LUNG.input_directory(False) files = sorted(data_dir.glob("100_HM10395_g0??.nii.gz")) - reference = next(i for i, p in enumerate(files) if "g070" in p.name) + reference = next((i for i, p in enumerate(files) if "g070" in p.name), -1) if not files: raise FileNotFoundError(f"No {dataset} phase images found under {data_dir}") + if reference < 0: + raise FileNotFoundError( + f"Lung reference phase g070 not found among {len(files)} phase images " + f"under {data_dir}" + ) return files, reference diff --git a/src/monai_physio/cli/convert_image_to_usd.py b/src/monai_physio/cli/convert_image_to_usd.py index 4a3c691b..6c7ff3da 100644 --- a/src/monai_physio/cli/convert_image_to_usd.py +++ b/src/monai_physio/cli/convert_image_to_usd.py @@ -160,7 +160,7 @@ def main() -> int: registration_method=registration_method, frames_per_second=args.frames_per_second, ) - except Exception as e: + except Exception as e: # noqa: BLE001 - CLI boundary: report and exit nonzero print(f"Error initializing workflow: {e}") return 1 @@ -180,7 +180,7 @@ def main() -> int: return 0 - except Exception as e: + except Exception as e: # noqa: BLE001 - CLI boundary: report and exit nonzero print(f"\nError during processing: {e}") import traceback diff --git a/src/monai_physio/cli/convert_vtk_to_usd.py b/src/monai_physio/cli/convert_vtk_to_usd.py index 1cee1200..ec99f169 100644 --- a/src/monai_physio/cli/convert_vtk_to_usd.py +++ b/src/monai_physio/cli/convert_vtk_to_usd.py @@ -183,8 +183,8 @@ def main() -> int: "Color values must all be in [0, 1] or all in [0, 255]." ) else: - raise ValueError("Color must be specified as a list of 3 float values") - except ValueError as e: + raise TypeError("Color must be specified as a list of 3 float values") + except (TypeError, ValueError) as e: print(f"Error: {e}") return 1 @@ -237,7 +237,7 @@ def main() -> int: print("\nConversion completed successfully.") print(f"Output: {result['usd_file']}") return 0 - except Exception as e: + except Exception as e: # noqa: BLE001 - CLI boundary: report and exit nonzero print(f"\nError during conversion: {e}") import traceback diff --git a/src/monai_physio/convert_vtk_to_usd.py b/src/monai_physio/convert_vtk_to_usd.py index 4b524078..7363e6a9 100644 --- a/src/monai_physio/convert_vtk_to_usd.py +++ b/src/monai_physio/convert_vtk_to_usd.py @@ -320,7 +320,7 @@ def from_files( report = validate_time_series_topology(mesh_data_seq) for w in report.get("warnings", []): instance.log_warning("%s", w) - except Exception as exc: + except Exception as exc: # noqa: BLE001 - validation is advisory only instance.log_debug("Topology validation skipped: %s", exc) return instance @@ -443,7 +443,7 @@ def list_available_arrays(self) -> dict: # Get point data arrays if hasattr(mesh, "point_data"): - for array_name in mesh.point_data.keys(): + for array_name in mesh.point_data: if array_name not in available_arrays: array_data = mesh.point_data[array_name] available_arrays[array_name] = { @@ -942,7 +942,7 @@ def _vtk_to_mesh_data( (vtk_mesh.point_data, "vertex"), (vtk_mesh.cell_data, "uniform"), ): - for array_name in source.keys(): + for array_name in source: array_data = source[array_name] num_components = int(array_data.shape[1] if array_data.ndim > 1 else 1) diff --git a/src/monai_physio/download_data.py b/src/monai_physio/download_data.py index af077de6..9f857036 100644 --- a/src/monai_physio/download_data.py +++ b/src/monai_physio/download_data.py @@ -16,7 +16,7 @@ import urllib.request import zipfile from pathlib import Path -from typing import Optional, Union +from typing import ClassVar, Optional, Union import itk import numpy as np @@ -91,7 +91,7 @@ def _DownloadFile(url: str, target_file: Path) -> None: Avoids partial files on interrupt and the indefinite hang that urlretrieve has without a timeout. """ - tmp_handle = tempfile.NamedTemporaryFile( + tmp_handle = tempfile.NamedTemporaryFile( # noqa: SIM115 - closed by the with below dir=str(target_file.parent), prefix=f".{target_file.name}.", suffix=".tmp", @@ -216,7 +216,7 @@ def _DownloadAndExtractTarMember( # subdirectory name -> release asset filename. Alterra and TPV25 are # each >1 GB (per-frame valve mesh time series); CT is the smaller # image/segmentation bundle. See data/CHOP-Valve4D/README.md. - CHOP_VALVE4D_ASSETS = { + CHOP_VALVE4D_ASSETS: ClassVar[dict[str, str]] = { "Alterra": "CHOP-Valve4D-Alterra.zip", "CT": "CHOP-Valve4D-CT.zip", "TPV25": "CHOP-Valve4D-TPV25.zip", @@ -360,7 +360,10 @@ def VerifyChestCTData(dirname: Union[str, Path]) -> bool: TCIA_4D_LUNG_RELEASE_URL = ( "https://github.com/Project-MONAI/monai-physio/releases/download/2026.07.1/" ) - TCIA_4D_LUNG_ASSETS = ["TCIA-4DLung-Part1.zip", "TCIA-4DLung-Part2.zip"] + TCIA_4D_LUNG_ASSETS: ClassVar[list[str]] = [ + "TCIA-4DLung-Part1.zip", + "TCIA-4DLung-Part2.zip", + ] @staticmethod def DownloadTCIA4DLungData(dirname: Union[str, Path]) -> Path: diff --git a/src/monai_physio/process_images.py b/src/monai_physio/process_images.py index 3f43507a..c9bb0eae 100644 --- a/src/monai_physio/process_images.py +++ b/src/monai_physio/process_images.py @@ -494,7 +494,7 @@ def pad_image( constant = ( float(background_value) if pixel_type in (itk.F, itk.D) - else int(round(background_value)) + else round(background_value) ) pad_filter = itk.ConstantPadImageFilter[ImageType, ImageType].New() diff --git a/src/monai_physio/process_labelmaps.py b/src/monai_physio/process_labelmaps.py index 19ed07fd..be956ffa 100644 --- a/src/monai_physio/process_labelmaps.py +++ b/src/monai_physio/process_labelmaps.py @@ -88,7 +88,7 @@ def convert_labelmap_to_mask( spacing = labelmap.GetSpacing() radius = itk.Size[3]() for i in range(3): - radius[i] = max(1, int(round(dilation_in_mm / float(spacing[i])))) + radius[i] = max(1, round(dilation_in_mm / float(spacing[i]))) structuring_element = itk.FlatStructuringElement[3].Ball(radius) return itk.binary_dilate_image_filter( mask, kernel=structuring_element, foreground_value=1 diff --git a/src/monai_physio/process_tests.py b/src/monai_physio/process_tests.py index 448a04a3..42b98b72 100644 --- a/src/monai_physio/process_tests.py +++ b/src/monai_physio/process_tests.py @@ -8,6 +8,7 @@ from __future__ import annotations +import contextlib import csv import json import logging @@ -81,7 +82,7 @@ def _read_metrics(path: Path) -> dict[str, Any]: rows = list(csv.DictReader(handle)) if not rows: return {} - fieldnames = [name for name in (rows[0].keys()) if name is not None] + fieldnames = [name for name in (rows[0]) if name is not None] label_column = fieldnames[0] labels = [row.get(label_column, "") or "" for row in rows] use_labels = len(set(labels)) == len(labels) @@ -594,11 +595,9 @@ def save_screenshot_mesh( import pyvista as pv xvfb_started = False - try: + with contextlib.suppress(Exception): pv.start_xvfb() xvfb_started = True - except Exception: - pass output_path = self._results_dir / filename plotter = pv.Plotter(off_screen=True, window_size=list(window_size)) diff --git a/src/monai_physio/process_transforms.py b/src/monai_physio/process_transforms.py index 111cbee1..72b24b7a 100644 --- a/src/monai_physio/process_transforms.py +++ b/src/monai_physio/process_transforms.py @@ -363,7 +363,7 @@ def invert_transform( """ try: analytic = tfm.GetInverseTransform() - except Exception: # pragma: no cover - transform type dependent + except Exception: # noqa: BLE001 # pragma: no cover - transform type dependent analytic = None if analytic is not None: return cast(itk.Transform, analytic) @@ -555,7 +555,7 @@ def transform_image( dtype = itk.GetArrayViewFromImage(image).dtype default_pixel_value: Union[int, float] if np.issubdtype(dtype, np.integer) or np.issubdtype(dtype, np.bool_): - default_pixel_value = int(round(background_value)) + default_pixel_value = round(background_value) low, high = ( (0, 1) if np.issubdtype(dtype, np.bool_) diff --git a/src/monai_physio/process_usd.py b/src/monai_physio/process_usd.py index d2c7eec0..b35cb2af 100644 --- a/src/monai_physio/process_usd.py +++ b/src/monai_physio/process_usd.py @@ -12,6 +12,7 @@ from __future__ import annotations +import contextlib import logging from collections.abc import Sequence from pathlib import Path @@ -502,7 +503,7 @@ def _copy_prim(src_prim: Any, target_path: str) -> None: if value is not None: try: new_attr.Set(value) - except Exception as e: + except Exception as e: # noqa: BLE001 - USD raises assorted errors self.log_warning( f"Failed to copy attribute {attr.GetName()}: {e}" ) @@ -818,7 +819,7 @@ def list_mesh_primvars( ) except (TypeError, ValueError): pass # Skip range for non-numeric data - except Exception as e: + except Exception as e: # noqa: BLE001 - USD raises assorted errors self.log_debug( f"Could not get value for primvar {pv_info['name']}: {e}" ) @@ -1120,14 +1121,12 @@ def apply_colormap_from_primvar( ) # If we're rewriting displayColor, clear any previously-authored time samples first. # This prevents leaving behind stale/corrupt samples at times we no longer author. - try: + # Silently ignore errors (e.g., if attribute doesn't exist yet or has no samples). + # This is expected on first-time creation or when no time samples are present. + with contextlib.suppress(Exception): dc_attr = display_color_pv.GetAttr() for t in list(dc_attr.GetTimeSamples()): dc_attr.ClearAtTime(t) - except Exception: - # Silently ignore errors (e.g., if attribute doesn't exist yet or has no samples). - # This is expected on first-time creation or when no time samples are present. - pass for idx, (tc, scalar) in enumerate(scalar_samples): # Normalize to [0, 1] @@ -1407,7 +1406,7 @@ def repair_mesh_primvar_element_sizes( "new_elementSize": int(inferred), } ) - except Exception as e: + except Exception as e: # noqa: BLE001 - USD raises assorted errors skipped.append( {"name": pv.GetName(), "reason": f"SetElementSize failed: {e}"} ) diff --git a/src/monai_physio/register_images_ants.py b/src/monai_physio/register_images_ants.py index 471d746d..25794487 100644 --- a/src/monai_physio/register_images_ants.py +++ b/src/monai_physio/register_images_ants.py @@ -457,8 +457,8 @@ def itk_transform_to_ANTSfile( ... fixed=fixed_ANTS, moving=moving_ANTS, initial_transform=transform_files ... ) """ - if isinstance(itk_tfm, itk.DisplacementFieldTransform) or isinstance( - itk_tfm, itk.CompositeTransform + if isinstance( + itk_tfm, (itk.DisplacementFieldTransform, itk.CompositeTransform) ): transform_tools = ProcessTransforms() disp_field_itk = transform_tools.convert_transform_to_displacement_field( diff --git a/src/monai_physio/register_images_fireants.py b/src/monai_physio/register_images_fireants.py index af05e738..de1b0e32 100644 --- a/src/monai_physio/register_images_fireants.py +++ b/src/monai_physio/register_images_fireants.py @@ -210,6 +210,10 @@ def configure_for_lung_ct(self) -> None: def _scales(self) -> list[float]: """Downsampling factor per level, coarse to fine.""" if self.scales is not None: + if len(self.scales) != len(self.number_of_iterations): + raise ValueError( + "scales and number_of_iterations must have equal length" + ) return list(self.scales) n = len(self.number_of_iterations) return [float(2 ** (n - 1 - i)) for i in range(n)] @@ -220,8 +224,12 @@ def _to_fireants_image( """Wrap an ITK image as a batched FireANTs image on the GPU. When ``masked`` is True the mask is appended as the last channel, as the - FireANTs masked losses expect; a missing mask becomes all ones. + FireANTs masked losses expect; a missing mask becomes all ones. For the + Mattes metric the intensities are first rescaled to [0, 1], the range + the FireANTs mutual-information loss clamps to. """ + if self.metric == "Mattes": + image = self._rescale_unit_range(image) fireants_image = _try_import_fireants("io.image") sitk_image = ProcessImages().convert_itk_image_to_sitk(image) wrapped = fireants_image.Image(sitk_image, device="cuda") @@ -236,6 +244,19 @@ def _to_fireants_image( ) return fireants_image.BatchedImages([wrapped]) + @staticmethod + def _rescale_unit_range(image: itk.Image) -> itk.Image: + """Return a float32 copy of ``image`` linearly rescaled to [0, 1].""" + array = itk.array_from_image(image).astype(np.float32) + low, high = float(array.min()), float(array.max()) + if high > low: + array = (array - low) / (high - low) + else: + array = np.zeros_like(array) + rescaled = itk.image_from_array(array) + rescaled.CopyInformation(image) + return rescaled + @staticmethod def _matrix_to_itk_affine(mat_4x4: NDArray[np.float64]) -> itk.Transform: """Convert a 4x4 physical-space ``y = A x + t`` matrix to an ITK affine.""" diff --git a/src/monai_physio/register_images_greedy.py b/src/monai_physio/register_images_greedy.py index 1230e40d..a77460dc 100644 --- a/src/monai_physio/register_images_greedy.py +++ b/src/monai_physio/register_images_greedy.py @@ -199,7 +199,7 @@ def _downsample_image( return image size = image.GetLargestPossibleRegion().GetSize() spacing = image.GetSpacing() - new_size = [max(1, int(round(int(size[i]) * scale))) for i in range(3)] + new_size = [max(1, round(int(size[i]) * scale)) for i in range(3)] new_spacing = [float(spacing[i]) * int(size[i]) / new_size[i] for i in range(3)] kwargs: dict[str, Any] = { "output_origin": image.GetOrigin(), @@ -337,7 +337,7 @@ def _registration_method_affine_or_rigid( try: ml = g.metric_log() loss = float(ml[-1]["TotalPerPixelMetric"][-1]) if ml else 0.0 - except Exception: + except (KeyError, IndexError, TypeError, ValueError, RuntimeError): loss = 0.0 self.log_info("Greedy affine/rigid registration loss: %s", loss) return mat, loss @@ -416,7 +416,7 @@ def _registration_method_deformable( try: ml = g.metric_log() loss = float(ml[-1]["TotalPerPixelMetric"][-1]) if ml else 0.0 - except Exception: + except (KeyError, IndexError, TypeError, ValueError, RuntimeError): loss = 0.0 self.log_info("Greedy deformable registration loss: %s", loss) return initial_affine, warp_out, loss diff --git a/src/monai_physio/register_images_icon.py b/src/monai_physio/register_images_icon.py index f81d1df4..1b32be8f 100644 --- a/src/monai_physio/register_images_icon.py +++ b/src/monai_physio/register_images_icon.py @@ -12,7 +12,7 @@ import logging from pathlib import Path -from typing import Any, Optional, Union +from typing import Any, ClassVar, Optional, Union import itk import numpy as np @@ -89,7 +89,7 @@ class RegisterImagesICON(RegisterImagesBase): # of times in a cohort run, and the freed host arenas are not necessarily # returned to the OS -- which is how a long run walks into the Linux OOM # killer. - _net_cache: dict[tuple[Optional[str], bool, bool], Any] = {} + _net_cache: ClassVar[dict[tuple[Optional[str], bool, bool], Any]] = {} def __init__(self, log_level: int | str = logging.INFO) -> None: """Initialize the ICON image registration class. diff --git a/src/monai_physio/register_models_distance_maps.py b/src/monai_physio/register_models_distance_maps.py index 3c04dd54..0bac5758 100644 --- a/src/monai_physio/register_models_distance_maps.py +++ b/src/monai_physio/register_models_distance_maps.py @@ -227,8 +227,10 @@ def _require_non_constant( return message = [ - f"The {description} is constant ({minimum:.6g} everywhere), so there " - "is nothing for a registration metric to align." + ( + f"The {description} is constant ({minimum:.6g} everywhere), so there " + "is nothing for a registration metric to align." + ) ] if greedy_loss is not None: message.append( diff --git a/src/monai_physio/register_models_icp.py b/src/monai_physio/register_models_icp.py index 74506f5d..3c756c1b 100644 --- a/src/monai_physio/register_models_icp.py +++ b/src/monai_physio/register_models_icp.py @@ -40,7 +40,7 @@ """ import logging -from typing import Optional +from typing import ClassVar, Optional import itk import numpy as np @@ -140,7 +140,7 @@ def __init__( self.registered_model: Optional[pv.PolyData] = None # ICP stages run for each transform type, in order. - _ICP_STAGES = { + _ICP_STAGES: ClassVar[dict[str, tuple[str, ...]]] = { "Rigid": ("Rigid",), "Similarity": ("Similarity",), "Affine": ("Similarity", "Affine"), diff --git a/src/monai_physio/register_time_series_images.py b/src/monai_physio/register_time_series_images.py index fb27f39c..9975b6bd 100644 --- a/src/monai_physio/register_time_series_images.py +++ b/src/monai_physio/register_time_series_images.py @@ -589,8 +589,8 @@ def _create_upsampled_reference( # Calculate new size to cover the same physical extent with isotropic spacing new_size = itk.Size[3]() for i in range(3): - new_size[i] = int( - round((moving_size[i] * moving_spacing[i]) / isotropic_spacing[i]) + new_size[i] = round( + (moving_size[i] * moving_spacing[i]) / isotropic_spacing[i] ) # Create reference image with combined properties diff --git a/src/monai_physio/report_evaluate_movement.py b/src/monai_physio/report_evaluate_movement.py index e7fcf178..8a710341 100644 --- a/src/monai_physio/report_evaluate_movement.py +++ b/src/monai_physio/report_evaluate_movement.py @@ -213,8 +213,10 @@ def _write_report( "## Run", "", f"- Hold-out case: `{provenance['case_id']}`", - f"- Stages evaluated: {len(stages)} " - f"({', '.join(f'{stage:.2f}' for stage in stages)})", + ( + f"- Stages evaluated: {len(stages)} " + f"({', '.join(f'{stage:.2f}' for stage in stages)})" + ), f"- Shape parameters: `{provenance['shape_parameters_file']}`", "- Shape parameters (standard deviations): " + json.dumps([round(value, 4) for value in coefficients]), @@ -225,20 +227,26 @@ def _write_report( f"- Deformation smoothing sigma: {smoothing_sigma_mm:.1f} mm", f"- Evaluation grid pitch: {evaluation_spacing_mm:.2f} mm isotropic", "", - "Every score compares the reference frame carried into that stage " - "with the inferred deformation against the frame acquired there.", + ( + "Every score compares the reference frame carried into that stage " + "with the inferred deformation against the frame acquired there." + ), "", "## Volume over the stages", "", f"![Structure volume against stage]({plot_file.name})", "", - "Solid: the volume acquired at that stage. Dashed: the volume the " - "prediction carries there.", + ( + "Solid: the volume acquired at that stage. Dashed: the volume the " + "prediction carries there." + ), "", "## Per structure, over the stages", "", - "Mean beside worst case: the stage a structure is predicted worst at " - "is what a plausible mean can hide.", + ( + "Mean beside worst case: the stage a structure is predicted worst at " + "is what a plausible mean can hide." + ), "", ] metrics = (["Dice"] if has_dice else []) + [ @@ -358,20 +366,24 @@ def _displacement_error_section( "", "## Displacement error", "", - "Distance between where the network puts each mesh point and where " - "the shape model fitted it in the frame acquired at that stage. " - "Unlike the labelmap metrics above, this is measured point by point, " - "so it does not average a displacement away against its opposite.", + ( + "Distance between where the network puts each mesh point and where " + "the shape model fitted it in the frame acquired at that stage. " + "Unlike the labelmap metrics above, this is measured point by point, " + "so it does not average a displacement away against its opposite." + ), "", f"- RMS over every point and stage: {pooled['rms_mm']:.3f} mm", f"- 95th percentile over every point and stage: {pooled['p95_mm']:.3f} mm", f"- Max over every point and stage: {pooled['max_mm']:.3f} mm", "", - "The 95th percentile is the figure to quote: the maximum is one " - "point of one stage, so it moves with a single badly placed vertex, " - "while the RMS is pulled down by the bulk of the surface that " - "hardly moves at all. Every figure here is over the whole shape " - "model; the same error split by structure is in the table above.", + ( + "The 95th percentile is the figure to quote: the maximum is one " + "point of one stage, so it moves with a single badly placed vertex, " + "while the RMS is pulled down by the bulk of the surface that " + "hardly moves at all. Every figure here is over the whole shape " + "model; the same error split by structure is in the table above." + ), "", ] lines += self._table_header( diff --git a/src/monai_physio/segment_heart_simpleware.py b/src/monai_physio/segment_heart_simpleware.py index e5a042b4..3fc2d989 100644 --- a/src/monai_physio/segment_heart_simpleware.py +++ b/src/monai_physio/segment_heart_simpleware.py @@ -227,6 +227,7 @@ def segmentation_method(self, preprocessed_image: itk.image) -> itk.image: subprocess.run( ["taskkill", "/F", "/T", "/PID", str(proc.pid)], capture_output=True, + check=False, ) else: os.killpg(os.getpgid(proc.pid), 9) diff --git a/src/monai_physio/train_physicsnemo_base.py b/src/monai_physio/train_physicsnemo_base.py index 21b709f9..0618bf0e 100644 --- a/src/monai_physio/train_physicsnemo_base.py +++ b/src/monai_physio/train_physicsnemo_base.py @@ -299,7 +299,7 @@ def train( # shape instead of doing that arithmetic symbolically. model = cast("torch.nn.Module", torch.compile(model, dynamic=False)) self._log_main(context, "torch.compile enabled.") - except Exception as exc: # pragma: no cover - platform dependent + except Exception as exc: # noqa: BLE001 # pragma: no cover - platform dependent self._log_main(context, "torch.compile skipped (%s).", exc) else: self._log_main(context, "torch.compile skipped on Windows.") diff --git a/src/monai_physio/vtk_to_usd/__init__.py b/src/monai_physio/vtk_to_usd/__init__.py index 5547145e..39858432 100644 --- a/src/monai_physio/vtk_to_usd/__init__.py +++ b/src/monai_physio/vtk_to_usd/__init__.py @@ -49,38 +49,38 @@ ) __all__ = [ - # File facade - "convert_vtk_file", # Data structures "ConversionSettings", "DataType", "GenericArray", + "LegacyVTKReader", "MaterialData", - "MeshData", - "TimeStepData", - "VolumeData", # Managers "MaterialManager", + "MeshData", + "PolyDataReader", + "TimeStepData", + "UnstructuredGridReader", "UsdMeshConverter", - # Utilities - "lps_to_usd", - "lps_points_to_usd", - "lps_normals_to_usd", - "create_primvar", - "sanitize_primvar_name", - "triangulate_face", - "compute_mesh_extent", + # Readers + "VTKReader", + "VolumeData", "add_framing_camera", # Mesh utils (cell type split) "cell_type_name_for_vertex_count", + "compute_mesh_extent", + # File facade + "convert_vtk_file", + "create_primvar", + "lps_normals_to_usd", + "lps_points_to_usd", + # Utilities + "lps_to_usd", + "read_vtk_file", + "sanitize_primvar_name", "split_mesh_data_by_cell_type", "split_mesh_data_by_connectivity", - # Readers - "VTKReader", - "PolyDataReader", - "LegacyVTKReader", - "UnstructuredGridReader", - "read_vtk_file", + "triangulate_face", "validate_time_series_topology", ] diff --git a/src/monai_physio/vtk_to_usd/primvar_derivations.py b/src/monai_physio/vtk_to_usd/primvar_derivations.py index 35a67c86..1f299f31 100644 --- a/src/monai_physio/vtk_to_usd/primvar_derivations.py +++ b/src/monai_physio/vtk_to_usd/primvar_derivations.py @@ -114,7 +114,7 @@ def derive_von_mises_from_stress(array: GenericArray) -> list[GenericArray]: return [] try: vm = compute_von_mises_stress(array.data) - except Exception as exc: + except Exception as exc: # noqa: BLE001 - derivation is optional logger.debug("Skipping VonMises derivation for '%s': %s", array.name, exc) return [] return [ diff --git a/src/monai_physio/vtk_to_usd/usd_mesh_converter.py b/src/monai_physio/vtk_to_usd/usd_mesh_converter.py index 4c3f1ecd..fa6ff3f4 100644 --- a/src/monai_physio/vtk_to_usd/usd_mesh_converter.py +++ b/src/monai_physio/vtk_to_usd/usd_mesh_converter.py @@ -168,10 +168,13 @@ def create_mesh( self._add_generic_arrays(mesh, mesh_data, time_code, triangulation_face_map) # Bind material (if material_id is provided and material exists in cache) - if bind_material and mesh_data.material_id: - if mesh_data.material_id in self.material_mgr.material_cache: - material = self.material_mgr.material_cache[mesh_data.material_id] - self.material_mgr.bind_material(mesh, material) + if ( + bind_material + and mesh_data.material_id + and mesh_data.material_id in self.material_mgr.material_cache + ): + material = self.material_mgr.material_cache[mesh_data.material_id] + self.material_mgr.bind_material(mesh, material) logger.info( f"Created mesh with {len(mesh_data.points)} points, " @@ -307,7 +310,7 @@ def _add_generic_arrays( ) create_primvar(mesh, chunk_arr, prefix, time_code) - except Exception as e: + except Exception as e: # noqa: BLE001 - USD raises assorted errors logger.warning("Failed to split primvar %s: %s", array.name, e) continue @@ -333,7 +336,7 @@ def _add_generic_arrays( try: create_primvar(mesh, array, prefix, time_code) - except Exception as e: + except Exception as e: # noqa: BLE001 - USD raises assorted errors logger.warning(f"Failed to create primvar for {array.name}: {e}") def create_time_varying_mesh( diff --git a/src/monai_physio/vtk_to_usd/usd_utils.py b/src/monai_physio/vtk_to_usd/usd_utils.py index efa06cdc..d8abe7b7 100644 --- a/src/monai_physio/vtk_to_usd/usd_utils.py +++ b/src/monai_physio/vtk_to_usd/usd_utils.py @@ -5,6 +5,7 @@ from __future__ import annotations +import contextlib import logging from typing import Any @@ -339,11 +340,9 @@ def create_primvar( Sdf.ValueTypeNames.UIntArray, Sdf.ValueTypeNames.UCharArray, ): - try: + # Not fatal; continue without elementSize. + with contextlib.suppress(Exception): primvar.SetElementSize(int(array.num_components)) - except Exception: - # Not fatal; continue without elementSize. - pass # Convert data to VtArray vt_array = numpy_to_vt_array(array.data, array.data_type) diff --git a/src/monai_physio/vtk_to_usd/vtk_reader.py b/src/monai_physio/vtk_to_usd/vtk_reader.py index 21526ffb..92f046e5 100644 --- a/src/monai_physio/vtk_to_usd/vtk_reader.py +++ b/src/monai_physio/vtk_to_usd/vtk_reader.py @@ -668,12 +668,11 @@ def validate_time_series_topology( f"got {array_size}, expected {n_points} points" ) - elif array.interpolation == "uniform": - if array_size != n_faces: - warnings.append( - f"{frame_label}: Cell data '{array.name}' size mismatch - " - f"got {array_size}, expected {n_faces} faces" - ) + elif array.interpolation == "uniform" and array_size != n_faces: + warnings.append( + f"{frame_label}: Cell data '{array.name}' size mismatch - " + f"got {array_size}, expected {n_faces} faces" + ) # Log summary if topology_changes: diff --git a/tests/conftest.py b/tests/conftest.py index d564856d..d332fd67 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,7 +8,7 @@ import logging import os import shutil -from datetime import datetime, timedelta +from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any, Optional @@ -170,7 +170,7 @@ def pytest_configure(config: pytest.Config) -> None: config._test_timings = { # type: ignore[attr-defined] "tests": [], "total_time": 0.0, - "start_time": datetime.now(), + "start_time": datetime.now(tz=UTC), } @@ -297,7 +297,7 @@ def pytest_terminal_summary( return # Calculate session duration - session_duration = datetime.now() - timings["start_time"] + session_duration = datetime.now(tz=UTC) - timings["start_time"] # Separate regular and tutorial tests regular_tests = [t for t in tests if not t["is_tutorial"]] @@ -497,7 +497,7 @@ def test_images( input_spacing = list(img.GetSpacing()) input_size = list(itk.size(img)) output_size = [ - int(round(input_size[i] * input_spacing[i] / target_spacing[i])) + round(input_size[i] * input_spacing[i] / target_spacing[i]) for i in range(3) ] interpolator = itk.LinearInterpolateImageFunction.New(img) @@ -539,7 +539,7 @@ def _downsample_image(source: Path, destination: Path, spacing_mm: float) -> Non input_spacing = list(image.GetSpacing()) input_size = list(itk.size(image)) output_size = [ - max(1, int(round(input_size[i] * input_spacing[i] / target_spacing[i]))) + max(1, round(input_size[i] * input_spacing[i] / target_spacing[i])) for i in range(3) ] resampler = itk.ResampleImageFilter.New(Input=image) @@ -560,7 +560,7 @@ def _downsample_labelmap(source: Path, destination: Path, spacing_mm: float) -> input_spacing = list(image.GetSpacing()) input_size = list(itk.size(image)) output_size = [ - max(1, int(round(input_size[i] * input_spacing[i] / target_spacing[i]))) + max(1, round(input_size[i] * input_spacing[i] / target_spacing[i])) for i in range(3) ] resampler = itk.ResampleImageFilter.New(Input=image) @@ -782,7 +782,7 @@ def test_transforms( "moving_to_fixed_transform": moving_to_fixed_transform, "fixed_to_moving_transform": fixed_to_moving_transform, } - except (RuntimeError, Exception) as e: + except Exception as e: # noqa: BLE001 - any unreadable cache is regenerated logger.warning("Error loading transforms: %s; regenerating", e) moving_to_fixed_transform_path.unlink(missing_ok=True) fixed_to_moving_transform_path.unlink(missing_ok=True) diff --git a/tests/test_register_images_fireants.py b/tests/test_register_images_fireants.py index d2ed0116..edbea18f 100644 --- a/tests/test_register_images_fireants.py +++ b/tests/test_register_images_fireants.py @@ -113,6 +113,7 @@ def test_s1_ac5_return_keys_and_ncc_not_worse( T_DEF_GAIN = 0.30 T_INV = 0.5 +T_POINT_MM = 1.0 def _smooth_volume(size: int = 64, spacing_mm: float = 2.0) -> itk.Image: @@ -383,3 +384,100 @@ def test_s4_tuned_options_deformable_gain_on_smooth_warp() -> None: assert _one_minus_ncc(fixed, warped) <= (1.0 - T_DEF_GAIN) * _one_minus_ncc( fixed, moving ) + + +@pytest.mark.parametrize("preset", ["configure_for_heart_ct", "configure_for_lung_ct"]) +def test_scales_length_validated_after_preset_iteration_change(preset: str) -> None: + """Changing the iteration count after a preset cannot pair mismatched scales.""" + registrar = RegisterImagesFireants() + getattr(registrar, preset)() + assert len(registrar._scales()) == len(registrar.number_of_iterations) + + registrar.set_number_of_iterations([1, 0]) + with pytest.raises(ValueError, match="equal length"): + registrar._scales() + + registrar.set_scales([2.0, 1.0]) + assert registrar._scales() == [2.0, 1.0] + + +def test_mattes_rescale_keeps_negative_intensities_distinct() -> None: + """An all-negative CT keeps its contrast after the unit-range rescale.""" + array = np.linspace(-1000.0, -900.0, 4 * 4 * 4, dtype=np.float32).reshape(4, 4, 4) + image = itk.image_from_array(array) + image.SetSpacing((2.0, 2.0, 2.0)) + + rescaled = RegisterImagesFireants._rescale_unit_range(image) + + values = itk.array_from_image(rescaled) + assert values.min() == 0.0 + assert values.max() == 1.0 + assert len(np.unique(values)) == len(np.unique(array)) + assert tuple(rescaled.GetSpacing()) == (2.0, 2.0, 2.0) + + +@pytest.mark.requires_gpu +def test_mattes_fireants_image_is_unit_range() -> None: + """The Mattes metric hands FireANTs intensities inside [0, 1].""" + registrar = RegisterImagesFireants() + registrar.set_metric("Mattes") + array = np.linspace(-1000.0, -900.0, 4 * 4 * 4, dtype=np.float32).reshape(4, 4, 4) + + tensor = registrar._to_fireants_image(itk.image_from_array(array))() + + assert float(tensor.min()) == 0.0 + assert float(tensor.max()) == 1.0 + + +def _known_geometry_pair() -> tuple[itk.Image, itk.Image, itk.Transform]: + """Fixed and moving images on different grids related by a known transform. + + The moving grid differs from the fixed one in size, spacing and direction. + The returned transform maps fixed physical points to moving physical points, + and the moving image samples the fixed content through its inverse. + """ + fixed = _smooth_volume() + center = [63.0, 63.0, 63.0] + known = itk.Euler3DTransform[itk.D].New() + known.SetCenter(center) + known.SetRotation(np.deg2rad(5.0), 0.0, np.deg2rad(8.0)) + known.SetTranslation([6.0, -4.0, 3.0]) + inverse = known.GetInverseTransform() + + size = np.array([96, 60, 80]) + spacing = np.array([1.5, 2.5, 2.0]) + direction = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + moving_center = np.array(known.TransformPoint(center)) + origin = moving_center - direction @ (spacing * (size - 1) / 2.0) + grid = itk.image_from_array(np.zeros(size[::-1], dtype=np.float32)) + grid.SetSpacing(spacing.tolist()) + grid.SetOrigin(origin.tolist()) + grid.SetDirection(itk.matrix_from_array(direction)) + moving = ProcessTransforms().transform_image( + fixed, inverse, grid, interpolation_method="linear" + ) + return fixed, moving, known + + +@pytest.mark.requires_gpu +def test_deformable_composition_matches_known_point_mapping_across_geometries() -> None: + """Fixed-to-moving and its inverse map physical points as the known transform.""" + fixed, moving, known = _known_geometry_pair() + registrar = _deformable_registrar(fixed) + registrar.set_number_of_iterations([100, 50, 25]) + registrar.set_init_affine("Affine") + + result = registrar.register(moving_image=moving) + + margin = 0.25 + axis = np.linspace(margin * 126.0, (1.0 - margin) * 126.0, 6) + points = np.stack(np.meshgrid(axis, axis, axis, indexing="ij"), -1).reshape(-1, 3) + expected = np.array([known.TransformPoint(p.tolist()) for p in points]) + forward = result["fixed_to_moving_transform"] + backward = result["moving_to_fixed_transform"] + mapped = np.array([forward.TransformPoint(p.tolist()) for p in points]) + returned = np.array([backward.TransformPoint(q.tolist()) for q in expected]) + + assert np.linalg.norm(expected - points, axis=1).mean() > 5.0 + assert np.linalg.norm(mapped - expected, axis=1).mean() < T_POINT_MM + assert np.linalg.norm(returned - points, axis=1).mean() < T_POINT_MM diff --git a/tests/test_workflow_create_statistical_model.py b/tests/test_workflow_create_statistical_model.py index 58d9ccb1..e0649373 100644 --- a/tests/test_workflow_create_statistical_model.py +++ b/tests/test_workflow_create_statistical_model.py @@ -8,7 +8,7 @@ from __future__ import annotations import inspect -from typing import Any, cast +from typing import Any, ClassVar, cast import itk import numpy as np @@ -49,7 +49,7 @@ class _IdentityRegistrar: thresholds the projection tests assert. """ - reference_images: list[itk.Image] = [] + reference_images: ClassVar[list[itk.Image]] = [] def __init__( self, diff --git a/tutorials/tutorial_06_duke_heart_create_statistical_model.py b/tutorials/tutorial_06_duke_heart_create_statistical_model.py index 45d12edf..8662a86f 100644 --- a/tutorials/tutorial_06_duke_heart_create_statistical_model.py +++ b/tutorials/tutorial_06_duke_heart_create_statistical_model.py @@ -37,6 +37,7 @@ # Imports from __future__ import annotations +import contextlib import json import logging from pathlib import Path @@ -238,11 +239,9 @@ mode_surface_files: list[Path] = [] xvfb_started = False - try: + with contextlib.suppress(Exception): pv.start_xvfb() xvfb_started = True - except Exception: - pass try: for mode_idx in range(mode_count): @@ -279,10 +278,8 @@ # Pair start_xvfb with cleanup, guarded like the startup above so # environments without Xvfb (e.g. Windows, pyvista >= 0.48) are unaffected. if xvfb_started: - try: + with contextlib.suppress(Exception): pv.stop_xvfb() - except Exception: - pass tutorial_results = { "pca_model": pca_model, diff --git a/tutorials/tutorial_06_heart_create_statistical_model.py b/tutorials/tutorial_06_heart_create_statistical_model.py index b2ec9d2f..da2061d6 100644 --- a/tutorials/tutorial_06_heart_create_statistical_model.py +++ b/tutorials/tutorial_06_heart_create_statistical_model.py @@ -16,6 +16,7 @@ # Imports from __future__ import annotations +import contextlib import json import logging from pathlib import Path @@ -196,11 +197,9 @@ def read_model_surface(path: Path) -> pv.DataSet: mode_count = min(2, number_of_pca_components, len(components), len(eigenvalues)) xvfb_started = False - try: + with contextlib.suppress(Exception): pv.start_xvfb() xvfb_started = True - except Exception: - pass mode_surface_files: list[Path] = [] @@ -242,10 +241,8 @@ def read_model_surface(path: Path) -> pv.DataSet: # Pair start_xvfb with cleanup, guarded like the startup above so # environments without Xvfb (e.g. Windows, pyvista >= 0.48) are unaffected. if xvfb_started: - try: + with contextlib.suppress(Exception): pv.stop_xvfb() - except Exception: - pass tutorial_results = { "pca_model": pca_model, diff --git a/tutorials/tutorial_06_lung_create_statistical_model.py b/tutorials/tutorial_06_lung_create_statistical_model.py index 36d47f21..ccc80924 100644 --- a/tutorials/tutorial_06_lung_create_statistical_model.py +++ b/tutorials/tutorial_06_lung_create_statistical_model.py @@ -32,6 +32,7 @@ # Imports from __future__ import annotations +import contextlib import json import logging from pathlib import Path @@ -271,11 +272,9 @@ mode_surface_files: list[Path] = [] xvfb_started = False - try: + with contextlib.suppress(Exception): pv.start_xvfb() xvfb_started = True - except Exception: - pass try: for mode_idx in range(mode_count): @@ -312,10 +311,8 @@ # Pair start_xvfb with cleanup, guarded like the startup above so # environments without Xvfb (e.g. Windows, pyvista >= 0.48) are unaffected. if xvfb_started: - try: + with contextlib.suppress(Exception): pv.stop_xvfb() - except Exception: - pass tutorial_results = { "pca_model": pca_model, diff --git a/tutorials/tutorial_07_duke_heart_fit_statistical_model_to_patient.py b/tutorials/tutorial_07_duke_heart_fit_statistical_model_to_patient.py index f717608a..1509dd49 100644 --- a/tutorials/tutorial_07_duke_heart_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_07_duke_heart_fit_statistical_model_to_patient.py @@ -33,6 +33,7 @@ # Imports from __future__ import annotations +import contextlib import json import logging from pathlib import Path @@ -224,10 +225,8 @@ log_level=log_level, ) - try: + with contextlib.suppress(Exception): pv.start_xvfb() - except Exception: - pass screenshots: list[Path] = [] diff --git a/tutorials/tutorial_07_heart_fit_statistical_model_to_patient.py b/tutorials/tutorial_07_heart_fit_statistical_model_to_patient.py index 3420ac57..e48f2e18 100644 --- a/tutorials/tutorial_07_heart_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_07_heart_fit_statistical_model_to_patient.py @@ -16,6 +16,7 @@ # Imports from __future__ import annotations +import contextlib import json import logging from pathlib import Path @@ -216,10 +217,8 @@ log_level=log_level, ) - try: + with contextlib.suppress(Exception): pv.start_xvfb() - except Exception: - pass screenshots: list[Path] = [] diff --git a/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py b/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py index d763d08a..1aeb0118 100644 --- a/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py @@ -24,6 +24,7 @@ # Imports from __future__ import annotations +import contextlib import json import logging from pathlib import Path @@ -210,10 +211,8 @@ log_level=log_level, ) - try: + with contextlib.suppress(Exception): pv.start_xvfb() - except Exception: - pass screenshots: list[Path] = [] diff --git a/tutorials/tutorial_13_heart_and_lung_motion.py b/tutorials/tutorial_13_heart_and_lung_motion.py index da6add63..4d4e1355 100644 --- a/tutorials/tutorial_13_heart_and_lung_motion.py +++ b/tutorials/tutorial_13_heart_and_lung_motion.py @@ -222,8 +222,10 @@ cardiac_stages = [round(0.1 * k, 2) for k in range(10)] # Fraction of a cardiac cycle advanced per respiratory phase. The heart is # decoupled from breathing, so a value < 1 means a single beat continues into - # the next phase (1.0 locks exactly one full beat to each phase). - cardiac_cycles_per_phase = 0.333333333333 + # the next phase (1.0 locks exactly one full beat to each phase). The product + # with the number of respiratory phases must be a whole number of beats, or + # the cardiac position jumps where the sequence loops (0.3 x 10 = 3 beats). + cardiac_cycles_per_phase = 0.3 # Gaussian sigma (mm) that spreads each rhythm's surface-sampled # displacements into a continuous field. The lung shell is sparse relative # to the thorax it has to fill, so it is spread further than the heart. diff --git a/tutorials/tutorial_15_duke_heart_leave_one_out.py b/tutorials/tutorial_15_duke_heart_leave_one_out.py index 2b920045..69bd313d 100644 --- a/tutorials/tutorial_15_duke_heart_leave_one_out.py +++ b/tutorials/tutorial_15_duke_heart_leave_one_out.py @@ -242,11 +242,13 @@ def _write_loo_report( f"Folds: {len(held_out_cases)}", f"Held-out cases: {', '.join(held_out_cases)}", "", - "Each fold rebuilt the shape model without its held-out case, refitted " - "the cohort, retrained the MeshGraphNet and scored that case against " - "its own acquired labelmaps. Every cell is the mean over folds of the " - "fold's own mean over frames, plus or minus the standard deviation " - "across folds.", + ( + "Each fold rebuilt the shape model without its held-out case, refitted " + "the cohort, retrained the MeshGraphNet and scored that case against " + "its own acquired labelmaps. Every cell is the mean over folds of the " + "fold's own mean over frames, plus or minus the standard deviation " + "across folds." + ), "", "| Structure | " + " | ".join(title for _, title, _ in metrics) + " |", "|---|" + "---|" * len(metrics), @@ -315,10 +317,12 @@ def _displacement_error_lines( "", "## Displacement error, per fold", "", - "Distance between where the network puts each mesh point of the held-out " - "case and where that fold's own fit put it, pooled over every point and " - "frame. Measured point by point, so a chamber predicted the right size " - "in the wrong place is scored here and nowhere above.", + ( + "Distance between where the network puts each mesh point of the held-out " + "case and where that fold's own fit put it, pooled over every point and " + "frame. Measured point by point, so a chamber predicted the right size " + "in the wrong place is scored here and nowhere above." + ), "", f"- RMS across folds: {mean:.3f} +/- {std:.3f} mm", f"- 95th percentile across folds: {p95_mean:.3f} +/- {p95_std:.3f} mm", diff --git a/tutorials/tutorial_15_lung_leave_one_out.py b/tutorials/tutorial_15_lung_leave_one_out.py index cd942c6b..82346271 100644 --- a/tutorials/tutorial_15_lung_leave_one_out.py +++ b/tutorials/tutorial_15_lung_leave_one_out.py @@ -240,14 +240,18 @@ def _write_loo_report( f"Folds: {len(held_out_cases)}", f"Held-out cases: {', '.join(held_out_cases)}", "", - "Each fold rebuilt the shape model without its held-out case, refitted " - "the cohort, retrained the MeshGraphNet and scored that case against " - "its own per-phase segmentations. Every cell is the mean over folds of " - "the fold's own mean over phases, plus or minus the standard deviation " - "across folds.", + ( + "Each fold rebuilt the shape model without its held-out case, refitted " + "the cohort, retrained the MeshGraphNet and scored that case against " + "its own per-phase segmentations. Every cell is the mean over folds of " + "the fold's own mean over phases, plus or minus the standard deviation " + "across folds." + ), "", - "Dice on a lobe is dominated by the size of the lobe rather than by how " - "far it moved; surface RMSE is what resolves the motion here.", + ( + "Dice on a lobe is dominated by the size of the lobe rather than by how " + "far it moved; surface RMSE is what resolves the motion here." + ), "", "| Structure | " + " | ".join(title for _, title, _ in metrics) + " |", "|---|" + "---|" * len(metrics), @@ -316,10 +320,12 @@ def _displacement_error_lines( "", "## Displacement error, per fold", "", - "Distance between where the network puts each mesh point of the held-out " - "case and where that fold's own fit put it, pooled over every point and " - "phase. Measured point by point, so a lobe predicted the right size in " - "the wrong place is scored here and nowhere above.", + ( + "Distance between where the network puts each mesh point of the held-out " + "case and where that fold's own fit put it, pooled over every point and " + "phase. Measured point by point, so a lobe predicted the right size in " + "the wrong place is scored here and nowhere above." + ), "", f"- RMS across folds: {mean:.3f} +/- {std:.3f} mm", f"- 95th percentile across folds: {p95_mean:.3f} +/- {p95_std:.3f} mm", diff --git a/utils/ai_agent_github_reviews.py b/utils/ai_agent_github_reviews.py index 815ecb8c..4cef65e8 100644 --- a/utils/ai_agent_github_reviews.py +++ b/utils/ai_agent_github_reviews.py @@ -470,7 +470,7 @@ def resolve_review_threads(thread_ids: list[str], repo: str) -> None: if not thread_ids: return print(f"[*] Resolving {len(thread_ids)} inline-comment thread(s)...") - owner, name = repo.split("/", 1) + _owner, _name = repo.split("/", 1) for tid in thread_ids: try: _gh_graphql(_RESOLVE_THREAD_MUTATION, {"threadId": tid}) diff --git a/utils/mp4_to_gif.py b/utils/mp4_to_gif.py index 27fb9898..d2f6bf23 100644 --- a/utils/mp4_to_gif.py +++ b/utils/mp4_to_gif.py @@ -71,8 +71,10 @@ def convert(mp4_path: Path, gif_path: Path, width: int, fps: int) -> None: "-i", str(palette), "-lavfi", - f"{filters}[v];[v][1:v]paletteuse=dither=bayer:bayer_scale=5" - ":diff_mode=rectangle", + ( + f"{filters}[v];[v][1:v]paletteuse=dither=bayer:bayer_scale=5" + ":diff_mode=rectangle" + ), "-loop", "0", str(gif_path), diff --git a/utils/setup_feature_worktree.py b/utils/setup_feature_worktree.py index 7810d1a9..a266726e 100644 --- a/utils/setup_feature_worktree.py +++ b/utils/setup_feature_worktree.py @@ -178,6 +178,7 @@ def get_current_branch() -> str: ["git", "symbolic-ref", "--short", "HEAD"], text=True, capture_output=True, + check=False, ) if result.returncode == 0: return result.stdout.strip() @@ -197,6 +198,7 @@ def branch_exists(branch_name: str) -> bool: ["git", "branch", "--list", branch_name], text=True, capture_output=True, + check=False, ) return bool(result.stdout.strip()) From a50b680a4dd970a685d364ea761a32548001e0e8 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Thu, 15 Oct 2026 12:20:23 -0400 Subject: [PATCH 3/4] ENH: Add fireants to nightly test config --- .github/workflows/ci.yml | 5 +++++ .github/workflows/nightly-health.yml | 2 ++ .github/workflows/test-slow.yml | 1 + 3 files changed, 8 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9cb3cefb..159cadfa 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -369,6 +369,11 @@ jobs: $env:CL = "/std:c++20" python -m uv pip install -e ".[dev_cuda13]" --no-build-isolation-package torch-scatter + - name: Install FireANTs + # --no-deps: FireANTs pins would otherwise override the CUDA torch stack. + run: | + python -m uv pip install fireants --no-deps + - name: Assert CUDA is accessible run: | python -c " diff --git a/.github/workflows/nightly-health.yml b/.github/workflows/nightly-health.yml index 0b1b3a6a..ba2a1610 100644 --- a/.github/workflows/nightly-health.yml +++ b/.github/workflows/nightly-health.yml @@ -138,6 +138,8 @@ jobs: if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } python -m uv pip install -e ".[dev_cuda13]" --no-build-isolation-package torch-scatter if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } + python -m uv pip install fireants --no-deps + if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } - name: Assert CUDA is accessible run: | diff --git a/.github/workflows/test-slow.yml b/.github/workflows/test-slow.yml index f0ab6f60..8834d7bb 100644 --- a/.github/workflows/test-slow.yml +++ b/.github/workflows/test-slow.yml @@ -60,6 +60,7 @@ jobs: # Invoke via python -m uv so uv targets the active venv interpreter. run: | python -m uv pip install -e ".[dev_cuda13]" --no-build-isolation-package torch-scatter + python -m uv pip install fireants --no-deps - name: Assert CUDA is accessible run: | From 3d52f46f8dd80c03c860e53f3652a5c2d0800a60 Mon Sep 17 00:00:00 2001 From: Stephen Aylward Date: Thu, 15 Oct 2026 12:33:27 -0400 Subject: [PATCH 4/4] ENH: Minor CI test improvement --- .github/workflows/test-slow.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test-slow.yml b/.github/workflows/test-slow.yml index 8834d7bb..3506055d 100644 --- a/.github/workflows/test-slow.yml +++ b/.github/workflows/test-slow.yml @@ -60,6 +60,7 @@ jobs: # Invoke via python -m uv so uv targets the active venv interpreter. run: | python -m uv pip install -e ".[dev_cuda13]" --no-build-isolation-package torch-scatter + if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } python -m uv pip install fireants --no-deps - name: Assert CUDA is accessible