feat(dynacell): migrate eval metrics to torch-fidelity, add probes

[alxndrkalinin]

Summary

• Migrate dynacell's eval pipeline to torch-fidelity for FID/KID/PRC/MIND, adding KID std + non-parametric P/R/F1 bootstrap + MIND (kernel-free, motivated by the paper's Appendix I mode-collapse analysis).
• Add CP feature selection (vendored pycytominer variance + correlation pruning, BSD-3), FOV-stratified linear probes (real-vs-pred + cross-condition), and the dynacell.evaluation.cross_condition_probe CLI with per-plate MAD scaling to cancel illumination/exposure offsets.
• Add high-level eval-runner scripts (evaluate_local.sh, evaluate_batch.sh) plus submit_evaluation_{job,batch}.py Hydra-override staging tools; save-dir naming routes through a new save_paths.py that stays bit-identical to the paper aggregation scripts.

Stacks on dynacell-models (PR #404).

Test plan

• uv run pytest applications/dynacell/src/dynacell/evaluation/ applications/dynacell/tests/test_evaluation_metrics.py applications/dynacell/tests/test_evaluation_pipeline.py --deselect applications/dynacell/tests/test_evaluation_metrics.py::test_gain_and_offset_errors_are_not_scale_invariant → 41/41 pass (10 feature_metrics, 8 feature_select, 9 linear_probe, 13 evaluation_metrics, 1 evaluation_pipeline). The one deselected test is a pre-existing failure on dynacell-models (commit bfc61891 made nrmse scale-invariant).
• uvx ruff check clean on all touched files.
• submit_evaluation_job.py --dry-run on membrane/fnet3d_paper/a549_mantis/predict__a549_mantis_mock.yml produces canonical eval_fnet3d_a549trained_membrane_mock save_dir with correct .cmd + .yml provenance.
• eval_save_dir matches paper/scripts/compute_all_organelle_precision_recall.py:eval_dir_for on 5 spot-checked (organelle, model, train_set, plate) combos including iPSC-trained, A549-trained, and joint flavors.
• Real-data cross_condition_probe on three celldiff joint membrane A549 plates: CP AUROCs in [0.46, 0.71] (no longer 1.0 after per-plate MAD + select_features), DINOv3/DynaCLR at ceiling consistent with cross-condition morphology being learnable in deep embedding space.
• End-to-end uv run dynacell evaluate ... on a real predict zarr (requires DINOv3 + DynaCLR weights cached + segmentation; left for reviewer to drive once stacked PR lands).

Notes

• torch-fidelity is pinned via git URL to commit 5e211a9 because MIND is not in any tagged release yet.
• save_paths.py is deliberately duplicated with the paper repo's compute_all_organelle_precision_recall.py:eval_dir_for; both files carry a comment pointing at the other.
• Per-(FOV, timepoint) KID columns on ~50-cell A549 cohorts are noisy and flagged as exploratory.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR migrates dynacell's evaluation pipeline to torch-fidelity for FID/KID/PRC/MIND metrics, adds a CELL-DINO foundation model wrapper, vendors pycytominer-style CP feature selection, and introduces FOV-stratified linear probes plus a cross-condition probe CLI. It also ships high-level eval-runner shell scripts, Hydra-override staging tools, and a new save_paths.py that mirrors paper-aggregation conventions.

Changes:

• New feature-similarity stack: feature_metrics.py (torch-fidelity FID/KID/PRC/MIND + bootstrap stds), feature_select.py (variance + correlation pruning), linear_probe.py (MAD-scaled FOV-stratified AUROC), and cross_condition_probe.py CLI; pipeline.py rewired to compose deep-feature crops once and threadpool the per-prefix metrics.
• New CELL-DINO foundation backbone (viscy_models.foundation.CellDinoModel + vendored DINOv2 ViT), wired through utils.CellDinoFeatureExtractor, the eval cache, and the eval Hydra config.
• Eval-runner ergonomics: submit_evaluation_{job,batch}.py, evaluate_local.sh, evaluate_batch.sh, plus save_paths.py for canonical save-dir naming bit-identical to the paper repo.

Reviewed changes

Copilot reviewed 24 out of 25 changed files in this pull request and generated 15 comments.

Show a summary per file

File	Description
applications/dynacell/pyproject.toml	Adds scikit-learn, pinned torch-fidelity git dep
uv.lock	Lock entries for new eval deps
applications/dynacell/src/dynacell/evaluation/feature_metrics.py	Torch-fidelity backed FID/KID/PRC/MIND + cosine, w/ pairwise variant
applications/dynacell/src/dynacell/evaluation/feature_metrics_test.py	End-to-end metric tests
applications/dynacell/src/dynacell/evaluation/feature_select.py	Vendored variance + correlation pruning
applications/dynacell/src/dynacell/evaluation/feature_select_test.py	Pruning tests
applications/dynacell/src/dynacell/evaluation/linear_probe.py	MADScaler + FOV-stratified AUROC + paired_auroc + indistinguishability
applications/dynacell/src/dynacell/evaluation/linear_probe_test.py	Linear probe tests
applications/dynacell/src/dynacell/evaluation/cross_condition_probe.py	Cross-condition diagnostic CLI on saved embeddings
applications/dynacell/src/dynacell/evaluation/metrics.py	Splits per-cell crop construction from extractor invocation; drops legacy cp_pairwise/deep_pairwise/dataset helpers
applications/dynacell/src/dynacell/evaluation/utils.py	Removes scipy-based FID/KID; adds *_batch extractor methods + CellDinoFeatureExtractor
applications/dynacell/src/dynacell/evaluation/pipeline.py	Re-wired to use new feature stack, threadpool, CELL-DINO branch, mask sidecar
applications/dynacell/src/dynacell/evaluation/pipeline_cache.py	CELL-DINO cache plumbing + sha hash hoist
applications/dynacell/src/dynacell/evaluation/cache.py	Adds celldino FeatureKind + group helpers
applications/dynacell/src/dynacell/evaluation/_configs/eval.yaml	Adds feature_extractor.celldino block
applications/dynacell/src/dynacell/evaluation/save_paths.py	Canonical eval save-dir naming (paper-aligned, intentionally duplicated)
applications/dynacell/tools/submit_evaluation_job.py	Stage one dynacell evaluate invocation from a predict leaf
applications/dynacell/tools/submit_evaluation_batch.py	Chain N per-leaf invocations into one sbatch
applications/dynacell/tools/evaluate_local.sh	Local-host wrapper around submit_evaluation_job.py
applications/dynacell/tools/evaluate_batch.sh	sbatch wrapper around submit_evaluation_batch.py
applications/dynacell/tests/test_evaluation_pipeline.py	Stub modules for new evaluation submodules
applications/dynacell/tests/test_evaluation_metrics.py	Drops legacy pairwise tests, adds batched-extractor tests
packages/viscy-models/src/viscy_models/foundation/cell_dino.py	CELL-DINO wrapper (preprocess + channel-mean cls pooling)
packages/viscy-models/src/viscy_models/foundation/_dinov2_vit.py	Vendored eval-only DINOv2 ViT
packages/viscy-models/src/viscy_models/foundation/init.py	Re-export CellDinoModel

Comments suppressed due to low confidence (5)

packages/viscy-models/src/viscy_models/foundation/cell_dino.py:83

• nn.Module.load_state_dict(state_dict, strict=True) raises RuntimeError itself when there are any missing or unexpected keys, so the subsequent if missing or unexpected: block is unreachable in normal operation. Either pass strict=False so the manual check actually runs (and provides the more descriptive error message) or drop the check entirely.

        missing, unexpected = self.model.load_state_dict(state_dict, strict=True)
        if missing or unexpected:
            raise RuntimeError(f"CELL-DINO state_dict mismatch — missing={missing}, unexpected={unexpected}")

applications/dynacell/src/dynacell/evaluation/cross_condition_probe.py:124

• MAD normalization is applied twice on the same feature matrix: first per-plate here via MADScaler().fit_transform(...), then again inside paired_auroc/fov_stratified_auroc whose Pipeline includes MADScaler() as the first step (linear_probe.py line 131-143). The second pass operates on already MAD-normalized data, which is unlikely to harm correctness but does undo part of the per-plate centering and produces non-obvious feature scales for the logistic regression. If the per-plate MAD is the intended normalization (per the PR description's "per-plate MAD scaling to cancel illumination/exposure offsets"), consider exposing a way for fov_stratified_auroc to skip the inner scaler when the caller has already normalized — or document that double normalization is intentional.

    # Per-plate MAD normalization: cancels per-plate intensity offsets
    # (illumination, exposure, gain) that would otherwise dominate the
    # classifier — especially on CP regionprops, where raw intensity
    # columns make different plates trivially separable (AUROC = 1.0).
    x0_scaled = MADScaler().fit_transform(x0.astype(np.float64))
    x1_scaled = MADScaler().fit_transform(x1.astype(np.float64))
    # Tag FOV ids by condition so the two sides cannot collide.
    fov0_tagged = np.asarray([f"{c0}::{f}" for f in fov0])
    fov1_tagged = np.asarray([f"{c1}::{f}" for f in fov1])
    result = paired_auroc(x0_scaled, x1_scaled, fov0_tagged, fov1_tagged, n_splits=n_splits, rng_seed=rng_seed)

applications/dynacell/tools/evaluate_local.sh:171

• This printf … | xargs … > /dev/null line is dead code: xargs -I {} echo {} echoes each token to stdout which is then sent to /dev/null, and any parse errors from xargs are not checked (set -e would catch a non-zero exit, but xargs -I will not error on simple newline-separated inputs). The actual command parse happens in the subshell below via the while IFS= read -r loop. Consider removing this line.

  printf '%s\n' "$cmd_text" | xargs -d '\n' -I {} echo {} > /dev/null  # validate parse

packages/viscy-models/src/viscy_models/foundation/_dinov2_vit.py:302

• x.shape for a 4D image tensor is (B, C, H, W), so this destructuring binds w to the height and h to the width. The downstream interpolate_pos_encoding(x, w, h) is then called with positional args that are consistently swapped throughout this module (e.g. w0 = w // self.patch_size, h0 = h // self.patch_size). For square inputs (CELL-DINO is trained at 224×224) the output is identical, but on any non-square input the positional encoding grid would be transposed. This mirrors the original DINOv2 source which has the same naming bug; consider adding a clarifying comment so future readers don't mistake w/h for actual width/height.

    def prepare_tokens_with_masks(self, x: Tensor, masks: Optional[Tensor] = None) -> Tensor:
        _, _, w, h = x.shape
        x = self.patch_embed(x)
        if masks is not None:
            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
        x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
        x = x + self.interpolate_pos_encoding(x, w, h)

applications/dynacell/tools/submit_evaluation_job.py:156

• The plate label inference relies on hardcoded suffix matching against _mock / _denv / _zikv and the literal stem predict__ipsc_confocal. This is duplicated verbatim in submit_evaluation_batch.py (lines 144-154). Any new test plate added later will require both files to be updated in lockstep. Consider extracting the plate-from-leaf-stem logic into dynacell.evaluation.save_paths (alongside the related path naming) so there is a single source of truth.

    # Plate label for save_dir: ipsc | mock | denv | zikv.
    # Predict leaves are named predict__<test_set>[_<cond>].yml. Extract the
    # trailing condition for a549; the bare ipsc_confocal predict leaf is iPSC.
    leaf_stem = leaf_path.stem  # predict__a549_mantis_mock
    if leaf_stem == "predict__ipsc_confocal":
        test_plate = "ipsc"
    elif leaf_stem.endswith("_mock"):
        test_plate = "mock"
    elif leaf_stem.endswith("_denv"):
        test_plate = "denv"
    elif leaf_stem.endswith("_zikv"):
        test_plate = "zikv"
    else:
        raise SystemExit(
            f"cannot infer test plate from leaf filename: {leaf_path.name} "
            f"(expected predict__ipsc_confocal.yml or predict__*_{{mock,denv,zikv}}.yml)"
        )

---

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