Skip to content

perf(dynacell): CELL-Diff on DynaCell-lite in 1.43 h per fit on one H100 - #520

Merged
alxndrkalinin merged 35 commits into
dynacell-modelsfrom
dynacell-celldiff-fast-inference
Oct 6, 2026
Merged

alxndrkalinin merged 35 commits into
dynacell-modelsfrom
dynacell-celldiff-fast-inference

Conversation

@alxndrkalinin

@alxndrkalinin alxndrkalinin commented Oct 3, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

CellDiff-3D inference of one fit over its whole DynaCell-lite test set (iPSC + A549 mock/DENV/ZIKV) now
takes 1.43 h on one H100 (5151 s for the nucleus iPSC-trained fit, measured; 5185 s for the A549-trained
one). The full benchmark's model orderings hold except two pairs that the full benchmark itself only barely
resolves (see Validation). The production recipe on the same set costs ~100 GPU-h per A549 condition. Evidence:
applications/dynacell/experiments/2026-10-01_celldiff-fast-inference/README.md (gitignored).

Where production spends its time (H100, measured): adaptive dopri5, ~56-66 velocity evaluations per
tile, over 512^2 tiles overlapping 4/256/256 (66 tiles per A549 volume), in fp32.

The lite recipe (model_overlays/celldiff_predict_lite.yml) keeps production's tile grid and changes:

  • the sampler: midpoint on 8 time points at the Beta(0.4, 0.3) quantiles, 14 velocity evaluations per
    tile. The grid is packed toward the data end (last step 0.006), where the fine texture the deep-feature
    metrics read is laid down;
  • Z overlap 4 -> 2 (8 instead of 11 Z tiles per A549 volume);
  • batch 4 tiles per forward;
  • precision and compile: bf16 autocast and torch.compile of a channels-last copy of the net. After
    compile the forward is conv-bound, and cuDNN spent ~19 % of it transposing NCDHW <-> NDHWC. The copy
    stores conv weights channels_last_3d in bf16 and computes GroupNorm over the NDHWC view. Same-GPU
    A100: 384.8 -> 308.7 ms per 4-patch forward, same 0.0033 relative L2 to the fp32 net;
  • cudnn.benchmark (trainer.benchmark): -6.7 % end to end on H100, pixels move by relative L2 0.0016.

Rejected on the way:

  • 12 cosine evaluations (and 10 with Z overlap 4): the A549-trained nucleus fit's deep-feature metrics shift
    enough to flip orderings (13 / 760 and 21 / 760 pairs, 1 and 8 confidently wrong).
  • Beta(0.5, 0.7) at 14 evaluations: lowest whole-tile error against dopri5, but a 0.13 last step loses fine
    texture (A549 DynaCLR KID +20 between-FOV SD; membrane 25 / 1786 pairs flipped, 4 confidently wrong; same on
    an A40, so not a GPU effect). Grids are now scored on high-pass error, which orders the known outcomes
    correctly.
  • multistep solvers (AB2, DPM-Solver++(2M)): worse than midpoint at every evaluation count.
  • larger tiles: whole volumes hallucinate on A549; full-field 8-plane slabs broke the iPSC-trained membrane
    fit on A549 (Dice 0.37 -> 0.24, 10 flips).

Changes

  • viscy-models: cosine and Beta(p, q) time_schedules for the flow-matching ODE sampler (Beta(0.5, 0.5) is
    the cosine grid); scipy joins the celldiff extra. Non-finite shapes and grids that are not strictly
    increasing in float32 raise at construction.
  • dynacell engine: predict_sampling_method, predict_time_schedule, predict_compile (defaults keep
    production behaviour). Compiled inference runs a channels-last copy that only predict_step samples
    through (a per-call compiled argument of the generate* methods); each predict run rebuilds it from the
    current weights, so validation never samples stale weights, and a module-supplied callback frees it
    from on_exception when a run raises.
  • dynacell: sampling noise drawn on the CPU (seeded CUDA randn differs by GPU model: one recipe on A40
    vs H100 = 0.44 relative L2).
  • lite generator: CELL-Diff joins the lite roster (12 fits incl. joint) with the fast overlay and an H100
    2 h / 32G profile; 48 lite predict leaves + grouped lite buckets regenerated; a composition test pins
    every CELL-Diff lite leaf to the overlay's recipe.
  • test fixes: wall-limit test judged the model family from the absolute path; C-cond spy kwargs.
  • internal review follow-ups: the compiled copy is scoped to one predict run; the CPU-noise test fails on a
    device-side draw; docstrings for step counts, the NDHWC GroupNorm and predict_compile.

Validation

Timing (one H100, nucleus iPSC-trained fit, whole lite test set, job 37040233): iPSC 260 s + A549 mock
1646 + DENV 1623 + ZIKV 1622 = 5151 s = 1.43 h (16 cosine evaluations: 6399 s without, 5971 s with
cudnn.benchmark).

Ordering. For each organelle, the lite metrics of all three CELL-Diff fits (iPSC-, A549-, joint-trained)
are replaced by the fast recipe's, and every model pair the full benchmark resolves (|gap| > 2 paired
FOV-bootstrap SD, 2000 draws) is re-ranked on the lite. A flip is a resolved pair whose lite sign reverses;
confidently wrong = the lite resolves the reversed sign too. In brackets: the 16-evaluation cosine recipe.

organelle vs other models: flips / resolvable (confidently wrong) among CELL-Diff fits
nucleus 4 / 2250 (1) [6 / 2250 (0)] 1 / 130 (0) [3 / 130 (0)]
membrane 2 / 1786 (0) [1 / 1786 (0)] 1 / 118 (0) [0 / 118 (0)]
ER 0 / 2071 (0) [0 / 2071 (0)] 1 / 86 (0) [2 / 86 (0)]
mitochondria 2 / 1877 (1) [3 / 1877 (1)] 0 / 113 (0) [0 / 113 (0)]

The two confidently wrong pairs, both CELL-Diff gaining on an A549 deep-feature cosine:

  • mitochondria, A549 ZIKV, MorphEm median cosine, joint-trained CELL-Diff vs pix2pix3d (joint): full gap
    -0.0108 (z = 2.02), production on the lite frames -0.0171, lite +0.0182 (16 cosine evaluations: +0.0251).
    It follows the Z overlap, not the evaluation count: Z overlap 4 keeps it (at ~2.4 h), Z overlap 3 and 18 or
    24 evaluations at Z overlap 2 do not.
  • nucleus, A549 mock, DINOv3 median cosine, iPSC-trained CELL-Diff vs FNet3D-VSCyto3DAug (iPSC): full gap
    -0.0207 (z = 2.05), production on the lite frames -0.0266, lite +0.0188 (z = 2.47). The 16-evaluation recipe
    flips it too (+0.0146, z = 1.73, below the confidence threshold).

On the most step-sensitive fit (A549-trained nucleus, A549 mock) the 14-evaluation grid tracks production
more closely than 16 cosine evaluations: DynaCLR KID shift +0.46 vs +0.52 between-FOV SD, CellDINO KID +0.44
vs +0.71, DINOv3 KID +0.19 vs +0.86.

Needs a decision: accept the two reversed pairs (documented in applications/dynacell/CLAUDE.md) at
1.43 h, or revert 4d7d863 for the 16-evaluation recipe at 1.66 h with one of them.

Test plan

  • pytest packages/viscy-models/tests/test_celldiff/test_transport.py (16 passed) and the engine Beta-grid and failed-predict tests
  • pytest applications/dynacell/tests/test_benchmark_config_composition.py (lite CELL-Diff leaves resolve to the validated recipe)
  • full suite on the merged branch (81df638): 5242 passed, 193 skipped, 1 xfailed, 0 failed

🤖 Generated with Claude Code

https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF

alxndrkalinin and others added 19 commits October 1, 2026 21:13
… patch

CELLDiffNet asserts its input equals input_spatial_size because the ViT
bottleneck's sincos positional embedding and token grid are fixed at
construction. The conv encoder/decoder are size-agnostic, so inference
over a larger field only needs those two rebuilt.

with_input_size() returns a shallow view that shares every learned
parameter and swaps in a bottleneck view with the embedding recomputed
for the new grid. The embedding is a function of integer grid
coordinates, so every position the training grid has keeps exactly its
trained value; training-time size checks stay strict on the original.
spatial_multiple exposes the sizes an input must be a multiple of. The
grid computation moves into a helper both paths share.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
…DE sampler

Fixed-grid solvers step on the num_steps output points, which were
always linspace(t0, t1). Adaptive dopri5 on CellDiff spends most of its
evaluations near both ends of [0, 1] (210 NFE over 3 patches: 40 in
t < 0.1 and 31 in t >= 0.9), so a uniform grid wastes fixed steps in the
middle. time_schedule="cosine" spaces the points as
(1 - cos(pi * s)) / 2, dense at both ends; "uniform" stays the default
and an unknown name raises.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
…hing

Production CellDiff predicts integrate adaptive dopri5 (~56-66 NFE per
tile) over overlapping 512^2 tiles in fp32: ~100 GPU-h per A549
condition on an H100. Four predict-time options, all defaulting to the
old behaviour:

- predict_sampling_method: any torchdiffeq method; fixed-grid ones
  (euler, midpoint, rk4) take num_generate_steps - 1 steps.
- predict_time_schedule: uniform or cosine time points.
- predict_tile_size: tile for sliding_window / iterative, -1 meaning the
  whole (padded) axis, e.g. [8, -1, -1] for full-field 8-plane slabs.
- predict_compile: torch.compile the velocity net. The forward is
  memory-bound (GroupNorm and elementwise ops ~70 % of GPU time), so
  compiled bf16 runs an 8x512x512 forward in 57 ms vs 294 ms in fp32 on
  an H100.

Tiles other than the training patch run through
CELLDiffNet.with_input_size, cached per size in
CELLDiff3DVS.inference_net. predict_method="generate" now pads to the
sizes the net accepts and runs the whole input in one pass.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
A seeded CUDA randn returns different numbers on different GPU models,
so a seeded CellDiff prediction depended on the GPU it ran on: the same
recipe on an A40 and an H100 came out 0.44 relative L2 apart, i.e. two
independent samples. Drawing the initial noise from the CPU generator
(seeded by seed_everything) and moving it to the device makes seeded
predictions reproducible across GPU types and lets two sampler settings
be compared pixel-for-pixel from the same noise.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
…hmarks/

The wall-limit test called a leaf CELL-Diff if any component of its
absolute path started with "celldiff". In a checkout whose directory
name starts with it (a worktree named celldiff-fast-inference), every
predict leaf counted as CELL-Diff and the test failed for all non-
CELL-Diff leaves. Only the path relative to the benchmarks root names
the model.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
predict_step passes sampling_method and time_schedule to generate since
the fast-inference options landed, so the spy's (cond, num_steps)
signature raised TypeError and test_cond_predict_reads_the_mask_source
failed on the branch head (948715e). Forward the extra kwargs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Gwn9hwgT8uJSoWaCxCAm8J
Drop predict_tile_size, whole-volume generate and the per-size network
views: every larger-than-patch geometry tried changed what the model
predicts on A549. Whole volumes (40-48 planes) hallucinate texture;
full-field 8-plane slabs and 512^2 tiles with narrow YX overlap both
pull the iPSC-trained membrane fit's A549 Dice from 0.37 to 0.24, even
with an exact dopri5 solve (0.20), while production's 512^2 tiles with
256 px anchored overlap and the fast sampler land at 0.42. The speed
now comes from the sampler, bf16 and compile alone.

inference_net() caches the compiled net in a dict so it never registers
as a submodule (that would duplicate the weights in the state dict).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
After torch.compile the CellDiff forward is conv-bound, and about a
fifth of the conv time was cuDNN transposing NCDHW tensors to its
native NDHWC layout and back around every conv. The compiled-inference
path now compiles a copy of the net whose conv weights are
channels_last_3d, stored in the active autocast dtype so no per-forward
cast lays them out contiguous again. Its GroupNorms reduce over the
NDHWC view: nn.GroupNorm on a channels-last input compiles to strided
reductions about 8x slower, which cancelled the gain.

Measured on one A100-80 (8x512x512, 4 patches per forward, bf16):
384.8 -> 308.7 ms per forward, relative L2 to the fp32 net unchanged
at 0.0033. End to end, the iPSC lite bucket of the nucleus fit took
468 -> 412 s. Predictions move 0.004 relative L2 from the old code,
against 0.032 between an A100 and an H100 running the same code.

The copy holds the weights it was built with, so entering train mode
drops it.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
CELL-Diff was left off the lite roster because its production predict
(adaptive dopri5 over 512^2 tiles, ~100 GPU-h per A549 condition) is
86 % of predict GPU-h. Its lite leaves now keep the checkpoint and the
production tile grid (512^2, YX overlap 256, iterative anchoring) but
swap in a fast overlay, celldiff_predict_lite.yml: Z overlap 2,
midpoint on a cosine time grid with 16 velocity evaluations, batch 4,
bf16 + torch.compile. A lite CELL-Diff prediction is therefore a new
sample, not a subset of the full one.

Measured on one H100: 6399 s (1.78 h) for the nucleus iPSC-trained
fit's whole lite test set. With all three CELL-Diff fits of an
organelle replaced, the full benchmark's model orderings hold:
nucleus 5 / 2245 resolvable pairs flip, membrane 1 / 1783, none
confidently wrong (resolved on both sides with opposite sign). At 12
evaluations the A549-trained nucleus fit flips 13 / 760 (1 confidently
wrong), so 16 is the floor; it is also the most that fits in 2 h.

The generator adds the joint fits for CELL-Diff only, keeps the
iterative leaves (denoise / sliding_window are variants), and drops
the leaf-level sampler keys that would override the overlay. A
dedicated H100 / 2 h / 32G profile carries the lite leaves (peak RSS
5.3-6.2 GB); the wall-limit test learns it. Evidence:
applications/dynacell/experiments/2026-10-01_celldiff-fast-inference/
(gitignored).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The compiled channels-last copy was cached until train() cleared it,
but Lightning leaves validation by assigning module.training directly
and never calls train(). With predict_compile on during fit, the copy
built for the sanity-check samples would have served every later
epoch's validation samples from the initial weights. The cache also
outlived its build context: a second predict run at another precision
failed on a bf16-weighted copy (Input type (float) and bias type
(c10::BFloat16) should be the same), and a device move or
load_state_dict between runs kept the old copy.

The engine now turns compile_inference on in on_predict_start and off
in on_predict_end, and setting the flag drops any cached copy, so each
predict run compiles from its own weights, device and autocast state
and validation sampling never uses the compiled path. The train()
override is gone.

Found in internal review of the branch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The test compared seeded noise for a CPU conditioning tensor, where a
draw on the tensor's own device gives the same numbers, so it passed
on the code the CPU-noise fix replaced. It now also passes a meta
tensor (standing in for a GPU) and records where torch.randn draws:
the draw must happen on the CPU and only the result move to meta.

Found in internal review of the branch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The docstring promised a channels_last_3d result, but a contiguous
input comes back contiguous; only inside the channels-last copy, where
every input is channels-last, is the output channels-last.

Found in internal review of the branch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The tiled and trajectory generate methods still called num_steps the
number of ODE integration steps. It is the number of time points the
sampler returns: dopri5 picks its own steps between them, and a
fixed-grid solver takes num_steps - 1 steps on them, as generate()
already says. The tiled methods now point there; the trajectory
methods say the points are the returned frames.

Found in internal review of the branch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
num_generate_steps and num_log_steps are the number of ODE time
points, not integration steps; with a fixed-grid predict sampler the
solver takes num_generate_steps - 1 steps. Matches the generate
methods' docstrings.

Found in internal review of the branch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
predict_method="denoise" calls the velocity net directly, so
predict_compile has no effect there; the docstring now names the two
methods it does apply to.

Found in internal review of the branch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The generator test only compares committed files with the generator's
own output, so a leaf that kept a production sampler key (the
production leaves set predict_method, predict_overlap and
num_generate_steps, which outrank the overlay) would pass it. Each of
the 48 CELL-Diff lite leaves now has to compose to the overlay's model
and trainer settings and batch size; the production leaves differ on
seven of them. z_window_size is left out: it follows each test set's
Z extent (48 on A549).

Found in internal review of the branch.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
Records what the lite CELL-Diff leaves run (production tile grid, Z
overlap 2, 16-evaluation midpoint on a cosine grid, bf16 + compiled
channels-last copy), what it costs (1.78 h per fit's whole lite test
set on an H100), how well it keeps the full benchmark's orderings,
and the one reversal it introduces: the fast sampler raises the
in-domain fits' A549 MorphEm scores, flipping mitochondria ZIKV
MorphEm median cosine for joint-trained CELL-Diff vs pix2pix3d, a pair
the full benchmark resolves only at z = 2.02. Also the two shortcuts
that were measured to break ordering.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🔵 Needs a closer look

Accepting the documented mitochondria ranking reversal requires a human benchmark decision.

Review effort: Balanced
Findings: 1 Medium severity · 1 Low severity

Open (2)
What changed in this PR

Adds a faster CELL-Diff inference recipe to DynaCell-lite, using shared sampling improvements in viscy-models and dedicated benchmark configurations.

Changes:

  • Adds cosine sampling, CPU-generated noise and per-run compiled channels-last inference.
  • Adds 48 lite prediction configurations, an H100 profile and grouped evaluation entries.
  • Extends regression tests and documents performance and benchmark-ordering trade-offs.
File Description
packages/​viscy-models/​tests/​test_celldiff/​test_transport.py Tests cosine scheduling and sampling.
packages/​viscy-models/​src/​viscy_models/​celldiff/​modules/​transport/​transport.py Exposes sampling schedules.
packages/​viscy-models/​src/​viscy_models/​celldiff/​modules/​transport/​integrators.py Implements cosine time grids.
applications/​dynacell/​tools/​generate_lite_benchmark_configs.py Generates CELL-Diff lite configurations.
applications/​dynacell/​tests/​test_training_integration.py Forwards sampling arguments through the spy.
applications/​dynacell/​tests/​test_engine.py Tests sampling, layout, lifecycle and noise.
applications/​dynacell/​tests/​test_benchmark_config_composition.py Checks lite recipes and hardware profiles.
applications/​dynacell/​src/​dynacell/​engine.py Adds prediction controls and compilation lifecycle.
applications/​dynacell/​src/​dynacell/​celldiff_wrapper.py Implements optimized inference and CPU noise.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__ipsc_confocal_lite.yml Nucleus: joint-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_zikv_lite.yml Nucleus: joint-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_mock_lite.yml Nucleus: joint-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_denv_lite.yml Nucleus: joint-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​ipsc_confocal/​predict__ipsc_confocal__iterative_lite.yml Nucleus: iPSC-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​ipsc_confocal/​predict__a549_mantis_zikv_lite.yml Nucleus: iPSC-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​ipsc_confocal/​predict__a549_mantis_mock_lite.yml Nucleus: iPSC-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​ipsc_confocal/​predict__a549_mantis_denv_lite.yml Nucleus: iPSC-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​a549_mantis/​predict__ipsc_confocal_lite.yml Nucleus: A549-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​a549_mantis/​predict__a549_mantis_zikv_lite.yml Nucleus: A549-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​a549_mantis/​predict__a549_mantis_mock_lite.yml Nucleus: A549-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​nucleus/​celldiff/​a549_mantis/​predict__a549_mantis_denv_lite.yml Nucleus: A549-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__ipsc_confocal_lite.yml Mitochondria: joint-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_zikv_lite.yml Mitochondria: joint-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_mock_lite.yml Mitochondria: joint-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_denv_lite.yml Mitochondria: joint-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​ipsc_confocal/​predict__ipsc_confocal__iterative_lite.yml Mitochondria: iPSC-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​ipsc_confocal/​predict__a549_mantis_zikv_lite.yml Mitochondria: iPSC-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​ipsc_confocal/​predict__a549_mantis_mock_lite.yml Mitochondria: iPSC-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​ipsc_confocal/​predict__a549_mantis_denv_lite.yml Mitochondria: iPSC-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​a549_mantis/​predict__ipsc_confocal_lite.yml Mitochondria: A549-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​a549_mantis/​predict__a549_mantis_zikv_lite.yml Mitochondria: A549-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​a549_mantis/​predict__a549_mantis_mock_lite.yml Mitochondria: A549-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​mito/​celldiff/​a549_mantis/​predict__a549_mantis_denv_lite.yml Mitochondria: A549-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__ipsc_confocal_lite.yml Membrane: joint-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_zikv_lite.yml Membrane: joint-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_mock_lite.yml Membrane: joint-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_denv_lite.yml Membrane: joint-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​ipsc_confocal/​predict__ipsc_confocal__iterative_lite.yml Membrane: iPSC-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​ipsc_confocal/​predict__a549_mantis_zikv_lite.yml Membrane: iPSC-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​ipsc_confocal/​predict__a549_mantis_mock_lite.yml Membrane: iPSC-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​ipsc_confocal/​predict__a549_mantis_denv_lite.yml Membrane: iPSC-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​a549_mantis/​predict__ipsc_confocal_lite.yml Membrane: A549-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​a549_mantis/​predict__a549_mantis_zikv_lite.yml Membrane: A549-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​a549_mantis/​predict__a549_mantis_mock_lite.yml Membrane: A549-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​membrane/​celldiff/​a549_mantis/​predict__a549_mantis_denv_lite.yml Membrane: A549-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__ipsc_confocal_lite.yml ER: joint-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_zikv_lite.yml ER: joint-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_mock_lite.yml ER: joint-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​joint_ipsc_confocal_a549_mantis/​predict__a549_mantis_denv_lite.yml ER: joint-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​ipsc_confocal/​predict__ipsc_confocal__iterative_lite.yml ER: iPSC-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​ipsc_confocal/​predict__a549_mantis_zikv_lite.yml ER: iPSC-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​ipsc_confocal/​predict__a549_mantis_mock_lite.yml ER: iPSC-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​ipsc_confocal/​predict__a549_mantis_denv_lite.yml ER: iPSC-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​a549_mantis/​predict__ipsc_confocal_lite.yml ER: A549-trained → lite iPSC.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​a549_mantis/​predict__a549_mantis_zikv_lite.yml ER: A549-trained → lite A549 ZIKV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​a549_mantis/​predict__a549_mantis_mock_lite.yml ER: A549-trained → lite A549 mock.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​er/​celldiff/​a549_mantis/​predict__a549_mantis_denv_lite.yml ER: A549-trained → lite A549 DENV.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​shared/​model/​model_overlays/​celldiff_predict_lite.yml Defines the fast inference recipe.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​shared/​model/​launcher_profiles/​hardware_predict_celldiff_lite.yml Allocates one H100 with a two-hour limit.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​nucleus_joint__lite/​eval_grouped.yaml Adds joint-trained nucleus evaluation.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​nucleus_ipsc_trained__lite/​eval_grouped.yaml Adds iPSC-trained nucleus CELL-Diff entries.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​nucleus_a549_trained__lite/​eval_grouped.yaml Adds A549-trained nucleus CELL-Diff entries.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​mitochondria_joint__lite/​eval_grouped.yaml Adds joint-trained mitochondria evaluation.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​mitochondria_ipsc_trained__lite/​eval_grouped.yaml Adds iPSC-trained mitochondria CELL-Diff entries.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​mitochondria_a549_trained__lite/​eval_grouped.yaml Adds A549-trained mitochondria CELL-Diff entries.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​membrane_joint__lite/​eval_grouped.yaml Adds joint-trained membrane evaluation.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​membrane_ipsc_trained__lite/​eval_grouped.yaml Adds iPSC-trained membrane CELL-Diff entries.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​membrane_a549_trained__lite/​eval_grouped.yaml Adds A549-trained membrane CELL-Diff entries.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​er_joint__lite/​eval_grouped.yaml Adds joint-trained ER evaluation.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​er_ipsc_trained__lite/​eval_grouped.yaml Adds iPSC-trained ER CELL-Diff entries.
applications/​dynacell/​configs/​benchmarks/​virtual_staining/​_internal/​leaf/​grouped/​er_a549_trained__lite/​eval_grouped.yaml Adds A549-trained ER CELL-Diff entries.
applications/​dynacell/​CLAUDE.md Documents the recipe and ordering caveats.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

alxndrkalinin and others added 2 commits October 3, 2026 05:14
The overlay header said model orderings were kept, but the fast recipe
reverses one marginal pair (mitochondria, A549 ZIKV, MorphEm median
cosine, joint CELL-Diff vs joint pix2pix3d). Say so and point to the
caveat in applications/dynacell/CLAUDE.md.

Raised in Copilot review of #520.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The predict_compile tests built the compiled copy but never called it,
and torch.compile is lazy, so a compile failure or a wrong compiled
forward would have passed. The new test runs Trainer.predict at
bf16-mixed on CPU, eager and then compiled, on the same weights, seed
and input: a callback sees the compiled copy during the batch, nothing
is cached after the run, and the samples agree to 0.0026 relative L2
(eager bf16 vs fp32 is 0.0028; the tolerance is 0.02). ~27 s on 4 CPUs.

Raised in Copilot review of #520 (and the internal review).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
@alxndrkalinin

Copy link
Copy Markdown
Collaborator Author

Follow-up on the one reversed mitochondria pair (A549 ZIKV MorphEm median cosine, joint CELL-Diff vs joint pix2pix3d; full z = 2.02). All variants below were run on the mito joint fit's ZIKV bucket and checked with 2000 bootstrap draws.

variant joint MorphEm cosine (production on lite 0.732) the pair
this PR: Z overlap 2, 16 evaluations 0.774 (+1.10 SD) reversed
Z overlap 2, 18 evaluations + cudnn.benchmark 0.777 (+1.17 SD) reversed
Z overlap 2, 24 evaluations 0.777 (+1.17 SD) reversed
Z overlap 3, 16 evaluations + cudnn.benchmark 0.760 (+0.73 SD) reversed (lite z 2.4)
Z overlap 4, 16 evaluations 0.726 (-0.16 SD) kept, but a CELL-Diff-internal CellDINO KID pair reverses instead

The shift comes from the Z overlap, not from solver accuracy: more evaluations move it further from production. Only Z overlap 4 keeps the pair, and at 16 evaluations it needs about 8800 s on an H100 (derived), over the 2 h budget. So no affordable recipe removes this reversal; the PR keeps Z overlap 2 / 16 evaluations with the caveat documented.

Separately, cudnn.benchmark speeds the shipped forward by 6.5 % on an H100 (41.3 -> 38.6 ms per tile; 8.2 % on H200). It is not in this PR, since it would change the validated samples; trainer.benchmark: true in the overlay would bring the 1.78 h to about 1.67 h.

🤖 Generated with Claude Code

alxndrkalinin and others added 2 commits October 4, 2026 13:16
on_predict_start turned a module-wide compile_inference flag on and
on_predict_end turned it off, but when predict raises Lightning runs
only on_exception hooks (call._interrupt), so the flag stayed on with
that run's snapshot cached. A later fit or validate on the same module
would then sample its validation images from the stale weights.

The compiled path is now a per-call argument: generate,
generate_sliding_window and generate_iterative take compiled (default
False), predict_step passes predict_compile, and the predict hooks only
drop the cached copy (start: rebuild from the current weights; end:
free memory). No state can leave compiled sampling switched on. A
regression test raises inside a compiled predict, moves the weights,
and checks validation sampling against a fresh module; it fails on the
previous code. predict_compile's docstring now also names "generate".

Raised in Codex adversarial review of #520.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
@alxndrkalinin

Copy link
Copy Markdown
Collaborator Author

Codex adversarial review finding addressed in 47e5403 (medium: compiled snapshot survives a failed predict).

Verified against Lightning 2.6.5: when trainer.predict raises, call._interrupt runs only callback/strategy on_exception hooks, so on_predict_end never ran and the module-wide compile_inference flag stayed on with that run's snapshot cached; validation sampling on the same module would then use stale weights.

Fix: no module-wide flag any more. generate, generate_sliding_window and generate_iterative take a per-call compiled argument (default False); predict_step passes predict_compile; on_predict_start / on_predict_end only drop the cached copy. A regression test raises inside a compiled Trainer.predict, moves the weights, and checks validation sampling against a fresh module with the same weights; it fails on the previous code ("Tensor-likes are not close"). Full suite: 5112 passed, 179 skipped, 1 xfailed, 0 failed.

🤖 Generated with Claude Code

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🔵 Needs a closer look

Accepting the documented mitochondria ranking reversal requires an explicit benchmark-owner decision.

Review effort: Balanced
Findings: 1 Medium severity

Open (1)
Resolved since last review (2)

alxndrkalinin and others added 9 commits October 5, 2026 23:26
Lightning's trainer.benchmark lets cuDNN time its convolution
algorithms per shape. On one H100 the nucleus iPSC-trained fit's whole
lite test set drops from 6399 s to 5971 s (1.78 h -> 1.66 h, jobs
36976986 vs 37036905). The predictions move by relL2 0.0015-0.0016
against the cudnn.benchmark-off run (6 FOVs per bucket, iPSC and A549
mock), 20x below the A100-vs-H100 drift the P2-M16c orderings were
validated across, so the ordering checks carry over.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
ODESolver (and through it Sampler.sample_ode, the CELLDiff3DVS generate
methods and DynacellFlowMatching.predict_time_schedule) now also takes
time_schedule=(p, q): the Beta(p, q) quantiles of num_steps uniform
points. The existing cosine grid is Beta(0.5, 0.5) and uniform is
Beta(1, 1), so the family contains both; p > q packs the points toward
the data end of the interval. Non-positive or malformed shapes raise.

The quantile comes from scipy.special.betaincinv, so scipy joins the
celldiff extra (already in the workspace through dynacell; the lock
gains two lines).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
…, 0.3) grid

The lite overlay drops from 9 cosine time points (16 midpoint
evaluations per tile) to 8 points at the Beta(0.4, 0.3) quantiles (14
evaluations), packed toward the data end (last step 0.006). One H100
now predicts the nucleus iPSC-trained fit's whole lite test set in
5151 s (1.43 h; was 5971 s with cosine 16; job 37040233).

Validated on all 12 fits (experiment P2-M14B43b), each organelle with
its three fits replaced, 2000-draw bootstrap: 8 of 7984 resolvable
pairs against other models flip (cosine 16: 10), 3 of 447 among the
CELL-Diff fits (5). Two flips are resolved the other way on the lite,
both marginal on the full benchmark (z 2.02 and 2.05); cosine 16 flips
both as well and resolves one.

The grid was chosen on high-pass error against dopri5, not whole-tile
error: Beta(0.5, 0.7) at 14 evaluations had the lowest whole-tile error
but a 0.13 last step that lost fine texture (A549 DynaCLR KID +20
between-FOV SD, 25 of 1786 membrane pairs flipped).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The (p, q) guard checked min(shape) > 0, which a trailing NaN passes
(min(0.5, nan) is 0.5): (0.5, nan) built an all-NaN time grid and
(0.5, inf) one with NaN interior points, and the failure only surfaced
later inside odeint. Both shapes must now be finite and positive.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The Beta schedule was only tested on ODESolver over t in [0, 1]. Add an
engine test that predict_step on the iterative path with
predict_time_schedule=(0.4, 0.3) evaluates the net at the Beta grid's
left points in every Z tile, and check that the grid scales onto a
(t0, t1) interval other than (0, 1).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
The overlay header's summary sentence still named only the Z overlap,
the sampler and the precision after trainer.benchmark was added.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
"Orderings hold except two pairs" blurred the 8 of 7984 flips against
other models and 3 of 447 among the CELL-Diff fits with the 2 pairs the
lite resolves the other way. State all three, as CLAUDE.md does.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
… rules

The DynaCell-lite bullet read "Z overlap 3 and 18 or 24 evaluations" as
one configuration and turned one counterexample into "pick grids by
high-pass error". Say what was measured: 18 and 24 evaluations were at
Z overlap 2, and high-pass error ranked the three measured grids in the
order of their deep-feature shifts while whole-tile error did not.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
@alxndrkalinin alxndrkalinin changed the title perf(dynacell): CELL-Diff on DynaCell-lite in 1.78 h per fit on one H100 perf(dynacell): CELL-Diff on DynaCell-lite in 1.43 h per fit on one H100 Oct 6, 2026
@alxndrkalinin
alxndrkalinin requested a balanced review from Copilot October 6, 2026 11:40

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🔵 Needs a closer look

Accepting the two documented ranking reversals requires a human decision on the speed–accuracy tradeoff.

Review effort: Balanced
Findings: 1 High severity · 1 Medium severity

Open (2)

alxndrkalinin and others added 3 commits October 6, 2026 04:50
A Beta shape can pass validation yet put distinct quantiles within
float32 rounding of each other: (1.0, 0.1) at 8 points ends in
[..., 1.0, 1.0], and odeint only rejects the non-increasing grid once
prediction is running. ODESolver now checks that every schedule's grid
is strictly increasing when it builds it and names the schedule in the
error.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
…ises

Lightning skips on_predict_end when Trainer.predict raises and runs only
callback on_exception hooks. The compiled channels-last copy is not a
registered submodule, so after a failed run it kept its device memory
(and moving the model to CPU did not release it) until the next predict.
DynacellFlowMatching.configure_callbacks now attaches a callback whose
on_exception drops the copy; on_predict_start still rebuilds it for each
run. The failed-predict test now checks the copy existed mid-run and is
gone afterwards, besides validation still sampling the live weights.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XvkM1gKifWmoDPq8Drd4JF
@alxndrkalinin
alxndrkalinin merged commit 19765e7 into dynacell-models Oct 6, 2026
26 checks passed
alxndrkalinin added a commit that referenced this pull request Oct 7, 2026
generate_trajectory always used the adaptive dopri5 default, so a
caller that wants every intermediate state could not use the fixed-grid
samplers generate() gained in #520. On one 8x512x512 demo patch, asking
dopri5 for 50 frames cost 80-104 network evaluations; midpoint on the
Beta(0.4, 0.3) grid with 8 time points costs 14 and lands 0.03-0.06
relative L2 from it. Add sampling_method/time_schedule with the old
defaults. The test pins the 14-evaluation midpoint path and that the
last frame equals generate() under the same seed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015kuht3npKtb2o5E6h19xbA
alxndrkalinin added a commit that referenced this pull request Oct 7, 2026
…on sampler

The Space asked dopri5 for 50 trajectory frames, which cost 80-104
network evaluations per patch: 61-71 s per call on ZeroGPU and 86-96 s
on a shared A40, close to the 120 s GPU cap. Use the DynaCell-lite
sampler (#520): midpoint on a Beta(0.4, 0.3) grid, 8 time points = 14
evaluations, under bf16 autocast (no torch.compile: ZeroGPU forks a
fresh process per call, so warm-up would be paid every time).

Same noise on the demo patches: 0.03-0.06 relative L2 to dopri5, equal
or higher Spectral PCC on all four organelles, no visible difference.
Through the Space's own code on the A40: 19-31 s per call including
model load. The slider now sets time points (8-24, default 8; each
extra point costs 2 evaluations). The Space's VisCy pin moves to
84cc1ae, the parent commit, which adds the generate_trajectory
sampler arguments on top of #520.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015kuht3npKtb2o5E6h19xbA
alxndrkalinin added a commit that referenced this pull request Nov 3, 2026
* feat(viscy-utils): add the segmentation-aware and background losses

Port the Spotlight-motivated losses from dynacell-models so they live
in the shared package before any application consumes them:
SegAuxDice (seg_aux.py) with the SAUNA-map and soft-clDice weight maps
(seg_aux_maps.py), the segmenter-consistency loss (seg_consistency.py),
the Track H background low-pass target (background_lowpass.py), and
SpotlightLoss(return_components=True). Defaults keep SpotlightLoss's
existing behavior; no package imports an application.

Source on dynacell-models: 40de793 (#498), 6c28a31 (#521),
0cbeda1 (#524), 050f553 (#530).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-utils): build foreground masks with CLAHE, median and Otsu

Port the Spotlight fg-mask recipe from dynacell-models.
generate_fg_masks now binarizes a CLAHE-equalized, 5^3-median-filtered
volume at an Otsu threshold picked on that volume after a sigma=5
Gaussian; compute_otsu uses the same full-resolution smoothing, so the
otsu_grid_spacing parameter is gone from generate_normalization_metadata
and VisCyTrainer.preprocess. The median and Gaussian run through cubic
on the GPU when cubic and a device are present, else SciPy; cubic stays
an optional import. torch is imported lazily so that light callers of
write_meta_field do not pay for it.

Behavior change: Spotlight masks and Otsu thresholds written by
`viscy preprocess` differ from main's.

Source on dynacell-models: 40de793 (#498), d087d96 (#510).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-utils): make prediction writes resumable, identified and write-once

Port the HCSPredictionWriter and prediction_metadata work from
dynacell-models:
- z_reduction="center" and z_planes (slab-only predicts write only the
  planes they predicted, and record them);
- per-FOV completion markers and the checkpoint/settings identity that
  predict resume and the eval caches key on (prediction_metadata.py);
- refuse writes into an existing store when a FOV outside the run has
  unmarked prediction channels, or when the array would grow;
- assemble each volume in memory and write it once on a background
  thread (bytes unchanged);
- raise NotImplementedError when trainer.world_size > 1, instead of
  letting ranks race on one store.

Source on dynacell-models: 2b848a6, 1be4466, 74fa18a (#497),
8e41498 (#511), 996adb4 (#526), 6c83fff (#531), 8e74a7e (#532).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-utils): guard every fit with checkpoint and health callbacks

Port the training guardrails from dynacell-models. VisCyCLI now
injects, on every `fit`:
- a latest-weights ModelCheckpoint that owns last.ckpt (the monitored
  checkpoint gets save_last=False), so last.ckpt always means the
  latest epoch, inheriting the leaf's write cadence;
- MonitorHealthCheck, which raises when the monitored metric is
  non-finite or bit-identical for 10 validation epochs;
- OptimizerHealthCheck, which raises after 5 consecutive non-finite
  losses or on GradScaler collapse.
Also adds the opt-in LinearLRCooldown callback and raises the CLI
subprocess test timeout to 120 s (it timed out on main).

Behavior change: this applies to cytoland, dynaclr and dynacell fits.

Source on dynacell-models: 20e54db, daa7477, bc0f444, 74fa18a
(#497), 92693ae, 2f9ef31 (#528).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): keep SSIM finite under mixed precision; lazy exports

Port two viscy-utils changes from dynacell-models:
- ssim_25d / ms_ssim_25d: clamp the variances, bound the covariance by
  Cauchy-Schwarz through a safe sqrt, floor the stability constants
  symmetrically, and run MS-SSIM outside autocast, so 16-mixed training
  cannot NaN the gradient. SSIM/MS-SSIM loss values change slightly for
  MixedLoss users.
- viscy_utils/__init__ resolves its exports lazily (PEP 562), keeping
  torch and lightning off light import paths.

Source on dynacell-models: a556bc9 (#484), 130e034, cee4e41,
272051c, 59343ad, 8b2413f (#509).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-data): slab-only and read-once predict, uint8 foreground masks

Port the viscy-data side of the dynacell-models predict and Spotlight
work:
- HCSDataModule.predict_z_planes and SlidingWindowDataset z_planes /
  read_z_planes, so a 2D model predicts only its focus slab;
- cache_volume: with num_workers == 0 predict reads each (position, t)
  volume once and, when every transform is pointwise, normalizes it
  once;
- foreground masks are read in their stored dtype (uint8) and cast to
  float on the device in on_after_batch_transfer.

Behavior change: every app's predict at num_workers=0 takes the
read-once path.

Source on dynacell-models: 40de793 (#498), 996adb4 (#526),
6c83fff (#531), 8e74a7e (#532).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-models): per-axis CellDiff patches and mask-conditioning channels

Port from dynacell-models:
- patch_size may be a per-axis (z, y, x) tuple through the patch
  embedding, transformer unpatchify and ViT bottleneck, which lets
  CELLDiffNet run Z-preserving at Z=1 for the 2D track (an int is still
  accepted);
- CELLDiffNet(cond_channels=1) adds extra conditioning input channels
  (default keeps the old single-channel net).

Source on dynacell-models: 0126fe6, 0617ed7, 40de793 (#498).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* perf(viscy-models): non-uniform ODE time grids for CellDiff sampling

Port the time_schedule option of the CellDiff ODE integrators from
dynacell-models: uniform (default, unchanged), cosine, or a Beta(p, q)
grid packed toward the data end, which is what lets the DynaCell-lite
recipe sample in 14 velocity evaluations. scipy joins the celldiff
extra for betaincinv.

Source on dynacell-models: 19765e7 (#520).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-models): add the MorphEm feature extractor

Port MorphEmModel from dynacell-models, the fifth deep-feature
extractor of the DynaCell evaluation. transformers is imported lazily
and the trust_remote_code model is pinned to a Hub commit.

Source on dynacell-models: d3f6422 (#479), 4918d2e, 13cb336,
c9c6914.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-models): let the PatchGAN discriminator preserve the Z axis

Port preserve_z from dynacell-models, so the 3D PatchGAN runs on Z=1
inputs for pix2pix2d (default off, unchanged).

Source on dynacell-models: c2d0481.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-models): UNet3D downsample_z for FNet-2D; reject non-5D input

Port from dynacell-models:
- UNet3D(downsample_z=True): False keeps Z through the encoder, which
  is how FNet2D runs on single planes (default unchanged);
- UNet3DBase.forward raises on non-5D input, and a comment flags the
  ConvTranspose checkerboard artifact.

Source on dynacell-models: a556bc9 (#484), 83a1d69, 16b1d31.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-transforms): stop BatchedRandAffined blanking singleton axes

Port 9640537 from dynacell-models: a size-1 spatial axis (Z=1 for 2D
models) gets an identity row in the affine, instead of being resampled
off the volume and zeroed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* perf(viscy-transforms): mark the sampled normalizers pointwise

NormalizeSampled and MinMaxSampled set pointwise = True, which lets the
read-once predict path in viscy-data normalize a whole volume once
instead of once per window.

Source on dynacell-models: 8e74a7e (#532).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(viscy-transforms): make the Gaussian-noise statistics test deterministic

Port 3409560 (#513) from dynacell-models.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build: parallel test runs with pytest-xdist

Port dfcf9a5 (#512) from dynacell-models: pytest-xdist joins the test
group, and a root conftest.py pins OMP/MKL/OpenBLAS to one thread under
each xdist worker. Without the pin every worker's torch sizes its pool
to all cores, which made `-n 11` 10x slower than it should be and timed
out subprocess tests. CLAUDE.md documents the usage.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test: register pytest markers at the root and make them strict

Port the marker half of c50300f (#489) from dynacell-models.
strict_markers = true turns an unregistered marker into a collection
error, so a typo'd marker can no longer silently decorate nothing.
hpc_integration moves from the cytoland and dynaclr conftests to the
root: an addinivalue_line registration only holds when that conftest
is in the invocation's ancestry. dynaclr drops its own testpaths so
the root ini applies.

requires_data is registered too: main's
applications/dynacell/tools/generate_grouped_eval_configs_test.py uses
it, and strict markers would otherwise drop that module at collection.
(dynacell-models no longer uses the marker; the dynacell PR that
follows removes the registration with it.)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build: put the dynacell hf_demo on the pytest pythonpath

Port the root half of 7b7cb43 (#523) from dynacell-models so the HF
demo's tests import its app module. The directory already exists on
main.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(dynaclr): join annotations on the full fov_name

Port d1ea492 from dynacell-models. The orchestrated linear-classifier
fixture joined on the first fov_name component, which failed 7 tests
on main once pytest collects dynaclr's colocated *_test.py.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(cytoland): correct the Spotlight recipe comments

Port the cytoland recipe comment fixes from 40de793 (#498): the
preprocess flags take explicit values (`--compute_otsu true
--compute_fg_masks true`), fg_threshold and nonzero_threshold are inert
while fg_mask_key is set, and the Otsu-centred target is what makes 0
the FG/BG boundary. Comments only.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(cytoland): install the VCP tutorials' unpublished deps from GitHub

Port the tutorial fix from d3f6422 (#479): cytoland and viscy-utils
are not on PyPI yet, so the tutorials install them from the GitHub
monorepo; their siblings resolve from PyPI. Adds the jupytext-paired
notebooks (no outputs).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs: cut CLAUDE.md content the code or global memory already states

Port 93ac234 from dynacell-models: drop the repo-layout tree, setup
and common-command blocks, lint/docstring rules (pyproject.toml is the
source) and the generic coding-philosophy sections. The DynaCell job
runbook from the same file arrives with the dynacell PR, since it
documents tools that only exist there.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* ci: run the viscy-utils tests in their own job

viscy-utils had no CI job on main or on dynacell-models: four of its
test files ran inside the dynacell job and the rest (losses, metrics,
meta_utils, optimizer health, CLI, lazy imports) ran nowhere. Mirror
the viscy-data job on ubuntu-latest for Python 3.12 and 3.13, and make
the gate job wait for it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* ci(dynacell): run the launcher, engine and prediction-writer suites

Port the test-dynacell-configs job from dynacell-models, minus the
eight test files that only arrive with the dynacell PR: sync
viscy-utils' test group with dynacell, and add test_submit_benchmark_batch,
test_engine, test_pipeline_cache, feature_metrics_test and the
viscy-utils prediction-writer, prediction-metadata, monitor-health and
LR-cooldown tests. The dynacell-models PR target is not carried over.

Source on dynacell-models: 33d2fee, 74fa18a (#497), 2d1ed32 (#506),
4450efd (#518), 2798dee (#519), 2f9ef31 (#528).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build: re-lock against the ported packages and main's dynacell

Start from dynacell-models' uv.lock and re-resolve with `uv lock`
against this tree, whose applications/dynacell/pyproject.toml is still
main's. Against dynacell-models' lock this drops the ten dynacell-only
packages (dinov3, etils, ftfy, jsonpath-rw, mlcroissant, pandas-stubs,
ply, rdflib, submitit, termcolor) and resolves cubic back to 0.7.0;
the dynacell PR restores them. Against main's lock it adds pytest-xdist
and moves boto3/botocore 1.43.0 -> 1.43.20, cellpose 4.1.1 -> 4.2.1.1,
pyarrow 19.0.1 -> 24.0.0, s3transfer and tracksdata; torch and
lightning are unchanged. `uv lock --check` exits 0 (main's own lock
does not).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): declare torchvision, which evaluation.metrics imports

viscy_utils.evaluation.metrics imports torchvision.ops at module level,
and losses.MixedLoss imports that module, but viscy-utils never listed
torchvision. Installs only worked because another workspace member
pulled it in. The new test-utils CI job syncs viscy-utils alone and
failed to collect six test modules with ModuleNotFoundError. The lock
gains only the dependency edge; no version moves.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* ci(dynacell): keep main's batch-launcher test and the writer tests out

Two corrections to the test-dynacell-configs job:
- main's tests/test_submit_benchmark_batch.py resolves leaves under
  /hpc and fails on a runner (5 PermissionErrors in CI). The
  dynacell-models version that runs there arrives with the DynaCell
  PR, which adds it back.
- the four viscy-utils test files (prediction writer, prediction
  metadata, monitor health, LR cooldown) now run in the test-utils job
  on both Python versions, so the dynacell job no longer repeats them
  or syncs viscy-utils' test group for them.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(viscy-utils): share the fit-callbacks lookup between CLI injectors

_inject_checkpoint_guardrails and _inject_optimizer_health_guard each
repeated the same walk from the config to fit's trainer.callbacks plus
the already-injected check; the `else config` branch of that walk could
never run behind the `subcommand != "fit"` return. Both now call
_fit_callbacks_without. _callback_field collapses its two identical
branches, and the monitored checkpoint's init_args are kept from the
search loop instead of being looked up twice. No behavior change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(viscy-utils): drop MonitorHealthCheck's unused check_finite flag

Nothing sets check_finite (the CLI injects monitor and patience only,
and no config or test passes it), so the callback always checks
finiteness; it now does so with math.isfinite instead of building a
tensor for one float.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(viscy-utils): stop describing fg masks as built from Otsu thresholds

generate_fg_masks does not read the stored otsu_threshold; it binarizes
a CLAHE-equalized, median-filtered volume. Its summary line and
VisCyTrainer.preprocess's compute_fg_masks entry still said "from Otsu
thresholds".

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build: leave the hf_demo pytest path to the DynaCell PR

Nothing on this branch imports from applications/dynacell/examples/
hf_demo: its only test, build_demo_data_test.py, arrives with the
DynaCell PR, which adds the pythonpath entry back with it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): guard a fit that sets no callbacks

Trainer.callbacks defaults to None, and the CLI injectors returned
early unless they found a list, so a fit that set no callbacks ran
without OptimizerHealthCheck, the guard meant for every fit. The
shared lookup now gives such a fit an empty list to append to (None
and [] mean the same to Trainer). The checkpoint guardrails still
need a monitored ModelCheckpoint and remain a no-op there.

Reported by Copilot on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(viscy-utils): say what the AMP skip counter actually counts

The counter resets only when GradScaler grows its scale, so successful
optimizer steps can sit between the reductions it counts; the error
message claimed every step was skipped "in a row". Describe it as
reductions without growth in between, which is what it detects: the
backward overflowing faster than growth_interval clean steps can
restore the scale. Behavior unchanged.

Reported by Copilot on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): resolve viscy_utils.callbacks names lazily

VisCyCLI now injects MonitorHealthCheck and OptimizerHealthCheck into
every fit by class path through viscy_utils.callbacks, whose __init__
eagerly imported OnlineEvalCallback and with it scikit-learn, which
only the eval extra installs. A fit on a base install would therefore
fail before training. The package now resolves its exports lazily
(PEP 562), like viscy_utils itself. The new test blocks sklearn and
imports the health callbacks: the eager __init__ fails it with
ModuleNotFoundError, the lazy one passes.

Reported by Copilot on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(cytoland): install viscy-data from GitHub in the VCP tutorials

The tutorials install viscy-utils from the GitHub monorepo and let the
siblings resolve from PyPI. The writer they predict with now reads
SlidingWindowDataset.z_planes, which the PyPI viscy-data 0.1.0 does not
have (checked in its wheel), so predict would fail with AttributeError
in on_predict_start. Install viscy-data from GitHub too, in the .py
sources and their paired notebooks.

Found by /code-review on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): return an empty CLAHE result when the clip range is zero

foreground_mask_from_volume guards a constant volume, but a blank
volume with a few hot pixels (under 0.2% of voxels) has equal 0.1/99.9
percentiles too, and equalize_clahe then divided by zero and fed NaN to
CLAHE (RuntimeWarning: divide by zero, reproduced on a (20, 64, 64)
volume). equalize_clahe now returns zeros there, which the mask's
existing flat-histogram guard turns into an empty mask, as for a
constant volume. Volumes with any contrast are unchanged.

Found by /code-review on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): keep a step-cadence leaf's last.ckpt on its step trigger

The injected latest-weights ModelCheckpoint takes over save_last from
the monitored one but copied only every_n_epochs (defaulting to 1). A
leaf that checkpoints every N steps, typical for max_steps runs with
long epochs, would get last.ckpt only at epoch boundaries, so a
preempted run resumed from much older weights. The injected callback
now copies the monitored callback's every_n_train_steps or
train_time_interval when set; epoch-cadence leaves are unchanged. No
current leaf combines a monitored checkpoint with a step trigger.

Found by /code-review on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): take last.ckpt from every monitored checkpoint

The guardrail turned save_last off only on the first monitored
ModelCheckpoint. A leaf with two monitored checkpoints (say loss and a
validation metric) kept a second last.ckpt writer next to the injected
latest-weights callback, so last.ckpt would alternate between "latest"
and "that monitor's cadence". All monitored callbacks now give it up;
the first still names the monitor, directory and cadence. No current
leaf has two monitored checkpoints.

Found by /code-review on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): write AnnData zarr without switching pandas string inference off

anndata 0.12.6 (capped below 0.12.9, which needs pandas<3) cannot write
pandas 3's default Arrow-backed strings to zarr. The workaround was
global: write_embedding_dataset set pd.options.future.infer_string =
False in the caller's process, and the test conftest set it for the
whole session. Tests therefore never ran under the configuration users
have, and the other writers (EmbeddingSnapshot's rewrite, annotation
convert) worked only after the embedding writer had flipped the option.

write_anndata_zarr converts obs/var Arrow strings (columns, categories,
index) to object just before writing, reusing the conversion
append_to_anndata_zarr already had; all three writers and the tests
use it, and both global switches are gone. New tests pin a round trip
under the default, that plain write_zarr still fails (so the TODO is
dropped deliberately once anndata 0.13 lands), and that the embedding
writer leaves the option alone.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MzpgExfYMKCgogfnhcSgDc

* fix(dynaclr): drop the session-wide pandas string-inference switch

dynaclr's pytest_configure set pd.options.future.infer_string = False
for the whole pytest session, so in a repo-wide run every package's
tests ran with pandas 3's default off (it hid the viscy-utils guard
added in the previous commit). split_embeddings flipped it in the
caller's process too. Both go; the eleven test writes and the split
writer use viscy_utils' write_anndata_zarr instead. The cellanome
embedding scripts keep their own switch: each owns its process and
has no test to check a change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MzpgExfYMKCgogfnhcSgDc

---------

Co-authored-by: Alexandr Kalinin <alxndrkalinin@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
alxndrkalinin added a commit that referenced this pull request Nov 4, 2026
* feat(viscy-utils): add the segmentation-aware and background losses

Port the Spotlight-motivated losses from dynacell-models so they live
in the shared package before any application consumes them:
SegAuxDice (seg_aux.py) with the SAUNA-map and soft-clDice weight maps
(seg_aux_maps.py), the segmenter-consistency loss (seg_consistency.py),
the Track H background low-pass target (background_lowpass.py), and
SpotlightLoss(return_components=True). Defaults keep SpotlightLoss's
existing behavior; no package imports an application.

Source on dynacell-models: 40de793 (#498), 6c28a31 (#521),
0cbeda1 (#524), 050f553 (#530).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-utils): build foreground masks with CLAHE, median and Otsu

Port the Spotlight fg-mask recipe from dynacell-models.
generate_fg_masks now binarizes a CLAHE-equalized, 5^3-median-filtered
volume at an Otsu threshold picked on that volume after a sigma=5
Gaussian; compute_otsu uses the same full-resolution smoothing, so the
otsu_grid_spacing parameter is gone from generate_normalization_metadata
and VisCyTrainer.preprocess. The median and Gaussian run through cubic
on the GPU when cubic and a device are present, else SciPy; cubic stays
an optional import. torch is imported lazily so that light callers of
write_meta_field do not pay for it.

Behavior change: Spotlight masks and Otsu thresholds written by
`viscy preprocess` differ from main's.

Source on dynacell-models: 40de793 (#498), d087d96 (#510).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-utils): make prediction writes resumable, identified and write-once

Port the HCSPredictionWriter and prediction_metadata work from
dynacell-models:
- z_reduction="center" and z_planes (slab-only predicts write only the
  planes they predicted, and record them);
- per-FOV completion markers and the checkpoint/settings identity that
  predict resume and the eval caches key on (prediction_metadata.py);
- refuse writes into an existing store when a FOV outside the run has
  unmarked prediction channels, or when the array would grow;
- assemble each volume in memory and write it once on a background
  thread (bytes unchanged);
- raise NotImplementedError when trainer.world_size > 1, instead of
  letting ranks race on one store.

Source on dynacell-models: 2b848a6, 1be4466, 74fa18a (#497),
8e41498 (#511), 996adb4 (#526), 6c83fff (#531), 8e74a7e (#532).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-utils): guard every fit with checkpoint and health callbacks

Port the training guardrails from dynacell-models. VisCyCLI now
injects, on every `fit`:
- a latest-weights ModelCheckpoint that owns last.ckpt (the monitored
  checkpoint gets save_last=False), so last.ckpt always means the
  latest epoch, inheriting the leaf's write cadence;
- MonitorHealthCheck, which raises when the monitored metric is
  non-finite or bit-identical for 10 validation epochs;
- OptimizerHealthCheck, which raises after 5 consecutive non-finite
  losses or on GradScaler collapse.
Also adds the opt-in LinearLRCooldown callback and raises the CLI
subprocess test timeout to 120 s (it timed out on main).

Behavior change: this applies to cytoland, dynaclr and dynacell fits.

Source on dynacell-models: 20e54db, daa7477, bc0f444, 74fa18a
(#497), 92693ae, 2f9ef31 (#528).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): keep SSIM finite under mixed precision; lazy exports

Port two viscy-utils changes from dynacell-models:
- ssim_25d / ms_ssim_25d: clamp the variances, bound the covariance by
  Cauchy-Schwarz through a safe sqrt, floor the stability constants
  symmetrically, and run MS-SSIM outside autocast, so 16-mixed training
  cannot NaN the gradient. SSIM/MS-SSIM loss values change slightly for
  MixedLoss users.
- viscy_utils/__init__ resolves its exports lazily (PEP 562), keeping
  torch and lightning off light import paths.

Source on dynacell-models: a556bc9 (#484), 130e034, cee4e41,
272051c, 59343ad, 8b2413f (#509).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-data): slab-only and read-once predict, uint8 foreground masks

Port the viscy-data side of the dynacell-models predict and Spotlight
work:
- HCSDataModule.predict_z_planes and SlidingWindowDataset z_planes /
  read_z_planes, so a 2D model predicts only its focus slab;
- cache_volume: with num_workers == 0 predict reads each (position, t)
  volume once and, when every transform is pointwise, normalizes it
  once;
- foreground masks are read in their stored dtype (uint8) and cast to
  float on the device in on_after_batch_transfer.

Behavior change: every app's predict at num_workers=0 takes the
read-once path.

Source on dynacell-models: 40de793 (#498), 996adb4 (#526),
6c83fff (#531), 8e74a7e (#532).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-models): per-axis CellDiff patches and mask-conditioning channels

Port from dynacell-models:
- patch_size may be a per-axis (z, y, x) tuple through the patch
  embedding, transformer unpatchify and ViT bottleneck, which lets
  CELLDiffNet run Z-preserving at Z=1 for the 2D track (an int is still
  accepted);
- CELLDiffNet(cond_channels=1) adds extra conditioning input channels
  (default keeps the old single-channel net).

Source on dynacell-models: 0126fe6, 0617ed7, 40de793 (#498).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* perf(viscy-models): non-uniform ODE time grids for CellDiff sampling

Port the time_schedule option of the CellDiff ODE integrators from
dynacell-models: uniform (default, unchanged), cosine, or a Beta(p, q)
grid packed toward the data end, which is what lets the DynaCell-lite
recipe sample in 14 velocity evaluations. scipy joins the celldiff
extra for betaincinv.

Source on dynacell-models: 19765e7 (#520).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-models): add the MorphEm feature extractor

Port MorphEmModel from dynacell-models, the fifth deep-feature
extractor of the DynaCell evaluation. transformers is imported lazily
and the trust_remote_code model is pinned to a Hub commit.

Source on dynacell-models: d3f6422 (#479), 4918d2e, 13cb336,
c9c6914.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-models): let the PatchGAN discriminator preserve the Z axis

Port preserve_z from dynacell-models, so the 3D PatchGAN runs on Z=1
inputs for pix2pix2d (default off, unchanged).

Source on dynacell-models: c2d0481.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(viscy-models): UNet3D downsample_z for FNet-2D; reject non-5D input

Port from dynacell-models:
- UNet3D(downsample_z=True): False keeps Z through the encoder, which
  is how FNet2D runs on single planes (default unchanged);
- UNet3DBase.forward raises on non-5D input, and a comment flags the
  ConvTranspose checkerboard artifact.

Source on dynacell-models: a556bc9 (#484), 83a1d69, 16b1d31.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-transforms): stop BatchedRandAffined blanking singleton axes

Port 9640537 from dynacell-models: a size-1 spatial axis (Z=1 for 2D
models) gets an identity row in the affine, instead of being resampled
off the volume and zeroed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* perf(viscy-transforms): mark the sampled normalizers pointwise

NormalizeSampled and MinMaxSampled set pointwise = True, which lets the
read-once predict path in viscy-data normalize a whole volume once
instead of once per window.

Source on dynacell-models: 8e74a7e (#532).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(viscy-transforms): make the Gaussian-noise statistics test deterministic

Port 3409560 (#513) from dynacell-models.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build: parallel test runs with pytest-xdist

Port dfcf9a5 (#512) from dynacell-models: pytest-xdist joins the test
group, and a root conftest.py pins OMP/MKL/OpenBLAS to one thread under
each xdist worker. Without the pin every worker's torch sizes its pool
to all cores, which made `-n 11` 10x slower than it should be and timed
out subprocess tests. CLAUDE.md documents the usage.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test: register pytest markers at the root and make them strict

Port the marker half of c50300f (#489) from dynacell-models.
strict_markers = true turns an unregistered marker into a collection
error, so a typo'd marker can no longer silently decorate nothing.
hpc_integration moves from the cytoland and dynaclr conftests to the
root: an addinivalue_line registration only holds when that conftest
is in the invocation's ancestry. dynaclr drops its own testpaths so
the root ini applies.

requires_data is registered too: main's
applications/dynacell/tools/generate_grouped_eval_configs_test.py uses
it, and strict markers would otherwise drop that module at collection.
(dynacell-models no longer uses the marker; the dynacell PR that
follows removes the registration with it.)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build: put the dynacell hf_demo on the pytest pythonpath

Port the root half of 7b7cb43 (#523) from dynacell-models so the HF
demo's tests import its app module. The directory already exists on
main.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(dynaclr): join annotations on the full fov_name

Port d1ea492 from dynacell-models. The orchestrated linear-classifier
fixture joined on the first fov_name component, which failed 7 tests
on main once pytest collects dynaclr's colocated *_test.py.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(cytoland): correct the Spotlight recipe comments

Port the cytoland recipe comment fixes from 40de793 (#498): the
preprocess flags take explicit values (`--compute_otsu true
--compute_fg_masks true`), fg_threshold and nonzero_threshold are inert
while fg_mask_key is set, and the Otsu-centred target is what makes 0
the FG/BG boundary. Comments only.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(cytoland): install the VCP tutorials' unpublished deps from GitHub

Port the tutorial fix from d3f6422 (#479): cytoland and viscy-utils
are not on PyPI yet, so the tutorials install them from the GitHub
monorepo; their siblings resolve from PyPI. Adds the jupytext-paired
notebooks (no outputs).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs: cut CLAUDE.md content the code or global memory already states

Port 93ac234 from dynacell-models: drop the repo-layout tree, setup
and common-command blocks, lint/docstring rules (pyproject.toml is the
source) and the generic coding-philosophy sections. The DynaCell job
runbook from the same file arrives with the dynacell PR, since it
documents tools that only exist there.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* ci: run the viscy-utils tests in their own job

viscy-utils had no CI job on main or on dynacell-models: four of its
test files ran inside the dynacell job and the rest (losses, metrics,
meta_utils, optimizer health, CLI, lazy imports) ran nowhere. Mirror
the viscy-data job on ubuntu-latest for Python 3.12 and 3.13, and make
the gate job wait for it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* ci(dynacell): run the launcher, engine and prediction-writer suites

Port the test-dynacell-configs job from dynacell-models, minus the
eight test files that only arrive with the dynacell PR: sync
viscy-utils' test group with dynacell, and add test_submit_benchmark_batch,
test_engine, test_pipeline_cache, feature_metrics_test and the
viscy-utils prediction-writer, prediction-metadata, monitor-health and
LR-cooldown tests. The dynacell-models PR target is not carried over.

Source on dynacell-models: 33d2fee, 74fa18a (#497), 2d1ed32 (#506),
4450efd (#518), 2798dee (#519), 2f9ef31 (#528).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build: re-lock against the ported packages and main's dynacell

Start from dynacell-models' uv.lock and re-resolve with `uv lock`
against this tree, whose applications/dynacell/pyproject.toml is still
main's. Against dynacell-models' lock this drops the ten dynacell-only
packages (dinov3, etils, ftfy, jsonpath-rw, mlcroissant, pandas-stubs,
ply, rdflib, submitit, termcolor) and resolves cubic back to 0.7.0;
the dynacell PR restores them. Against main's lock it adds pytest-xdist
and moves boto3/botocore 1.43.0 -> 1.43.20, cellpose 4.1.1 -> 4.2.1.1,
pyarrow 19.0.1 -> 24.0.0, s3transfer and tracksdata; torch and
lightning are unchanged. `uv lock --check` exits 0 (main's own lock
does not).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(dynacell): DynaCell v1.1 benchmark configs from dynacell-models

Port every DynaCell config path of dynacell-models verbatim: the model
leaves (baselines, A549/joint/HEK/brightfield, the 2D track with its
seed replicates and warmup-matched refit, the temporal and
patch-vs-augmentation ablations, CELL-Diff-R2 and the DynaCell-lite
leaves), the grouped and instance-AP eval buckets of the paper models,
the shared overlays, train/predict sets and launcher profiles, and the
deletion of 155 stale config files. The leaves are generated; review
the generators and guard tests in the code commits, not each leaf.

Spotlight v2 arm leaves, their four fit overlays and their 231
spotlight_v2_* eval buckets are left out; they follow in their own PR.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* feat(dynacell): DynaCell v1.1 application code from dynacell-models

Port every DynaCell code path of dynacell-models verbatim (226
non-merge commits since 2026-06-16; PRs #479-#530): the evaluation
pipeline (cpdino/MorphEm/focus-aware 2D, focus2d slab metrics,
foreground-limited metrics, MicroMS3IM, scale-invariant pixel metrics,
CP shared space, provenance and cache invalidation), the paths
registry, preprocessing/registry/croissant/distribution, the launcher
and stall watchdog, the lite benchmark and CELL-Diff fast inference,
the 2D track, the HF demo, docs, and the engine, celldiff_wrapper and
mask_conditioning modules exactly as on dynacell-models.

The three LightningModules keep their Spotlight training hooks
(seg_aux, seg_consistency, mask_mode, target_bg_lowpass). Every hook
defaults off, the bit-identity tests guard that, and only Spotlight
leaves turn them on; two DynaCell-paper 2D checkpoints store
seg_aux_weight in their hparams, so a hook-free engine could not
resume them. mask_conditioning.py comes along because engine.py and
celldiff_wrapper.py import it.

The files that import the Spotlight leaf generator are trimmed in the
commits that follow.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build(dynacell): lock the dynacell eval, croissant and preprocess deps

Take dynacell-models' uv.lock verbatim, which this tree's pyprojects now
match (`uv lock --check` exits 0). Against the shared-packages lock it
moves cubic 0.7.0 -> 0.9.0a4 and adds dinov3 (pinned commit), etils,
ftfy, jsonpath-rw, mlcroissant, pandas-stubs, ply, rdflib, submitit and
termcolor. A fresh `uv lock` would instead float dinov3 to its current
HEAD and bump mlcroissant/pandas-stubs, so the verbatim lock keeps the
versions the paper evaluations ran with.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test: drop the requires_data marker registration

Its only user was main's generate_grouped_eval_configs_test.py, which
the DynaCell port replaced: the live-data tests there are now marked
slow. With strict markers, nothing needs the registration.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* ci(dynacell): run the release-path, provenance and watchdog suites

Restore the eight dynacell test files the shared-packages PR left out
of the test-dynacell-configs job, now that they exist:
test_evaluation_microssim_calibration, segmentation_test,
test_distribution, test_croissant, test_main_cli_routing,
test_evaluation_fid_provenance, provenance_test and
watch_stalled_jobs_test. The job now matches dynacell-models'.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs: add the runbook for jobs that hang while holding an allocation

Port the DynaCell job runbook of the root CLAUDE.md from dynacell-models
(6193b9c, bcbfbe6, c4c5ed0, 74fa18a #497, e1d75cf): how to tell a
finished-but-hung job from healthy compute by the CPU-time derivative,
the stall watchdog (tools/watch_stalled_jobs.py), per-family predict
and fit wall limits, and the forensics to capture before cancelling.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(dynacell): scope the grouped eval walk to registered model dirs

The grouped generator imported ARMS from the Spotlight-v2 leaf
generator, which is not part of this PR, to skip the Spotlight arms that
share the canonical prediction tree. Drop that import and the Spotlight
model names from _CODE_TO_PAPER, _DETERMINISTIC_MODELS and
_CELLDIFF_MODELS, keeping the four 2D seed replicates the 2D appendix
scores.

The shared tree still holds the arms' stores, and parse_zarr_name
raises on an unknown model token (e.g. the CellDiff "variant"
segaux_iterative). walk_predictions now skips any model dir outside
_REGISTERED_MODEL_DIRS before parsing and prints how many stores it
skipped. The test that pinned the Spotlight skip becomes one that pins
the unregistered-dir skip and its count.

evaluation/paths.py keeps its Spotlight names: without them
canonical_model_name would silently resolve an arm's run dir onto its
baseline.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(dynacell): drop the Spotlight leaf entries from the config guards

The Spotlight arm leaves are not part of this PR, so:
- test_benchmark_config_composition: drop the eight v1 *_spotlight
  train leaves, their iPSC predict leaves and data-hparam entries, the
  _UNMIGRATED_MODELS filter that only they needed, and the Spotlight
  members of _WIDE_GPU_TRAIN_LEAVES and _LONG_WALL_TRAIN_LEAVES (the
  2D pix2pix seed1 replicate keeps its wide-GPU opt-out);
- test_paths: drop the two tests that import the Spotlight leaf
  generator or depend on the grouped generator registering the arms.
  The canonical_model_name / resolve_model tests of the arm names stay,
  since evaluation/paths.py keeps those names.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(dynacell): describe the 2D seed1 leaves without the Spotlight generator

The 40 leaves of the four 2D baseline seed replicates were emitted by
the Spotlight-v2 leaf generator, which is not part of this PR, and
their headers told the reader to edit that generator. Say instead what
they are for: the second draw scored by the *__2d_focus2d_seed1
buckets. Comments only; the configs are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(dynacell): drop the Spotlight eval generator from the generator lists

evaluation/README.md and the hand-written-bucket guard's comment in
test_evaluate_compose.py named generate_spotlight_v2_eval_configs.py
among the generators that write grouped buckets; it is not part of this
PR.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): declare torchvision, which evaluation.metrics imports

viscy_utils.evaluation.metrics imports torchvision.ops at module level,
and losses.MixedLoss imports that module, but viscy-utils never listed
torchvision. Installs only worked because another workspace member
pulled it in. The new test-utils CI job syncs viscy-utils alone and
failed to collect six test modules with ModuleNotFoundError. The lock
gains only the dependency edge; no version moves.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* ci(dynacell): keep main's batch-launcher test and the writer tests out

Two corrections to the test-dynacell-configs job:
- main's tests/test_submit_benchmark_batch.py resolves leaves under
  /hpc and fails on a runner (5 PermissionErrors in CI). The
  dynacell-models version that runs there arrives with the DynaCell
  PR, which adds it back.
- the four viscy-utils test files (prediction writer, prediction
  metadata, monitor health, LR cooldown) now run in the test-utils job
  on both Python versions, so the dynacell job no longer repeats them
  or syncs viscy-utils' test group for them.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(viscy-utils): share the fit-callbacks lookup between CLI injectors

_inject_checkpoint_guardrails and _inject_optimizer_health_guard each
repeated the same walk from the config to fit's trainer.callbacks plus
the already-injected check; the `else config` branch of that walk could
never run behind the `subcommand != "fit"` return. Both now call
_fit_callbacks_without. _callback_field collapses its two identical
branches, and the monitored checkpoint's init_args are kept from the
search loop instead of being looked up twice. No behavior change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(viscy-utils): drop MonitorHealthCheck's unused check_finite flag

Nothing sets check_finite (the CLI injects monitor and patience only,
and no config or test passes it), so the callback always checks
finiteness; it now does so with math.isfinite instead of building a
tensor for one float.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(viscy-utils): stop describing fg masks as built from Otsu thresholds

generate_fg_masks does not read the stored otsu_threshold; it binarizes
a CLAHE-equalized, median-filtered volume. Its summary line and
VisCyTrainer.preprocess's compute_fg_masks entry still said "from Otsu
thresholds".

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* build: leave the hf_demo pytest path to the DynaCell PR

Nothing on this branch imports from applications/dynacell/examples/
hf_demo: its only test, build_demo_data_test.py, arrives with the
DynaCell PR, which adds the pythonpath entry back with it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): guard a fit that sets no callbacks

Trainer.callbacks defaults to None, and the CLI injectors returned
early unless they found a list, so a fit that set no callbacks ran
without OptimizerHealthCheck, the guard meant for every fit. The
shared lookup now gives such a fit an empty list to append to (None
and [] mean the same to Trainer). The checkpoint guardrails still
need a monitored ModelCheckpoint and remain a no-op there.

Reported by Copilot on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(viscy-utils): say what the AMP skip counter actually counts

The counter resets only when GradScaler grows its scale, so successful
optimizer steps can sit between the reductions it counts; the error
message claimed every step was skipped "in a row". Describe it as
reductions without growth in between, which is what it detects: the
backward overflowing faster than growth_interval clean steps can
restore the scale. Behavior unchanged.

Reported by Copilot on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): resolve viscy_utils.callbacks names lazily

VisCyCLI now injects MonitorHealthCheck and OptimizerHealthCheck into
every fit by class path through viscy_utils.callbacks, whose __init__
eagerly imported OnlineEvalCallback and with it scikit-learn, which
only the eval extra installs. A fit on a base install would therefore
fail before training. The package now resolves its exports lazily
(PEP 562), like viscy_utils itself. The new test blocks sklearn and
imports the health callbacks: the eager __init__ fails it with
ModuleNotFoundError, the lazy one passes.

Reported by Copilot on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(cytoland): install viscy-data from GitHub in the VCP tutorials

The tutorials install viscy-utils from the GitHub monorepo and let the
siblings resolve from PyPI. The writer they predict with now reads
SlidingWindowDataset.z_planes, which the PyPI viscy-data 0.1.0 does not
have (checked in its wheel), so predict would fail with AttributeError
in on_predict_start. Install viscy-data from GitHub too, in the .py
sources and their paired notebooks.

Found by /code-review on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): return an empty CLAHE result when the clip range is zero

foreground_mask_from_volume guards a constant volume, but a blank
volume with a few hot pixels (under 0.2% of voxels) has equal 0.1/99.9
percentiles too, and equalize_clahe then divided by zero and fed NaN to
CLAHE (RuntimeWarning: divide by zero, reproduced on a (20, 64, 64)
volume). equalize_clahe now returns zeros there, which the mask's
existing flat-histogram guard turns into an empty mask, as for a
constant volume. Volumes with any contrast are unchanged.

Found by /code-review on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): keep a step-cadence leaf's last.ckpt on its step trigger

The injected latest-weights ModelCheckpoint takes over save_last from
the monitored one but copied only every_n_epochs (defaulting to 1). A
leaf that checkpoints every N steps, typical for max_steps runs with
long epochs, would get last.ckpt only at epoch boundaries, so a
preempted run resumed from much older weights. The injected callback
now copies the monitored callback's every_n_train_steps or
train_time_interval when set; epoch-cadence leaves are unchanged. No
current leaf combines a monitored checkpoint with a step trigger.

Found by /code-review on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(dynacell): skip only PAPER_KEY-known arms in the grouped eval walk

The scope filter skipped every model dir outside the generator's
registry and only printed a count. That also silenced the loud
"unknown model code-name" a new, never-registered model used to raise,
so its stores would drop out of every paper bucket without an error.
Now the walk skips only dirs that paths.PAPER_KEY knows but this
generator does not register (the Spotlight arms on this branch);
anything unknown to both still raises in parse_zarr_name, as on
dynacell-models. Tests pin both sides.

Found by /code-review on #534.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(dynacell): say the instance-AP audit does not share the walk's arm skip

audit_prediction_coverage claimed to apply the same skip filters as
walk_predictions. It applies the same test-set filter, but on purpose
not the walk's skip of PAPER_KEY models the grouped generator leaves
unregistered: listing every unregistered model is what it is for.

Found by /code-review on #534.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(dynacell): carry the seg-consistency engine and focus-parity tests

test_seg_consistency_training.py exercises the seg_consistency hook of
DynacellUNet and DynacellGAN with a stand-in segmenter, and
test_seg_consistency_focus.py pins the loss's focus-plane rule to
dynacell.evaluation.focus. Both test code this PR ships and import
nothing from the Spotlight PR, so they belong here rather than with
the Spotlight arms. Verbatim from dynacell-models.

Found by /code-review on #534.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(viscy-utils): take last.ckpt from every monitored checkpoint

The guardrail turned save_last off only on the first monitored
ModelCheckpoint. A leaf with two monitored checkpoints (say loss and a
validation metric) kept a second last.ckpt writer next to the injected
latest-weights callback, so last.ckpt would alternate between "latest"
and "that monitor's cadence". All monitored callbacks now give it up;
the first still names the monitor, directory and cadence. No current
leaf has two monitored checkpoints.

Found by /code-review on #533.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(dynacell): one model-dir acceptance rule for parse and walk filter

_REGISTERED_MODEL_DIRS restated, as a separate set, the rule
_split_model_variant uses to accept a dir name, so the two could drift.
_match_model_variant is now that rule (None when the dir does not
parse); _split_model_variant wraps it and raises the same errors as
before, and the walk filter calls it. No behavior change.

Found by /code-review on #534.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(dynacell): share the foreign-model skip between the walk and the audit

walk_predictions skipped the Spotlight arms' dirs on the shared tree,
but the instance-AP coverage audit still parsed them, so on this branch
it reported 468 coverage errors on the live tree instead of
dynacell-models' 102, and its docstring claimed it applied the walk's
filters. Both now call one predicate, is_foreign_model_dir: a dir whose
model paths.PAPER_KEY knows but the grouped generator does not parse.
Measured on the live tree, the audit's error set is identical to
dynacell-models' (102 entries; the remainder are pre-existing). The new
audit test holds one registered, one foreign and one unknown store and
expects exactly the unknown one.

Found by /code-review on #534 (findings 1, 9, 10).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(dynacell): keep raising on campaign models the generator cannot parse

pix2pix2d_unetvit_wmatch{,_seed1} are in _CODE_TO_PAPER (and in
PAPER_KEY) but not in _DETERMINISTIC_MODELS, so their stores do not
parse. On dynacell-models the walk raises on them and the audit lists
them. The foreign-model predicate would have skipped them as another
campaign's arms. A model the generator itself lists in _CODE_TO_PAPER
is not foreign: the walk raises and the audit flags it, as on
dynacell-models. With this, the audit's error set on the live tree is
identical to dynacell-models' (102 entries, compared path by path).

Registering wmatch properly is a separate decision (its comment says
"evaluate in focused sibling buckets", and adding it to
_DETERMINISTIC_MODELS would fold it into the canonical buckets), so it
is left as on dynacell-models.

Found by /code-review on #534 (findings 2, 7).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(dynacell): stop pointing at the Spotlight leaf generator from this branch

evaluation/paths.py and tests/test_paths.py cited
tools/generate_spotlight_v2_leaves.py, which this branch does not
have, and test_paths claimed a pin to it that only exists with the
Spotlight v2 configs. Say where the generator lives instead. The
Spotlight model names stay: without them canonical_model_name would
resolve an arm onto its baseline. Comments only.

Found by /code-review on #534 (finding 8).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(dynacell): parse the warmup-matched pix2pix2d stores in the grouped walk

_CODE_TO_PAPER declares pix2pix2d_unetvit_wmatch{,_seed1} a parse-only
registration "same rule as fnet2d_voxelmatched", but unlike that model
they were never added to _DETERMINISTIC_MODELS. Their 56 stores on the
live tree therefore raised "unknown model code-name", which alone stops
walk_predictions. Register them like fnet2d_voxelmatched: they now parse,
and, as for every parse-only ablation, they are scored in their focused
sibling buckets (the committed *__2d_focus2d_pix2pix2d_wmatch leaves),
not by regenerating the canonical ones.

Fixes part of FOLLOWUPS #1 (live-tree walk crash).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(dynacell): skip the listed checkpoint re-predicts in the grouped walk

The live tree holds second-checkpoint predictions of registered models
beside their canonical stores, under train-set dirs paths.py does not
know: fnet3d_paper/ipsc_step{8400,11k}, fnet2d_voxelmatched/ipsc_step200k,
fcmae_vscyto{2d_scratch_samplematched,3d_scratch}/ipsc_final and
fcmae_vscyto3d_scratch/ipsc_repredict (34 stores). key_from_prediction_store
raises on them ("unknown train_set token 'ipsc_final'"), which stopped
walk_predictions and filled the instance-AP audit. They come from the
voxel-matched FNet-2D and sample-matched UNeXt2-2D ablations, are
scored only by hand-written experiment buckets, and no committed config
references them, so they are not canonical inputs. Each
(model dir, train-set dir) pair is now listed in _CHECKPOINT_REPREDICTS
with its campaign; the walk skips and counts them and the audit skips
them. A train-set dir that is neither canonical nor listed still raises.

Fixes part of FOLLOWUPS #1 (live-tree walk crash).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(dynacell): name the P1 segcons replicate dirs as foreign in the walk

nucleus/pix2pix2d_unetvit_segcons_{s1,ctrl_s1,x10} (12 stores, Spotlight
P1 follow-ups of 2026-10-19) sit in the canonical tree but are in
neither this generator's registry nor paths.PAPER_KEY, so the
PAPER_KEY-based foreign test cannot see them and they raised "unknown
model code-name". They are now listed by name in _FOREIGN_MODEL_DIRS
and skipped like the PAPER_KEY arms; the skip message no longer says
"PAPER_KEY". Any other unknown dir still raises. The new audit test
holds one store of each kind on the live tree (registered, checkpoint
re-predict, named foreign, wmatch, unknown train set) and expects only
the unknown one to be flagged.

Fixes the rest of FOLLOWUPS #1 (live-tree walk crash).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(dynacell): list checkpoint re-predicts per organelle

_CHECKPOINT_REPREDICTS matched (model dir, train-set dir) in every
organelle, so a future mito/fnet3d_paper/ipsc_step8400 store would have
been skipped as a known re-predict instead of raising, though only the
ER one exists. Entries are now (organelle, model, train set) triples
for exactly the stores on the live tree; a listed token in another
organelle raises "unknown train_set token" again (new test, fails
before this change). The audit test now uses the nucleus step-11k
store, since the step-8400 one is ER-only.

Found by /code-review on the FOLLOWUPS #1 fix.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor(dynacell): one skip rule for the grouped walk and the audit

walk_predictions and audit_prediction_coverage each called
is_foreign_model_dir and is_checkpoint_repredict inline, so a third skip
category added to one could silently miss the other and split them
again. Both now call skip_reason(zarr_path), which names the category;
the walk counts per category. No behavior change.

Found by /code-review on the FOLLOWUPS #1 fix.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(dynacell): guard the grouped registries against the walk crash class

Two guards on the hand-kept registries behind the live-tree crash:
- every _CODE_TO_PAPER model must parse. At a04b3eb this lists
  pix2pix2d_unetvit_wmatch{,_seed1}, the omission that stopped
  walk_predictions;
- the name-based skip lists stay disjoint from what the campaign
  registers: no _FOREIGN_MODEL_DIRS entry in _CODE_TO_PAPER or the
  parse tuples (it would be dropped silently, since the foreign check
  runs first), and no re-predict token among the canonical train sets.

Found by /code-review on the FOLLOWUPS #1 fix.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(dynacell): reflow the checkpoint re-predict comment

The organelle-scoping edit left a broken line wrap in the
_CHECKPOINT_REPREDICTS comment. Comment only.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(dynacell): keep the wmatch refit out of canonical and instance-AP leaves

Registering pix2pix2d_unetvit_wmatch{,_seed1} for parsing (97c30a1)
fixed the live-tree walk, but a default regeneration would then also
fold their 56 stores into the canonical {organelle}_{train_set}
buckets and the instance-AP leaves, against DON'T-FOLD: the committed
*__2d_focus2d_pix2pix2d_wmatch buckets are where they are scored.
They are now listed in _SIBLING_ONLY_MODELS: walk_predictions still
returns them, so sibling buckets can be built from it, but main()
buckets canonical_predictions() (the walk minus that list) and the
instance-AP in_scope() drops them. New tests: canonical_predictions
leaves them out, in_scope rejects them.

MEASURED on the live tree, both generators run the default way (no
--out-root) emit 0 wmatch blocks, and the canonical pool (628 stores)
is identical, store by store, to what the pre-fix rules parse, so a
default regeneration writes exactly what a04b3eb would have written
had its walk not crashed.

Found by /code-review on the FOLLOWUPS #1 fix (finding 2).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* docs(dynacell): retitle the paper and describe DynaCell as the broader framework

The camera-ready title is now "Evaluating Dynamic 3D Virtual Staining
of Live Cells with DynaCell": DynaCell names a broader framework for
image-based profiling of cell-state dynamics, and this paper's
evaluation framework is part of it. Update the title in CITATION.cff
(software and preferred citation), the README and both HF card BibTeX
entries, and the Croissant citeAs string; rewrite the CITATION.cff
abstract and the README tagline to describe the evaluation framework
as part of DynaCell.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015kuht3npKtb2o5E6h19xbA

* docs(dynacell): call the contribution the DynaCell evaluation suite

Mirror paper 1c85e84: "This evaluation framework is part of DynaCell, a
broader framework" repeated the noun. The paper's contribution is now
the DynaCell evaluation suite, and DynaCell stays the broader framework.
The CITATION.cff abstract also says "a set of baseline models", so
"suite" means one thing.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015kuht3npKtb2o5E6h19xbA

* docs(dynacell): retitle to The DynaCell Evaluation Suite

The paper's title is now "The DynaCell Evaluation Suite for Dynamic 3D
Virtual Staining of Live Cells" (Overleaf b73f25f). Updates
CITATION.cff, the README and HF card BibTeX, and the Croissant citeAs.
test_croissant.py: 17 passed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015kuht3npKtb2o5E6h19xbA

* chore(dynacell): delete finished one-shot tools

Each did a one-time job that is done, and nothing else imports them:
- build_migration_manifest, run_migration, restandardize_config_paths
  (+ tests): the canonical-path migration was applied 2026-07-13; a
  rebuilt manifest on 07-24 had nothing left to move;
- retokenize_deconv_predictions (+ test): the 24 a549__deconv /
  joint__legacy_deconvgt renames were applied 2026-07-22; paths.py
  still reads both tokens for the dirs kept on purpose;
- scan_nan_microssim_caches (+ test): one audit for the cubic 0.9.0a1
  NaN MicroMS3IM caches that #506 fixed;
- add_instance_ap_to_ablation_leaves: one codemod, already applied;
- smoke_joint_leaf: an April CPU smoke the config-composition tests
  now cover.
The provenance builders (cpdino seg masks, HEK stores, A549 dual fuse,
random-init checkpoints) stay. README entries and one bake_best_ckpt
docstring line follow. Full suite: 5251 passed, 0 failed (the 18
fewer are the two deleted tests/ files; the 14 colocated tool tests
were never collected by the root run).

Raised in review by edyoshikun.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MzpgExfYMKCgogfnhcSgDc

---------

Co-authored-by: Alexandr Kalinin <alxndrkalinin@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants