Skip to content

feat(dynacell): 2D in-focus VS model track - #487

Merged
alxndrkalinin merged 13 commits into
dynacell-2dfrom
dynacell-2d-models
Jul 22, 2026
Merged

alxndrkalinin merged 13 commits into
dynacell-2dfrom
dynacell-2d-models

Conversation

@alxndrkalinin

Copy link
Copy Markdown
Collaborator

Adds the 2D in-focus virtual-staining model track to the DynaCell benchmark — models trained on the in-focus slice extracted from the 3D data, for a 2D-vs-3D comparison. Stacks on #484 (the eval-side 2D plumbing: FRC + exclude_fov_names).

Scope

3 models × 3 train_sets × 4 organelles = 36 fit leaves:

  • fcmae_vscyto2d_scratch (UNeXt2-2D, random init)
  • fcmae_vscyto2d_pretrained (VSCyto2D, FCMAE-pretrained encoder init, lr 2e-4)
  • fnet2d (FNet with downsample_z=false, Z-preserving)

2D geometry: in_stack_depth=1, stem_kernel_size=[1,2,2], z_window_size=1. In-focus data = offline 5-plane _focus.zarr slab stores (waveorder-phase focus on Phase3D), extracted from the campaign's 3D train stores.

Commits

  • 5d2671b8 expose downsample_z on Unet3d (Z-preserving FNet2D; +tests)
  • ce22258a offline focus-slab store extractor (extract_focus_slab_store.py +tests)
  • eb68a45d 2D config overlays (fcmae_vscyto2d + fnet2d, fit + predict)
  • 98398330 nucleus/a549 smoke fit+predict leaves (exemplar)
  • b7a1a454 fix: identity Z scale in FCMAE-2D affine (kornia scale factor 0 → singular-matrix crash)
  • f701de71 fix: single-GPU batch 32→16 to fit the 48 GB pool (keeps 2D fits off scarce Hopper)
  • 3cc012ab / eef10cbb / ff982afc the 36 fit leaves (scratch / pretrained / fnet2d)

Verification

All 36 leaves render clean via submit_benchmark_job.py --print-resolved-config (correct 2D geometry, _focus.zarr data paths, per-organelle target channels). The nucleus/a549 FCMAE-2D config trained end-to-end in a 1-epoch smoke fit (both overlay bugs above were found + fixed there).

Not in this PR

Phase E (training the 36 fits), Phase F/G (predict + register + re-eval — gated behind training + the campaign's Phase 10). Predict leaves (144) trail training.

🤖 Generated with Claude Code

https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

alxndrkalinin and others added 9 commits July 22, 2026 10:48
… use

Unet3d hardcoded downsample_z=True into UNet3DBase, so it could not run
Z-preserving at Z=1 (the base D%2**depth divisor check trips on D=1) and
rejected a downsample_z kwarg with TypeError. Expose downsample_z (default
True, preserving FNet3D) and forward it to the base. With downsample_z=False
the encoder uses stride (1,2,2) / Z-preserving ConvTranspose3d and the base
skips the D divisor check, enabling the FNet2D track to run the FNet
recursive U-Net at Z=1.

Regression tests: default keeps downsamples_z True; downsample_z=False runs
a (B,1,1,64,64) forward preserving Z; and skips the Z divisor check for a
non-divisible Z. 11/11 pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
Pre-extract <name>_focus.zarr for 2D in-focus VS training: per position/T,
estimate the in-focus plane on the phase channel by reusing the eval
estimator (dynacell.evaluation.focus.estimate_focus_plane + the
MIDBAND_FRACTIONS constant, no reimplementation, so train and eval land
on the same plane), then cut a 2*hw+1-plane slab (hw=2 -> 5 planes) of all
channels. A new clamp_shift_slab shifts the whole window inward at edge-focus
FOVs (vs the eval's clipping focus_slab_from_plane) so every position yields
a uniform rectangular depth for SlidingWindowDataset. omero/normalization and
custom zattrs are preserved; a focus_slab_extraction provenance zattr records
the params. pixel_size defaults to the store lateral spacing (--pixel-size
override). zarr v3 / OME-Zarr v0.5, context-managed I/O.

9 integration tests (synthetic OME-Zarr with injected focus planes: shape,
channel/zattr preservation, plane centering, clamp-shift edges, limit/overwrite,
missing-channel raise, CLI) pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
Overlays for the 2D in-focus VS model track (fit + predict), mirroring the
3D FCMAE/FNet overlays with the 2D deltas:
- recipes/models/fnet2d.yml: architecture FNet3D + model_config.downsample_z:
  false, in_stack_depth 1 (Z-preserving FNet run at Z=1 via the new Unet3d
  downsample_z passthrough).
- data_overlays fcmae_vscyto2d_fit / fnet2d_fit: z_window_size 1; flatten Z
  augmentations (drop Z rotate/scale/sigma_z + the [15,...]/[32,...] z-crop;
  keep XY); RandWeightedCrop / center-crop roi Z -> 1.
- model_overlays fcmae_vscyto2d_fit / fnet2d_fit: 2D model_config (FCMAE stem
  [1,2,2], in_stack_depth 1; FNet2D downsample_z false), single_gpu topology
  (drop the DDP 3h-timeout block), FCMAE lr 4e-4/200 epochs/16-mixed, FNet2D
  200k steps/32-true.
- model_overlays fcmae_vscyto2d_predict / fnet2d_predict: re-declare the full
  2D model_config + data.z_window_size 1 so predict rebuilds a matching model
  before load_state_dict (engine.py:308).

Render-verified via submit_benchmark_job.py --print-resolved-config.

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

Phase-B smoke pair for the 2D track: fcmae_vscyto2d_scratch, nucleus, a549.
Fit reads the offline in-focus slab store dual_nucl_memb_all_focus.zarr inline
(bypasses dataset_ref, as A549 leaves do), z_window_size 1, target_channel
Nuclei, RandWeightedCrop re-keyed to Nuclei. Predict runs the 2D model over
the unchanged 3D test store at z_window_size 1 -> full-Z prediction.zarr under
the canonical directory grammar. Predict target_channel Nuclei satisfies the
Nuclei_prediction eval channel contract. ckpt_path/warmup_steps are Phase-B/F
TODOs (render clean; not loaded at compose time).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
The FCMAE-2D fit data overlay set scale_range Z to [0, 0], intending
"no Z scaling". kornia's RandomAffine3D treats scale as an absolute
factor (1.0 = identity), so 0 zeroes diagonal element 3 of the 4x4
affine matrix; warp_affine3d's torch.linalg.inv then raises a singular-
matrix _LinAlgError on the first GPU-augmented training batch. Use the
identity factor [1, 1] and correct the mapping comment. Affects both
FCMAE-2D leaves (scratch + pretrained), which compose this overlay.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
The FCMAE-2D fit overlay ran single-GPU (2D nets are light) but kept the
3D overlay's batch_size: 32, which there was split across ddp_4gpu
(8/GPU). On one device 32 samples/step OOMs a 48 GB a40/a6000/l40s at
~47 GB, and the long GPU profile sets constraint: null, so any FCMAE-2D
fit could land there and die on the first step. Lower to 16 (a verified
bs=8 uses ~half the 48 GB, so 16 keeps ~2x headroom) so the fit stays on
the cheap pool instead of being forced onto scarce 80 GB Hopper GPUs
across all 36 planned fits. Effective batch is intentionally not matched
to the 3D 4-GPU global batch; the 2D nets are tuned on their own.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
Author the remaining 11 UNeXt2-2D (FCMAE, pretraining=False, random init)
fit leaves — 4 organelles x 3 train_sets, minus the pre-existing
nucleus/a549_mantis smoke-validated exemplar. Each mirrors the 3D
fcmae_vscyto3d_scratch leaf with 2D geometry: composes the
fcmae_vscyto2d fit overlays (in_stack_depth=1, stem [1,2,2],
z_window_size=1), single-GPU topology, and reads the offline 5-plane
in-focus _focus.zarr slab stores.

iPSC single-set leaves set benchmark.dataset_ref: null and inline
source_channel/target_channel/data_path, because the _compose_hook
resolver raises when a leaf co-declares data_path alongside a full
inherited dataset_ref. Joint leaves inline a BatchedConcatDataModule
over the ipsc + a549 focus stores with batch_size 4 (joint mode does
not divide by num_samples, so 4 x 4 = 16 samples/step matches the
single-set 16) and the Z-flattened augmentation policy.

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

Author the 12 VSCyto2D (FCMAE, pretraining=False, FCMAE-pretrained encoder
init) fit leaves — 4 organelles x 3 train_sets. Identical to the
fcmae_vscyto2d_scratch leaves except the model block loads the 2D FCMAE
encoder (encoder_only: true, ckpt from the fcmae-2d
hek-a549-bj5a-livecell-ctmc-400ep pretrain) and finetunes at lr 0.0002,
mirroring the cytoland VSCyto2D recipe. Same 2D geometry, focus stores,
iPSC dataset_ref: null handling, and joint BatchedConcatDataModule as the
scratch counterpart. er/mito keep the 3D source's _ws8500 wandb-name
suffix so 2D and 3D pretrained runs stay greppable together.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
Author the 12 FNet2D fit leaves — 4 organelles x 3 train_sets. Each
mirrors the 3D fnet3d_paper leaf but composes the fnet2d fit overlays
(FNet3D architecture with downsample_z: false + in_stack_depth 1 =
Z-preserving 2D net, z_window_size 1, [1,64,64] crops), keeps the
single-GPU topology, and reads the offline _focus.zarr slab stores.
Retains FNet's mean/std Structure normalization and 8-crops-per-FOV
sampling. iPSC single-set leaves use dataset_ref: null + inlined
channels/path (resolver co-declaration guard); joint leaves inline a
BatchedConcatDataModule over the ipsc + a549 focus stores. sbatch mem
drops from the 3D leaves' 512G to 128G since the 5-plane focus stores
are far smaller than the full-Z sources.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
@alxndrkalinin alxndrkalinin changed the title feat(dynacell): 2D in-focus VS model track (fcmae_vscyto2d + fnet2d, 36 fits) feat(dynacell): 2D in-focus VS model track Jul 22, 2026
@alxndrkalinin
alxndrkalinin requested a review from Copilot July 22, 2026 22:44

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.

Pull request overview

Adds the 2D in-focus virtual-staining model track to DynaCell benchmarks by (1) enabling a Z-preserving “2D-style” FNet variant via downsample_z=False, (2) introducing an offline focus-slab extractor to generate _focus.zarr training stores, and (3) adding the corresponding model recipes/overlays and benchmark fit/predict leaves for the new 2D track.

Changes:

  • Expose downsample_z on Unet3d and add tests validating Z-preserving forward for Z=1 and skipping the Z divisor check.
  • Add an offline OME-Zarr focus-slab extractor (+ integration tests) to generate uniform-depth in-focus _focus.zarr slab stores.
  • Add FNet2D + FCMAE-VSCyto2D overlays and a large set of benchmark config leaves (fit + exemplar predict) for the 2D in-focus track.

Reviewed changes

Copilot reviewed 48 out of 48 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
packages/viscy-models/tests/test_unet/test_unet3d.py Adds coverage for downsample_z default behavior and the Z-preserving Z=1 forward path.
packages/viscy-models/src/viscy_models/unet/unet3d.py Exposes downsample_z on Unet3d to support a Z-preserving (2D-style) configuration.
applications/dynacell/tools/extract_focus_slab_store.py New CLI/tool to generate _focus.zarr in-focus slab stores from 3D training stores.
applications/dynacell/tools/extract_focus_slab_store_test.py Integration tests for the focus-slab extractor on a synthetic OME-Zarr plate.
applications/dynacell/configs/recipes/models/fnet2d.yml Adds a model recipe for “FNet2D” (FNet3D with downsample_z: false, in_stack_depth: 1).
applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/fnet2d_fit.yml Introduces the shared fit-time model/trainer overlay for FNet2D track.
applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/fnet2d_predict.yml Introduces the shared predict-time model/trainer overlay for FNet2D track.
applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/data_overlays/fnet2d_fit.yml Adds shared fit-time datamodule settings for the 2D FNet2D track (Z flattened to 1).
applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/fcmae_vscyto2d_fit.yml Introduces shared FCMAE VSCyto2D fit overlay (2D geometry + single-GPU topology).
applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/fcmae_vscyto2d_predict.yml Introduces shared FCMAE VSCyto2D predict overlay (2D geometry + z_window_size=1).
applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/data_overlays/fcmae_vscyto2d_fit.yml Adds shared fit-time datamodule settings for FCMAE VSCyto2D (flattened Z augs, bs=16).
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet2d/ipsc_confocal/train.yml Adds nucleus FNet2D single-set fit leaf for iPSC confocal focus-slab store.
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet2d/a549_mantis/train.yml Adds nucleus FNet2D single-set fit leaf for A549 mantis focus-slab store.
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet2d/joint_ipsc_confocal_a549_mantis/train.yml Adds nucleus FNet2D joint fit leaf for pooled iPSC+A549 focus-slab stores.
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto2d_scratch/ipsc_confocal/train.yml Adds nucleus FCMAE-2D scratch single-set fit leaf (iPSC focus store + nucleus re-keying).
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto2d_scratch/a549_mantis/train.yml Adds nucleus FCMAE-2D scratch single-set fit leaf (A549 focus store + nucleus re-keying).
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto2d_scratch/joint_ipsc_confocal_a549_mantis/train.yml Adds nucleus FCMAE-2D scratch joint fit leaf (pooled iPSC+A549 focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto2d_scratch/a549_mantis/predict__a549_mantis_mock.yml Adds exemplar nucleus FCMAE-2D scratch predict leaf (A549 mock test set).
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto2d_pretrained/ipsc_confocal/train.yml Adds nucleus FCMAE-2D pretrained single-set fit leaf (encoder-only init + lr override).
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto2d_pretrained/a549_mantis/train.yml Adds nucleus FCMAE-2D pretrained single-set fit leaf (A549 focus store).
applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto2d_pretrained/joint_ipsc_confocal_a549_mantis/train.yml Adds nucleus FCMAE-2D pretrained joint fit leaf (pooled iPSC+A549 focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fnet2d/ipsc_confocal/train.yml Adds mito FNet2D single-set fit leaf (iPSC TOMM20 focus store).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fnet2d/a549_mantis/train.yml Adds mito FNet2D single-set fit leaf (A549 TOMM20 focus store).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fnet2d/joint_ipsc_confocal_a549_mantis/train.yml Adds mito FNet2D joint fit leaf (pooled iPSC+A549 TOMM20 focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fcmae_vscyto2d_scratch/ipsc_confocal/train.yml Adds mito FCMAE-2D scratch single-set fit leaf (iPSC TOMM20 focus store).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fcmae_vscyto2d_scratch/a549_mantis/train.yml Adds mito FCMAE-2D scratch single-set fit leaf (A549 TOMM20 focus store).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fcmae_vscyto2d_scratch/joint_ipsc_confocal_a549_mantis/train.yml Adds mito FCMAE-2D scratch joint fit leaf (pooled iPSC+A549 TOMM20 focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fcmae_vscyto2d_pretrained/ipsc_confocal/train.yml Adds mito FCMAE-2D pretrained single-set fit leaf (encoder-only init + lr override).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fcmae_vscyto2d_pretrained/a549_mantis/train.yml Adds mito FCMAE-2D pretrained single-set fit leaf (A549 TOMM20 focus store).
applications/dynacell/configs/benchmarks/virtual_staining/mito/fcmae_vscyto2d_pretrained/joint_ipsc_confocal_a549_mantis/train.yml Adds mito FCMAE-2D pretrained joint fit leaf (pooled iPSC+A549 TOMM20 focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fnet2d/ipsc_confocal/train.yml Adds membrane FNet2D single-set fit leaf (iPSC focus store + membrane re-keying).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fnet2d/a549_mantis/train.yml Adds membrane FNet2D single-set fit leaf (A549 dual focus store + membrane re-keying).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fnet2d/joint_ipsc_confocal_a549_mantis/train.yml Adds membrane FNet2D joint fit leaf (pooled iPSC+A549 dual focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fcmae_vscyto2d_scratch/ipsc_confocal/train.yml Adds membrane FCMAE-2D scratch single-set fit leaf (iPSC focus store + membrane re-keying).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fcmae_vscyto2d_scratch/a549_mantis/train.yml Adds membrane FCMAE-2D scratch single-set fit leaf (A549 dual focus store + membrane re-keying).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fcmae_vscyto2d_scratch/joint_ipsc_confocal_a549_mantis/train.yml Adds membrane FCMAE-2D scratch joint fit leaf (pooled iPSC+A549 dual focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fcmae_vscyto2d_pretrained/ipsc_confocal/train.yml Adds membrane FCMAE-2D pretrained single-set fit leaf (encoder-only init + lr override).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fcmae_vscyto2d_pretrained/a549_mantis/train.yml Adds membrane FCMAE-2D pretrained single-set fit leaf (A549 dual focus store).
applications/dynacell/configs/benchmarks/virtual_staining/membrane/fcmae_vscyto2d_pretrained/joint_ipsc_confocal_a549_mantis/train.yml Adds membrane FCMAE-2D pretrained joint fit leaf (pooled iPSC+A549 dual focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/er/fnet2d/ipsc_confocal/train.yml Adds ER FNet2D single-set fit leaf (iPSC SEC61B focus store).
applications/dynacell/configs/benchmarks/virtual_staining/er/fnet2d/a549_mantis/train.yml Adds ER FNet2D single-set fit leaf (A549 SEC61B focus store).
applications/dynacell/configs/benchmarks/virtual_staining/er/fnet2d/joint_ipsc_confocal_a549_mantis/train.yml Adds ER FNet2D joint fit leaf (pooled iPSC+A549 SEC61B focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/er/fcmae_vscyto2d_scratch/ipsc_confocal/train.yml Adds ER FCMAE-2D scratch single-set fit leaf (iPSC SEC61B focus store).
applications/dynacell/configs/benchmarks/virtual_staining/er/fcmae_vscyto2d_scratch/a549_mantis/train.yml Adds ER FCMAE-2D scratch single-set fit leaf (A549 SEC61B focus store).
applications/dynacell/configs/benchmarks/virtual_staining/er/fcmae_vscyto2d_scratch/joint_ipsc_confocal_a549_mantis/train.yml Adds ER FCMAE-2D scratch joint fit leaf (pooled iPSC+A549 SEC61B focus stores).
applications/dynacell/configs/benchmarks/virtual_staining/er/fcmae_vscyto2d_pretrained/ipsc_confocal/train.yml Adds ER FCMAE-2D pretrained single-set fit leaf (encoder-only init + lr override).
applications/dynacell/configs/benchmarks/virtual_staining/er/fcmae_vscyto2d_pretrained/a549_mantis/train.yml Adds ER FCMAE-2D pretrained single-set fit leaf (A549 SEC61B focus store).
applications/dynacell/configs/benchmarks/virtual_staining/er/fcmae_vscyto2d_pretrained/joint_ipsc_confocal_a549_mantis/train.yml Adds ER FCMAE-2D pretrained joint fit leaf (pooled iPSC+A549 SEC61B focus stores).
Comments suppressed due to low confidence (1)

applications/dynacell/tools/extract_focus_slab_store.py:155

  • _custom_zattrs() currently excludes only the "ome" key; if a store uses top-level omero/multiscales metadata (a common layout), those keys will be copied and may overwrite destination metadata. Update this filter to use the full managed-key set.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

alxndrkalinin and others added 4 commits July 22, 2026 15:56
The class docstring unconditionally claimed the model downsamples all
three spatial dimensions and that Z must be divisible by 2**depth. With
the downsample_z parameter, Z is preserved (stride (1,2,2)) when
downsample_z=False, so the Z-divisibility constraint applies only when
downsample_z=True. Also replace the hardcoded "D % 16" (correct only for
depth=4) with the depth-general "D % 2**depth". Docstring only; no
behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
The module docstring states the output is zarr v3 / OME-Zarr v0.5, but
the open_ome_zarr write call relied on iohub's default version. Pin
version="0.5" explicitly so the output format is deterministic regardless
of the iohub default, matching build_cpdino_seg_cleaned.py.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
The batch size in the sibling data_overlays/fcmae_vscyto2d_fit.yml was
lowered 32->16 for single-GPU OOM headroom, but two comments in the model
overlay still referenced bs=32. Update both to bs=16 to match the actual
batch_size the composed leaves use.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
An explicit `if x.ndim != 5` guard once lived in Unet3d.forward but was
dropped in refactor 3bad6cd when that method was removed in favor of the
shared UNet3DBase.forward, which never re-added it. A 4D input (e.g. a
squeezed Z=1 dim) then hits Conv3d and raises a confusing "expected N
channels" error instead of a clear message. The 2D track's Z=1 /
downsample_z=False path newly exercises this, so restore the guard and
document both raises. Add a regression test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg
@alxndrkalinin
alxndrkalinin merged commit 5431ca3 into dynacell-2d Jul 22, 2026
@alxndrkalinin
alxndrkalinin deleted the dynacell-2d-models branch July 22, 2026 23:06
alxndrkalinin added a commit that referenced this pull request Jul 23, 2026
…D evaluation (#484)

* feat(dynacell/eval): support 2D pixel metrics (FRC) alongside 3D (FSC)

Make the pixel-metric path dimension-flexible so the same code scores a
single in-focus 2D projection as well as the full 3D volume, following
cubic's convention of dispatching on input rank:

- ssim() accepts 2-D (H, W) or 3-D (D, H, W) input and derives
  spatial_dims from ndim, instead of hard-rejecting non-3D.
- compute_pixel_metrics() uses ring-based frc_resolution (FRC_Resolution)
  for 2D and keeps shell-based fsc_resolution (*_FSC_Resolution) for 3D,
  and slices spacing to the trailing spatial axes. PCC/NRMSE/PSNR/
  Spectral_PCC are shape-agnostic. The 3D path is byte-identical.

This is the metric-side enabler for adding in-focus 2D columns to the
benchmark without replacing the existing full-3D pixel metrics.

The test change also fixes a pre-existing stale stub harness: the cubic
stubs lacked glcm_features / cubic.scipy / cubic.skimage, so the guarded
top-level import fell back to all-None and every stubbed test failed at
HEAD. Adds those stubs plus 2D/3D dispatch, FRC-vs-FSC, and real-cubic
integration tests.

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

* feat(dynacell/eval): add io.exclude_fov_names to skip incomplete-prediction FOVs

evaluate_predictions can now drop specific FOVs (matched by full name
0/0/fov0011 or leaf fov0011) via io.exclude_fov_names, applied uniformly to
pred/gt/seg positions so the strict position-count alignment still holds.
Unlike limit_positions it does not set partial_walk, so deep-feature caches
self-heal normally. Eval-side analog of the predict-side exclude_fov_names
(--resume-predict): evaluate the finished FOVs of a partially-run predict.

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

* test(dynacell/eval): cover io.exclude_fov_names FOV exclusion

End-to-end test driving evaluate_predictions on the 3-position CPU fixture:
excluding one position by its full name (A/1/1) or its leaf (1) leaves
N_POSITIONS-1 positions' worth of rows and keeps the strict pred/gt
count-alignment. Mirrors the limit_positions integration test.

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

* feat(dynacell): 2D in-focus VS model track (#487)

* feat(viscy-models): expose downsample_z on Unet3d for Z-preserving 2D use

Unet3d hardcoded downsample_z=True into UNet3DBase, so it could not run
Z-preserving at Z=1 (the base D%2**depth divisor check trips on D=1) and
rejected a downsample_z kwarg with TypeError. Expose downsample_z (default
True, preserving FNet3D) and forward it to the base. With downsample_z=False
the encoder uses stride (1,2,2) / Z-preserving ConvTranspose3d and the base
skips the D divisor check, enabling the FNet2D track to run the FNet
recursive U-Net at Z=1.

Regression tests: default keeps downsamples_z True; downsample_z=False runs
a (B,1,1,64,64) forward preserving Z; and skips the Z divisor check for a
non-divisible Z. 11/11 pass.

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

* feat(dynacell): add offline focus-slab store extractor for the 2D track

Pre-extract <name>_focus.zarr for 2D in-focus VS training: per position/T,
estimate the in-focus plane on the phase channel by reusing the eval
estimator (dynacell.evaluation.focus.estimate_focus_plane + the
MIDBAND_FRACTIONS constant, no reimplementation, so train and eval land
on the same plane), then cut a 2*hw+1-plane slab (hw=2 -> 5 planes) of all
channels. A new clamp_shift_slab shifts the whole window inward at edge-focus
FOVs (vs the eval's clipping focus_slab_from_plane) so every position yields
a uniform rectangular depth for SlidingWindowDataset. omero/normalization and
custom zattrs are preserved; a focus_slab_extraction provenance zattr records
the params. pixel_size defaults to the store lateral spacing (--pixel-size
override). zarr v3 / OME-Zarr v0.5, context-managed I/O.

9 integration tests (synthetic OME-Zarr with injected focus planes: shape,
channel/zattr preservation, plane centering, clamp-shift edges, limit/overwrite,
missing-channel raise, CLI) pass.

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

* feat(dynacell): add 2D-track config overlays (fcmae_vscyto2d + fnet2d)

Overlays for the 2D in-focus VS model track (fit + predict), mirroring the
3D FCMAE/FNet overlays with the 2D deltas:
- recipes/models/fnet2d.yml: architecture FNet3D + model_config.downsample_z:
  false, in_stack_depth 1 (Z-preserving FNet run at Z=1 via the new Unet3d
  downsample_z passthrough).
- data_overlays fcmae_vscyto2d_fit / fnet2d_fit: z_window_size 1; flatten Z
  augmentations (drop Z rotate/scale/sigma_z + the [15,...]/[32,...] z-crop;
  keep XY); RandWeightedCrop / center-crop roi Z -> 1.
- model_overlays fcmae_vscyto2d_fit / fnet2d_fit: 2D model_config (FCMAE stem
  [1,2,2], in_stack_depth 1; FNet2D downsample_z false), single_gpu topology
  (drop the DDP 3h-timeout block), FCMAE lr 4e-4/200 epochs/16-mixed, FNet2D
  200k steps/32-true.
- model_overlays fcmae_vscyto2d_predict / fnet2d_predict: re-declare the full
  2D model_config + data.z_window_size 1 so predict rebuilds a matching model
  before load_state_dict (engine.py:308).

Render-verified via submit_benchmark_job.py --print-resolved-config.

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

* feat(dynacell): add 2D nucleus/a549 smoke fit+predict leaves (fcmae_vscyto2d)

Phase-B smoke pair for the 2D track: fcmae_vscyto2d_scratch, nucleus, a549.
Fit reads the offline in-focus slab store dual_nucl_memb_all_focus.zarr inline
(bypasses dataset_ref, as A549 leaves do), z_window_size 1, target_channel
Nuclei, RandWeightedCrop re-keyed to Nuclei. Predict runs the 2D model over
the unchanged 3D test store at z_window_size 1 -> full-Z prediction.zarr under
the canonical directory grammar. Predict target_channel Nuclei satisfies the
Nuclei_prediction eval channel contract. ckpt_path/warmup_steps are Phase-B/F
TODOs (render clean; not loaded at compose time).

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

* fix(dynacell): use identity Z scale in FCMAE-2D affine augmentation

The FCMAE-2D fit data overlay set scale_range Z to [0, 0], intending
"no Z scaling". kornia's RandomAffine3D treats scale as an absolute
factor (1.0 = identity), so 0 zeroes diagonal element 3 of the 4x4
affine matrix; warp_affine3d's torch.linalg.inv then raises a singular-
matrix _LinAlgError on the first GPU-augmented training batch. Use the
identity factor [1, 1] and correct the mapping comment. Affects both
FCMAE-2D leaves (scratch + pretrained), which compose this overlay.

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

* fix(dynacell): size FCMAE-2D single-GPU batch to fit the 48 GB pool

The FCMAE-2D fit overlay ran single-GPU (2D nets are light) but kept the
3D overlay's batch_size: 32, which there was split across ddp_4gpu
(8/GPU). On one device 32 samples/step OOMs a 48 GB a40/a6000/l40s at
~47 GB, and the long GPU profile sets constraint: null, so any FCMAE-2D
fit could land there and die on the first step. Lower to 16 (a verified
bs=8 uses ~half the 48 GB, so 16 keeps ~2x headroom) so the fit stays on
the cheap pool instead of being forced onto scarce 80 GB Hopper GPUs
across all 36 planned fits. Effective batch is intentionally not matched
to the 3D 4-GPU global batch; the 2D nets are tuned on their own.

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

* feat(dynacell): add fcmae_vscyto2d_scratch fit leaves for the 2D track

Author the remaining 11 UNeXt2-2D (FCMAE, pretraining=False, random init)
fit leaves — 4 organelles x 3 train_sets, minus the pre-existing
nucleus/a549_mantis smoke-validated exemplar. Each mirrors the 3D
fcmae_vscyto3d_scratch leaf with 2D geometry: composes the
fcmae_vscyto2d fit overlays (in_stack_depth=1, stem [1,2,2],
z_window_size=1), single-GPU topology, and reads the offline 5-plane
in-focus _focus.zarr slab stores.

iPSC single-set leaves set benchmark.dataset_ref: null and inline
source_channel/target_channel/data_path, because the _compose_hook
resolver raises when a leaf co-declares data_path alongside a full
inherited dataset_ref. Joint leaves inline a BatchedConcatDataModule
over the ipsc + a549 focus stores with batch_size 4 (joint mode does
not divide by num_samples, so 4 x 4 = 16 samples/step matches the
single-set 16) and the Z-flattened augmentation policy.

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

* feat(dynacell): add fcmae_vscyto2d_pretrained fit leaves for the 2D track

Author the 12 VSCyto2D (FCMAE, pretraining=False, FCMAE-pretrained encoder
init) fit leaves — 4 organelles x 3 train_sets. Identical to the
fcmae_vscyto2d_scratch leaves except the model block loads the 2D FCMAE
encoder (encoder_only: true, ckpt from the fcmae-2d
hek-a549-bj5a-livecell-ctmc-400ep pretrain) and finetunes at lr 0.0002,
mirroring the cytoland VSCyto2D recipe. Same 2D geometry, focus stores,
iPSC dataset_ref: null handling, and joint BatchedConcatDataModule as the
scratch counterpart. er/mito keep the 3D source's _ws8500 wandb-name
suffix so 2D and 3D pretrained runs stay greppable together.

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

* feat(dynacell): add fnet2d fit leaves for the 2D track

Author the 12 FNet2D fit leaves — 4 organelles x 3 train_sets. Each
mirrors the 3D fnet3d_paper leaf but composes the fnet2d fit overlays
(FNet3D architecture with downsample_z: false + in_stack_depth 1 =
Z-preserving 2D net, z_window_size 1, [1,64,64] crops), keeps the
single-GPU topology, and reads the offline _focus.zarr slab stores.
Retains FNet's mean/std Structure normalization and 8-crops-per-FOV
sampling. iPSC single-set leaves use dataset_ref: null + inlined
channels/path (resolver co-declaration guard); joint leaves inline a
BatchedConcatDataModule over the ipsc + a549 focus stores. sbatch mem
drops from the 3D leaves' 512G to 128G since the 5-plane focus stores
are far smaller than the full-Z sources.

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

* docs: correct Unet3d Z-divisibility docstring for downsample_z

The class docstring unconditionally claimed the model downsamples all
three spatial dimensions and that Z must be divisible by 2**depth. With
the downsample_z parameter, Z is preserved (stride (1,2,2)) when
downsample_z=False, so the Z-divisibility constraint applies only when
downsample_z=True. Also replace the hardcoded "D % 16" (correct only for
depth=4) with the depth-general "D % 2**depth". Docstring only; no
behavior change.

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

* fix(dynacell): pin OME-Zarr version=0.5 on focus-slab output store

The module docstring states the output is zarr v3 / OME-Zarr v0.5, but
the open_ome_zarr write call relied on iohub's default version. Pin
version="0.5" explicitly so the output format is deterministic regardless
of the iohub default, matching build_cpdino_seg_cleaned.py.

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

* docs: fix stale bs=32 comments in fcmae_vscyto2d model overlay

The batch size in the sibling data_overlays/fcmae_vscyto2d_fit.yml was
lowered 32->16 for single-GPU OOM headroom, but two comments in the model
overlay still referenced bs=32. Update both to bs=16 to match the actual
batch_size the composed leaves use.

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

* fix(viscy-models): restore 5D input guard in UNet3DBase.forward

An explicit `if x.ndim != 5` guard once lived in Unet3d.forward but was
dropped in refactor 3bad6cd when that method was removed in favor of the
shared UNet3DBase.forward, which never re-added it. A 4D input (e.g. a
squeezed Z=1 dim) then hits Conv3d and raises a confusing "expected N
channels" error instead of a clear message. The 2D track's Z=1 /
downsample_z=False path newly exercises this, so restore the guard and
document both raises. Add a regression test.

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

---------

Co-authored-by: Alexandr Kalinin <alxndrkalinin@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* refactor(dynacell/eval): share focus estimator defaults via focus constants

Hoist the mantis acquisition defaults (na_det=1.35, lambda_ill=0.450)
into module constants DEFAULT_NA_DET / DEFAULT_LAMBDA_ILL in
evaluation/focus.py, use them as the read_focus_compute_config defaults,
and import them into the focus-slab extractor instead of re-declaring the
literals. Single source of truth so the training slab and the eval slab
share the same in-focus plane (train/eval cannot silently diverge).

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

* refactor(dynacell): simplify focus-slab extractor internals

Tidy ExtractionSummary and the extraction loop without behavior change:
derive n_positions as a property of focus_planes (drop the redundant
field and its manual counter), build the summary once the pixel size is
resolved (drop the float("nan") placeholder + reassignment), inline the
single-use _split_position_name helper into name.split("/"), and unpack
the leading shape dims via data.shape[:3].

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

* refactor(dynacell/eval): simplify freq_spacing rank inference

compute_pixel_metrics always receives a (Z, Y, X) spacing tuple, so
reduce it to the input rank unconditionally with list(spacing)[-ndim:]
and drop the dead scalar-spacing guard. Also drop the redundant int()
around pred_xp.ndim (already int). Byte-identical for the 3D path.

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

* fix(dynacell): train FCMAE-2D in 32-true to stop fp16 NaN divergence

All 24 FCMAE_VSCyto2D fits (scratch + pretrained, every train_set ×
organelle) diverged to NaN loss under precision: 16-mixed — top-k
checkpoints frozen at epoch 0-15, final loss/validate=NaN despite
reaching epoch 199/200. FNet2D (32-true) and the 3D FCMAE (same
16-mixed + MixedLoss) both train fine, so the failure is specific to the
2D FCMAE under fp16.

Verified empirically (scratchpad/repro_fcmae2d_nan.py): MixedLoss /
ms_ssim_25d on a real median/iqr-normalized target with near-untrained
predictions is finite in both fp32 and fp16 autocast, so the NaN is not
the loss forward (nor zero-divisor normalization) but fp16 training
dynamics — activation/gradient overflow the GradScaler misses. 32-true
removes fp16 entirely; 2D nets are light so the cost is negligible.

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

* fix(viscy-utils): stop bf16 SSIM NaN that diverged 2D FCMAE training

_compute_ssim_and_cs_bf16 runs its mean / mean-of-squares convolutions in
bf16 regardless of trainer precision. On large-magnitude z-scored FCMAE
targets (a channel with a tiny iqr, e.g. nucleus iqr~1.8, normalizes to
values in the hundreds) the variance estimate mu_xx - mu_x**2
catastrophically cancels and rounds *negative*. At depth=1 (2D FCMAE) each
SSIM window has only 1*H_k*W_k samples, so the cancellation is far more
likely than at depth=15 (3D, where the extra Z samples keep the estimate
stable) — a near-flat window then collapses the SSIM denominator to ~0 and
yields intermittent NaN loss that trains fine for a while then diverges.
clamp=True in MixedLoss cannot repair it (clamp bounds magnitude, not NaN),
so the poisoned gradient corrupts every weight.

Fix, in the shared helper:
- clamp both diagonal variances to >=0 (a variance is non-negative by
  definition), flooring the CS denominator at c2;
- bound the covariance by Cauchy-Schwarz (|sigma_xy| <= sqrt(sigma_x*sigma_y)).
  Exact fp32 recovers cs=1 on a flat window because the equally-negative
  sigma_x/sigma_y/sigma_xy cancel in the ratio; clamping only the diagonal
  terms leaves an unbalanced negative sigma_xy and inflates |cs| past 10 on
  background windows (which dominate a VS image). The bound restores cs=1;
- eps-floor both denominators so a zero data_range (c1=c2=0) cannot yield 0/0.

On well-conditioned windows the bound is slack and the clamps are no-ops:
3D correlated-pair aggregate SSIM and gradient-flow cosine vs MONAI fp32 are
unchanged within existing tolerances.

Two deterministic, hardware-independent regression tests are added: zero
data_range stays finite (baseline NaN on every device), and near-flat
large-magnitude windows keep cs~=1 (the unbounded-covariance variant fails
with |cs-1| > 10). The negative-variance NaN itself is a bf16 rounding
knife-edge not unit-reproducible across GPU architectures; it is validated
end-to-end (the exact failing fit ran clean after the clamp).

* fix(dynacell): revert FCMAE-2D fit to 16-mixed now that the SSIM NaN is fixed

The 32-true switch (93990d1) was a wrong hypothesis: the 2D FCMAE NaN was
the bf16 SSIM denominator collapse in ms_ssim_25d, not fp16 training
dynamics. _compute_ssim_and_cs_bf16 runs bf16 convolutions regardless of
trainer precision, so 32-true never addressed the real cause. The preceding
metrics fix clamps the variances and bounds the covariance at the source, so
16-mixed is safe again — restoring the lighter, faster precision the 3D
FCMAE also uses.

---------

Co-authored-by: Alexandr Kalinin <alxndrkalinin@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <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