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

applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/data_overlays/fcmae_vscyto2d_fit.yml

@@ -0,0 +1,78 @@
+# FCMAE-2D (VSCyto2D-geometry FullyConvolutionalMAE) fit-time HCS data
+# hparams for the in-focus 2D benchmark track. Split out from
+# model_overlays/fcmae_vscyto2d_fit.yml so the model+trainer half there
+# stays composable by joint-dataset (BatchedConcatDataModule) leaves that
+# author their own data: block.
+#
+# Mirrors data_overlays/fcmae_vscyto3d_fit.yml but flattens all Z
+# augmentations to the cytoland 2D convention (z_window_size=1 reads the
+# offline `_focus.zarr` 5-plane slab stores; each window is one plane):
+#   - z_window_size: 20 -> 1
+#   - RandWeightedCropd spatial_size [20,600,600] -> [1,600,600]
+#   - BatchedRandAffined rotate_range [3.14,0,0] -> [0,0,0],
+#     scale_range Z [0.7,1.3] -> [1,1] (identity = no Z scale; kornia
+#     affine scale is an absolute factor, so 1.0 is the no-op, not 0)
+#   - BatchedCenterSpatialCropd roi_size [15,384,384] -> [1,384,384]
+#   - BatchedRandGaussianSmoothd sigma_z [0.25,0.75] -> [0,0]
+# XY augmentation policy (shear/scale/contrast/intensity/noise/smooth) is
+# kept identical to the 3D overlay.
+data:
+  init_args:
+    z_window_size: 1
+    # 16 (not the 3D overlay's 32) so a single_gpu fit fits the 48 GB pool
+    # (a40|l40s|a6000): bs=32 OOMs at ~47 GB there, bs=8 uses ~half, so 16
+    # keeps ~2x headroom and constraint:null stays valid. Effective batch is
+    # deliberately not forced to match the 3D 4-GPU global batch — the 2D
+    # nets are trained to their own best on cheap hardware (plan Phase-B).
+    batch_size: 16
+    num_workers: 4
+    yx_patch_size: [384, 384]
+    augmentations:
+      # CPU weighted crop on the 5-plane focus slab: spatial_size Z=1
+      # samples a single target-weighted plane per crop.
+      - class_path: viscy_transforms.RandWeightedCropd
+        init_args:
+          keys: [Phase3D, Structure]
+          w_key: Structure
+          spatial_size: [1, 600, 600]
+          num_samples: 4
+    gpu_augmentations:
+      - class_path: viscy_transforms.BatchedRandAffined
+        init_args:
+          keys: [source, target]
+          prob: 0.8
+          rotate_range: [0, 0, 0]
+          shear_range: [0.0, 0.05, 0.05]
+          scale_range: [[1, 1], [0.5, 1.5], [0.5, 1.5]]
+      - class_path: viscy_transforms.BatchedCenterSpatialCropd
+        init_args:
+          keys: [source, target]
+          roi_size: [1, 384, 384]
+      - class_path: viscy_transforms.BatchedRandAdjustContrastd
+        init_args:
+          keys: [source]
+          prob: 0.5
+          gamma: [0.8, 1.2]
+      - class_path: viscy_transforms.BatchedRandScaleIntensityd
+        init_args:
+          keys: [source]
+          prob: 0.5
+          factors: 0.5
+      - class_path: viscy_transforms.BatchedRandGaussianNoised
+        init_args:
+          keys: [source]
+          prob: 0.5
+          mean: 0.0
+          std: 0.3
+      - class_path: viscy_transforms.BatchedRandGaussianSmoothd
+        init_args:
+          keys: [source]
+          prob: 0.5
+          sigma_x: [0.25, 0.75]
+          sigma_y: [0.25, 0.75]
+          sigma_z: [0, 0]
+    val_gpu_augmentations:
+      - class_path: viscy_transforms.BatchedCenterSpatialCropd
+        init_args:
+          keys: [source, target]
+          roi_size: [1, 384, 384]

applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/data_overlays/fnet2d_fit.yml

@@ -0,0 +1,58 @@
+# FNet2D fit-time HCS data hparams for the in-focus 2D benchmark track.
+# Split out from model_overlays/fnet2d_fit.yml so the model+trainer half
+# there stays composable by joint-dataset (BatchedConcatDataModule)
+# leaves that author their own data: block.
+#
+# Mirrors data_overlays/fnet3d_paper_fit.yml (including the mean/std
+# Structure normalization and the 8-crops-per-FOV sampling that diverge
+# from the CellDiff/UNetViT conventions) but flattens Z to the in-focus
+# `_focus.zarr` 5-plane slab stores (z_window_size=1, each window one
+# plane):
+#   - z_window_size: 32 -> 1
+#   - RandWeightedCropd spatial_size [32,64,64] -> [1,64,64]
+#   - CenterSpatialCropd (val) roi_size [32,64,64] -> [1,64,64]
+# The YX flips are unchanged (spatial_axes 1=Y, 2=X are valid at Z=1).
+data:
+  init_args:
+    z_window_size: 1
+    batch_size: 48
+    num_workers: 8
+    yx_patch_size: [64, 64]
+    normalizations:
+      - class_path: viscy_transforms.NormalizeSampled
+        init_args:
+          keys: [Phase3D]
+          level: fov_statistics
+          subtrahend: mean
+          divisor: std
+      - class_path: viscy_transforms.NormalizeSampled
+        init_args:
+          keys: [Structure]
+          level: fov_statistics
+          subtrahend: mean
+          divisor: std
+    augmentations:
+      # CPU: 8 patches per FOV (amortizes zarr decompression).
+      # batch_size=48 -> DataLoader loads 6 FOVs, each yields 8 patches = 48.
+      - class_path: viscy_transforms.RandWeightedCropd
+        init_args:
+          keys: [Phase3D, Structure]
+          w_key: Structure
+          spatial_size: [1, 64, 64]
+          num_samples: 8
+    gpu_augmentations:
+      - class_path: viscy_transforms.BatchedRandFlipd
+        init_args:
+          keys: [source, target]
+          spatial_axes: [1]
+          prob: 0.5
+      - class_path: viscy_transforms.BatchedRandFlipd
+        init_args:
+          keys: [source, target]
+          spatial_axes: [2]
+          prob: 0.5
+    val_augmentations:
+      - class_path: viscy_transforms.CenterSpatialCropd
+        init_args:
+          keys: [Phase3D, Structure]
+          roi_size: [1, 64, 64]

applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/fcmae_vscyto2d_fit.yml

@@ -0,0 +1,57 @@
+# Shared FCMAE-2D (FullyConvolutionalMAE, pretraining=False) fit overlay
+# for the in-focus 2D benchmark track. Geometry is the published VSCyto2D
+# recipe (applications/cytoland/examples/configs/recipes/models/fcmae_2d.yml,
+# plan M5): in_stack_depth=1, stem_kernel_size=[1,2,2]. Used by both
+# fcmae_vscyto2d_scratch and fcmae_vscyto2d_pretrained leaves —
+# encoder_only + ckpt_path are set only in the pretrained leaf so init is
+# the only difference between the two (lr drops to 0.0002 there,
+# mirroring cytoland vscyto2d finetuning).
+#
+# Mirrors model_overlays/fcmae_vscyto3d_fit.yml but with 2D geometry and
+# single-GPU topology:
+#   - stem_kernel_size [5,4,4] -> [1,2,2], in_stack_depth 15 -> 1
+#   - topology ddp_4gpu -> single_gpu (2D nets are light; a Phase-B
+#     throughput decision — see plan §2d). No explicit DDPStrategy /
+#     find_unused_parameters / 3h-timeout block is needed at single-device
+#     (SingleDeviceStrategy has no DDP reducer, so FCMAE pretraining=False
+#     unused decoder/head paths are fine). If a later throughput probe
+#     switches this to DDP, re-add the fcmae_vscyto3d_fit.yml
+#     DDPStrategy(find_unused_parameters=true, timeout="3:00:00") block.
+#
+# HCS data hparams (bs=16, z=1, yx=384, flattened augs) live in
+# data_overlays/fcmae_vscyto2d_fit.yml; single-store train leaves compose
+# both, joint (BatchedConcatDataModule) leaves compose only this one and
+# author data: themselves.
+base:
+  - ../../../../../../recipes/trainer/fit.yml
+  - ../../../../../../recipes/topology/single_gpu.yml
+model:
+  class_path: dynacell.engine.DynacellUNet
+  init_args:
+    architecture: fcmae
+    model_config:
+      in_channels: 1
+      out_channels: 1
+      encoder_blocks: [3, 3, 9, 3]
+      encoder_drop_path_rate: 0.1
+      dims: [96, 192, 384, 768]
+      decoder_conv_blocks: 2
+      stem_kernel_size: [1, 2, 2]
+      in_stack_depth: 1
+      pretraining: false
+    loss_function:
+      class_path: viscy_utils.losses.MixedLoss
+      init_args:
+        l1_alpha: 0.5
+        l2_alpha: 0.0
+        ms_dssim_alpha: 0.5
+    lr: 0.0004
+    schedule: WarmupCosine
+    # Phase-B TODO: re-size warmup_steps to ~1 epoch once the focus-store
+    # FOV count is known (single-GPU bs=16 at z_window_size=1 emits ~5
+    # windows/FOV/T, so steps/epoch differ from the 3D 4-GPU baseline).
+    warmup_steps: 4000
+    warmup_multiplier: 1e-3
+trainer:
+  precision: 16-mixed
+  max_epochs: 200

applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/fcmae_vscyto2d_predict.yml

@@ -0,0 +1,49 @@
+# FCMAE-2D (FullyConvolutionalMAE, pretraining=False) predict overlay for
+# the in-focus 2D benchmark track. Mirrors fcmae_vscyto3d_predict.yml but
+# with 2D geometry.
+#
+# CRITICAL: predict rebuilds the model from THIS overlay's model_config
+# and then loads the trained weights (engine.py:308), so the full 2D
+# model_config MUST be re-declared here (in_stack_depth=1,
+# stem_kernel_size=[1,2,2]) and data.z_window_size MUST be 1 — otherwise
+# load_state_dict fails on a stem/conv shape mismatch against the 2D
+# checkpoint. Used by both fcmae_vscyto2d_pretrained and
+# fcmae_vscyto2d_scratch predict leaves — predict loads the full trained
+# checkpoint, so the pretrained-encoder warm-start path (encoder_only +
+# ckpt_path) from the fit leaf does NOT belong here.
+#
+# The 2D model predicts each z-plane of the full-Z 3D test store from its
+# own phase plane (z_window_size=1 -> one window per z; the writer keys
+# each output to its z-index, z_padding=0). 640x960 / 512x512 need no YX
+# pad (num_blocks=4 -> multiple of 16).
+base:
+  - ../../../../../../recipes/trainer/predict.yml
+  - ../../../../../../recipes/topology/single_gpu.yml
+model:
+  class_path: dynacell.engine.DynacellUNet
+  init_args:
+    architecture: fcmae
+    model_config:
+      in_channels: 1
+      out_channels: 1
+      encoder_blocks: [3, 3, 9, 3]
+      encoder_drop_path_rate: 0.1
+      dims: [96, 192, 384, 768]
+      decoder_conv_blocks: 2
+      stem_kernel_size: [1, 2, 2]
+      in_stack_depth: 1
+      pretraining: false
+    loss_function:
+      class_path: viscy_utils.losses.MixedLoss
+      init_args:
+        l1_alpha: 0.5
+        l2_alpha: 0.0
+        ms_dssim_alpha: 0.5
+    predict_method: full_image
+    predict_overlap: [4, 256, 256]
+data:
+  init_args:
+    z_window_size: 1
+    batch_size: 1
+    num_workers: 0
+    yx_patch_size: [512, 512]

applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/fnet2d_fit.yml

@@ -0,0 +1,25 @@
+# FNet2D fit overlay — model + trainer only — for the in-focus 2D
+# benchmark track. HCS data hparams (the mean/std Structure normalization
+# and 8-crops-per-FOV sampling that diverge from the CellDiff/UNetViT
+# conventions) live in data_overlays/fnet2d_fit.yml; single-store train
+# leaves compose both, joint (BatchedConcatDataModule) leaves compose only
+# this one and author data: themselves.
+#
+# Mirrors model_overlays/fnet3d_paper_fit.yml but binds the fnet2d model
+# recipe (in_stack_depth=1, downsample_z=false). FNet3D was already
+# single-GPU, so topology is unchanged. lr/schedule/loss/precision/
+# max_steps all match the 3D FNet paper baseline.
+base:
+  - ../../../../../../recipes/models/fnet2d.yml
+  - ../../../../../../recipes/trainer/fit.yml
+  - ../../../../../../recipes/topology/single_gpu.yml
+seed_everything: 0
+model:
+  init_args:
+    loss_function:
+      class_path: torch.nn.MSELoss
+    lr: 0.001
+    schedule: Constant
+trainer:
+  precision: 32-true
+  max_steps: 200000

applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/fnet2d_predict.yml

@@ -0,0 +1,26 @@
+# FNet2D predict overlay for the in-focus 2D benchmark track. Binds the
+# fnet2d model recipe + predict trainer recipe, then layers predict-time
+# model hparams and data-loader settings. Predict-time normalizations and
+# data_path are leaf-owned (leaf overrides target-inherited values to
+# match each organelle's test_cropped store).
+#
+# CRITICAL: predict rebuilds the model from the fnet2d recipe's
+# model_config (in_stack_depth=1, downsample_z=false) and then loads the
+# trained weights (engine.py:308); data.z_window_size MUST be 1 so the
+# datamodule emits one window per z-plane of the full-Z 3D test store,
+# matching the Z-preserving 2D checkpoint. A z_window_size mismatch would
+# feed a multi-plane window into a Z=1 network and error.
+base:
+  - ../../../../../../recipes/models/fnet2d.yml
+  - ../../../../../../recipes/trainer/predict.yml
+  - ../../../../../../recipes/topology/single_gpu.yml
+model:
+  init_args:
+    predict_method: full_image
+    predict_overlap: [4, 256, 256]
+data:
+  init_args:
+    z_window_size: 1
+    batch_size: 1
+    num_workers: 0
+    yx_patch_size: [512, 512]

applications/dynacell/configs/benchmarks/virtual_staining/er/fcmae_vscyto2d_pretrained/a549_mantis/train.yml

@@ -0,0 +1,66 @@
+# FCMAE-2D (FullyConvolutionalMAE, pretraining=False) with FCMAE-2D
+# pretrained encoder init on ER/SEC61B, in-focus 2D benchmark track,
+# A549 mantis pooled (mock + DENV + ZIKV). Paper key: VSCyto2D. Companion
+# to fcmae_vscyto2d_scratch — the two leaves are identical except this one
+# loads the 2D FCMAE encoder weights and drops lr to 0.0002 (mirroring
+# cytoland VSCyto2D finetuning).
+#
+# Reads the offline 5-plane in-focus slab store SEC61B_all_focus.zarr
+# (z_window_size=1 emits one in-focus window per plane). target_channel
+# "Structure" matches the FCMAE-2D data overlay's default augmentation
+# keys, so no augmentation override is needed. Mirror of the 3D
+# er/fcmae_vscyto3d_pretrained/a549_mantis fit leaf.
+base:
+  - ../../../_internal/shared/model/train_sets/a549_mantis.yml
+  - ../../../_internal/shared/model/targets/er_sec61b.yml
+  - ../../../_internal/shared/model/data_overlays/fcmae_vscyto2d_fit.yml
+  - ../../../_internal/shared/model/model_overlays/fcmae_vscyto2d_fit.yml
+  - ../../../_internal/shared/model/launcher_profiles/mode_fit.yml
+  - ../../../_internal/shared/model/launcher_profiles/hardware_gpu_any_long.yml
+  - ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml
+
+benchmark:
+  task: virtual_staining
+  organelle: er
+  train_set: a549_mantis
+  model_name: fcmae_vscyto2d_pretrained
+  experiment_id: er__a549_mantis__fcmae_vscyto2d_pretrained
+
+model:
+  init_args:
+    # Load only the encoder from the 2D FCMAE ckpt (400 ep on HEK + A549 +
+    # BJ5a + LIVECell + CTMC phase data); decoder/head stay at fresh init.
+    # lr 0.0002 overrides the overlay's 0.0004 for pretrained finetuning.
+    encoder_only: true
+    ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/fcmae-2d/pretrain-logs/hek-a549-bj5a-livecell-ctmc-400ep/checkpoints/last.ckpt
+    lr: 0.0002
+
+trainer:
+  logger:
+    init_args:
+      name: FCMAE_VSCyto2D_Pretrained_A549_SEC61B_ws8500
+      save_dir: /hpc/projects/comp.micro/virtual_staining/models/dynacell/a549/er/fcmae_vscyto2d_pretrained
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: step
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: loss/validate
+        filename: "epoch={epoch}-step={step}-loss={loss/validate:.3f}"
+        auto_insert_metric_name: false
+        every_n_epochs: 1
+        save_top_k: 5
+        save_last: true
+        dirpath: /hpc/projects/comp.micro/virtual_staining/models/dynacell/a549/er/fcmae_vscyto2d_pretrained/checkpoints
+
+data:
+  init_args:
+    # A549 pooled in-focus slab store + target_channel — no resolver in
+    # this train_set.
+    target_channel: Structure
+    data_path: /hpc/projects/virtual_staining/training/dynacell/a549/mantis/train/SEC61B_all_focus.zarr
+
+launcher:
+  job_name: FCMAE_VSCyto2D_Pretrained_A549_SEC61B_ws8500
+  run_root: /hpc/projects/comp.micro/virtual_staining/models/dynacell/a549/er/fcmae_vscyto2d_pretrained