Skip to content

Rescale and LOT correction - #424

Open
Soorya19Pradeep wants to merge 113 commits into
mainfrom
rescale-lot-correction
Open

Soorya19Pradeep wants to merge 113 commits into
mainfrom
rescale-lot-correction

Conversation

@Soorya19Pradeep

@Soorya19Pradeep Soorya19Pradeep commented May 19, 2026 •

Copy link
Copy Markdown
Contributor

The following have been implemented in this branch:

  • TripletDataModule that rescales the patches to a reference pixel and a uses the Z-MIP projection transform
  • computation and use of the computed linear optical transfer function for batch correction between datasets.
  • a simple case of computation of Maximum Mean Discrepancy (MMD) with sample config.

@edyoshikun
edyoshikun force-pushed the rescale-lot-correction branch from a93d071 to 37b81ff Compare June 29, 2026 17:59
@edyoshikun
edyoshikun changed the base branch from modular-viscy-staging to main June 29, 2026 20:22
Soorya19Pradeep and others added 6 commits June 29, 2026 13:28
…staging

Scale-aware patch extraction (packages/viscy-data):
- Add _read_pixel_size() helper to read X pixel size from OME-Zarr metadata
- Add reference_pixel_size parameter to TripletDataModule: when set, computes
  initial_yx_patch_size from the pixel-size ratio so the same physical area is
  covered at inference time
- Replace BatchedRescaleYXd (removed per review) with existing BatchedZoomd
  using scale_factor=(1.0, final_y/initial_y, final_x/initial_x) and
  mode="bilinear" with antialias; import is lazy to avoid a hard dep

LOT batch correction (applications/dynaclr):
- Add dynaclr.evaluation.lot_correction submodule with core logic (fit, apply,
  save, load), Pydantic configs, and Click CLI entry points
- Register fit-lot-correction and apply-lot-correction in cli.py via LazyCommand
- Add pot and joblib to optional-dependencies.eval in pyproject.toml

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Plots n cell patches side-by-side:
- Left: raw patch at initial_yx_patch_size (larger physical area)
- Right: same patch bilinearly downscaled to final_yx_patch_size (model input)

Both columns share the same percentile contrast window so differences are
spatial. Physical scale (µm × µm) is shown in each panel title.

Usage:
  python visualize_triplet_rescaling.py       --data-path /path/to/data.zarr       --tracks-path /path/to/tracks       --source-channel Phase3D       --z-range 0 5       --final-yx-patch-size 224 224       --reference-pixel-size 0.325       --output rescaling_comparison.png

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…meta

_collate_norm_meta iterated every normalization level and called
torch.stack on each stat, but timepoint_statistics is nested
{timepoint: {stat: tensor}} rather than flat {stat: tensor}. Any zarr
carrying timepoint_statistics (alongside fov/dataset stats) crashed
batch collation in TripletDataModule with "expected Tensor as element 0
... but got dict". Stack within each timepoint sub-dict instead.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…etDataModule

Lets a 3D OME-Zarr feed a 2D model without materializing a separate MIP
dataset, and centers the extracted Z window on each FOV's focus plane.

- z_reduction ("mip"/"center"): collapse the extracted z_range to one
  slice via BatchedChannelWiseZReductiond. Label-free channels (resolved
  by parse_channel_name) take the center slice; others are max-projected.
  on_after_batch_transfer expects Z=1 when reduction is on.
- z_extraction_window/z_focus_offset/focus_channel: resolve a per-FOV
  focus-centered window from each position's
  focus_slice[ch].fov_statistics.z_focus_mean (fallback z_total//2), all
  windows the same width. z_range stays as an explicit override; exactly
  one of z_range / z_extraction_window must be given. Per-FOV windows are
  resolved at setup() and looked up per patch in the dataset.

Tests cover both reduction strategies (discriminating center vs MIP),
the normalize-then-reduce order, per-FOV focus resolution, and the
z_range/z_extraction_window XOR guard.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Documents the TripletDataModule predict path (zarr + tracking, not
parquet) and adds a runnable sample config demonstrating z_reduction +
reference_pixel_size to feed a 3D dataset to a 2D model.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Notebook-style script (no CLI) that loads a 3D OME-Zarr + tracking,
extracts a per-FOV focus-centered Z window, collapses it via z_reduction,
applies a random affine so anchor/positive diverge, and visualizes a
couple of batches to a PNG.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

This PR expands the Triplet-based data loading path to better support cross-dataset inference (pixel-size aware patch extraction + on-the-fly Z reduction + focus-centered per-FOV Z windows), and adds a LOT (Linear Optimal Transport) batch-correction pipeline with CLI entrypoints for embedding-space correction.

Changes:

  • Extend TripletDataModule/TripletDataset to support per-FOV z_range resolution via z_extraction_window + focus_slice metadata, plus optional z_reduction and pixel-size aware rescaling.
  • Add LOT correction core functions + Pydantic config models + dynaclr CLI commands for fitting/applying correction pipelines.
  • Add tests for the new Triplet Z-window/Z-reduction behavior and update docs/config examples for inference.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 7 comments.

Show a summary per file
File Description
packages/viscy-data/src/viscy_data/triplet.py Adds pixel-size rescaling, per-FOV focus-centered Z windows, and optional Z-reduction in the Triplet datamodule/dataset.
packages/viscy-data/tests/test_triplet.py Adds tests for z_extraction_window XOR validation, per-FOV focus-centered windows, and Z-reduction behavior/order.
packages/viscy-data/src/viscy_data/_utils.py Extends norm-meta collation to correctly stack nested timepoint_statistics.
applications/dynaclr/src/dynaclr/evaluation/lot_correction/lot_correction.py Implements LOT correction fit/apply + pipeline save/load helpers.
applications/dynaclr/src/dynaclr/evaluation/lot_correction/config.py Adds Pydantic models for fit/apply configs and filter specs.
applications/dynaclr/src/dynaclr/evaluation/lot_correction/fit_lot_correction.py CLI wrapper for fitting and saving a LOT pipeline from YAML config.
applications/dynaclr/src/dynaclr/evaluation/lot_correction/apply_lot_correction.py CLI wrapper for applying a saved LOT pipeline to an embedding zarr.
applications/dynaclr/src/dynaclr/evaluation/lot_correction/init.py Exposes LOT correction functions at the package level.
applications/dynaclr/src/dynaclr/cli.py Registers new fit-lot-correction and apply-lot-correction CLI subcommands.
applications/dynaclr/pyproject.toml Adds joblib + pot to eval extras for LOT correction dependencies.
applications/dynaclr/scripts/dataloader_inspection/visualize_triplet_rescaling.py Adds a visualization utility for the new pixel-size rescaling behavior.
applications/dynaclr/scripts/dataloader_inspection/triplet_dataloader_zprojection.py Adds an inspection notebook/script for focus-centered Z windows + Z-reduction.
applications/dynaclr/docs/DAGs/inference_triplet.md Documents Triplet inference path and new options (needs a small wording update for z-window semantics).
applications/dynaclr/configs/prediction/predict_triplet_2d_from_3d.yml Provides a sample prediction config demonstrating reference_pixel_size + z_reduction.

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

edyoshikun and others added 17 commits June 29, 2026 14:07
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
- Declare viscy-transforms in viscy-data's triplet extra and hoist the
  BatchedZoomd/BatchedChannelWiseZReductiond imports to module top; the
  inline imports were masking a missing runtime dependency that would
  ImportError for anyone using z_reduction or reference_pixel_size.
- Raise in _resolve_per_fov_z_ranges when a FOV's Z is smaller than
  z_extraction_window, instead of silently emitting a narrower window
  that breaks cross-FOV batch stacking.
- Remove a dead duplicate break in test_focus_centered_z_range.
- Correct the inference DAG doc: embeddings live in .X (the embedding_key
  array), mirrored to obsm["X_backbone"]/["X_projections"].

uv.lock intentionally omitted: the workspace glob currently entangles the
untracked applications/eet package, so a clean regen of the single
viscy-transforms edge is not possible until eet is committed or removed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the NumPy/scipy.cdist backend with PyTorch so the pooled RBF
kernel matrix is built once on the available device (CUDA if usable,
else CPU) and reused across all permutations. Device selection probes a
trivial kernel launch and falls back to CPU, so a present-but-unusable
GPU (e.g. compute capability older than the torch build) does not crash.

Public API is unchanged: median_heuristic/compute_mmd_unbiased return the
same bandwidth-convention values, and mmd_permutation_test still returns
(mmd2, p_value, null_distribution) so dynaclr's effect-size and activity
z-score computations keep working.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add witness_function to the MMD module: the per-point empirical witness
w(z) = mean_i k(z, x_i) - mean_j k(z, y_j), the RKHS direction along
which distributions P (X) and Q (Y) differ. Positive scores lean toward
X, negative toward Y. Reuses the GPU RBF-kernel helpers and bandwidth
(median-heuristic) convention already in the module, and chunks over
query points to bound the intermediate kernel matrices.

Tests cover distribution separation, X<->Y antisymmetry, and chunk-size
invariance.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Adds `label_source: witness` to the linear-classifier eval as a second
option to annotation-CSV labels. Per marker, cells are pooled across
experiments, control/perturbed references are built from per-experiment
wells, the MMD witness is fit and scored per cell, and scores are gated
(sign + dead-zone) into control/perturbed pseudo-labels. Everything
downstream (train_linear_classifier, publish, append-predictions, plots)
is unchanged — only the label source differs.

- evaluate_config.py: WitnessLabelSource, WitnessSettings, label_source
  switch on LinearClassifiersStepConfig (annotations stays the default).
- witness_labels.py: well->reference mask (path-component match so C/1
  != C/10), witness fit/score, dead-zone gating.
- orchestrated.py: branch label assembly into _annotation_run_specs /
  _witness_run_specs / _build_labeled_adata; training loop unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Five integration tests exercising the real code path:
- well-prefix match rejects C/1 vs C/10 confusion
- witness scores separate control/perturbed clusters
- dead-zone drops ambiguous near-origin cells
- empty AnnData when a reference group has no cells
- end-to-end run_linear_classifiers witness mode → metrics + summary PDF

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Mirrors linear_classifiers_infectomics.yml but with label_source: witness.
Per-experiment control/perturbed wells replace annotation CSVs; wells are
templates to edit per plate layout. One classifier per marker (G3BP1,
SEC61B, Phase3D, viral_sensor), dead_zone 0.1, track-level split.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- witness_score_classifiers.md: full DAG (why, step-by-step, gating rule,
  config, notes) plus embedded visuals.
- visuals/: three Graphviz flow diagrams (data flow, gating, pipeline) as
  .dot + png + pdf, and four mock example-output plots (witness-score
  histogram, per-marker metrics bar, ROC, F1-over-time) with a reproducible
  generator (mock_witness_plots.py). Mock plots are watermarked synthetic.
- Cross-links from evaluation.md and linear_classifiers/README.md so the
  new DAG is discoverable from the annotation path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The witness label is a deterministic function of the embedding, so
evaluating a witness-trained classifier against that same label is
circular — it reports a trivial ~1.0 (verified on real infectomics
embeddings: 1.000 vs witness label, 0.712 vs true infection_state).

WitnessSettings gains eval_against (default "infection_state") and
eval_class_map ({control: uninfected, perturbed: infected}). In witness
mode, val metrics are recomputed on the val split against the ground-truth
obs column via the class map; the summary records eval_source. When the
column is absent, metrics fall back to the witness label flagged as
eval_source="witness_label". Adds two tests (annotation eval + fallback).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Score the witness-trained classifier against infection_state (not the
self-referential witness label) with the control→uninfected /
perturbed→infected class map, matching the leakage fix.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- witness_score_classifiers.md: add "Evaluation: avoid the circularity
  trap" section with the real per-marker metrics table (scored vs
  infection_state), eval_against/eval_class_map in the config block, and
  step 7 in the flow; refresh the Notes caveat (metrics are honest now,
  check eval_source; note class imbalance).
- witness_dataflow.dot/png/pdf: add the EVALUATE-vs-annotations node.
- mock_witness_plots.py + bar/ROC png/pdf: use real run values (viral_sensor
  0.87, SEC61B 0.76, Phase3D 0.58, G3BP1 0.49) instead of invented all-high
  numbers; watermark clarifies bar/ROC are real, hist/F1 synthetic.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
edyoshikun pushed a commit that referenced this pull request Aug 7, 2026
edyoshikun and others added 28 commits August 12, 2026 15:15
Segment cells from the projected virtual-stain membrane with Cellpose
CP-DINO, assign tracked nuclei to cells, derive cytoplasm masks, and
measure per-cell fluorescence for track eligibility. Includes the
0.58 um recipes (April pilot and organelle_remodeling rollout), the
per-FOV Slurm array scripts, the review tools behind the pilot
diagnostics, and the recipe doc. The reader accepts OME-Zarr identity
transforms used by legacy tracking stores.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Fit the pooled pAL17 teacher used for pAL-anchored alignment: per-dataset
control median/MAD normalization, pooled PCA80, an RBF witness against
controls, and an all-cell joint control+ZIKV GMM with tied covariance and
0.8/0.2 posterior gates. Leave-one-dataset-out pAL->pAL and pAL->Phase
classifiers measure agreement with the teacher. The threshold sweep
writes predictions_by_threshold.parquet, the per-cell posterior table
that the onset alignment reads.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Align ZIKV organelle remodeling to each lineage's pAL response. Track
segments are grouped into parent-linked families within dataset and FOV,
descendants are collapsed by median pAL posterior, and tau=0 is the first
of three frames with posterior >=0.8 after three frames <=0.2. SEC61B,
G3BP1, and TOMM20 embeddings are centered on the smoothed exact-HPI
control median, scaled by the control MAD, and reduced to PCA80.
Null-standardized MMD then compares HPI and tau on the same onset-bearing
cohort in 30-minute bins.

quantify_pal_alignment_gain.py computes the spline-fit alignment metrics,
and bootstrap_pal_alignment_gain.py resamples whole families within
dataset, well, and condition.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Plot pAL posterior curves on the HPI and tau axes, onset support, and
five-bin organelle-divergence calls. The DENV module holds the shared
plotting functions, and the ZIKV script reuses them. The per-dataset
audit freezes normalization, PCA, bandwidth, and onsets, reconstructs the
pooled curves, and then compares each marker's datasets individually.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Render the HPI-versus-tau MMD figure, the per-marker lineage montages
sampled every 3 hours from tau=-9 to +9 h, and the shared organelle UMAP.
The assembler combines the ZIKV and DENV panels into the pseudotime paper
figure under paper_figure_pseudotime/.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- evaluation/temporal/forecast.py, config_forecast.py: multi-event,
  multi-horizon onset forecasting from causal embedding histories
  (TCN and RoPE transformer heads, FOV-grouped folds).
- evaluation/temporal/dynaformer.py, dynaformer_data.py,
  dynaformer_model.py, config_dynaformer.py: self-supervised RoPE
  temporal head over lineage paths on frozen embeddings, plus the
  baseline ladder and probes.
- evaluation/temporal/model.py: RotaryEmbedding,
  ContinuousRotaryEmbedding, RoPESelfAttention, RoPETransformerLayer,
  BidirectionalRoPETransformer (exported from temporal/__init__.py).
- evaluation/geometry.py: embedding geometry and probe diagnostics.
- cli.py: register `temporal-forecast` and `dynaformer` commands.
- Phase3D 2026 annotation recipes, tools/dynaformer.sbatch, and tests
  (test_dynaformer, test_temporal_forecast, test_geometry).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Legacy tracking label images hold nuclei with no track row, and a
concave nucleus can have its rounded centroid on background. Both
rejected whole frames as track_label_encoding_unverified, skipping
almost every 2024/2025 frame. An encoding is now valid when at least
one centroid hits its own label and none lands on a different label.
Untracked nuclei stay in the label image and are still removed from
the cytoplasm. Adds centroid_in_footprint and
nuclear_footprints_in_cell (all nuclei, tracked or not, per cell).

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

Renders each completed frame of a pilot attempt next to a reference
attempt (VS membrane and Phase3D with cell, shared-mask and nucleus
outlines) and writes per-frame and pooled mask metrics.

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

Feeds cached TOMM20 X_normalized_pca80 embeddings as per-class dumps to the
pinned OPS set classifier (czbiohub-sf/ops_monorepo@ab11b56, run unchanged
from an external copy). Model I classifies onset-aligned pseudotime bins
against bystander and control references with phase-matched pools; Model C
learns G1/S/G2/M from uninfected cells. Includes bag-size sweep, held-out
evaluation, per-cell scoring, baselines and montage rendering, plus the
attempt-13 crop helpers it loads (style-only lint fixes; AST unchanged
apart from import order, one unused assignment and one variable rename).

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

Onset lineages now need an onset from both the Phase->pAL TCN and the pAL17
sensor, and bystanders must be not-infected by both. Renames the two models
to the infection classifier and the cell-cycle classifier, adds the
sensor-verified config (attempt 26) and an attempt 25 vs 26 comparison
table, and extends the tests to the new lineage rule.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
TripletDataset can mark predict-time patches whose QC channels are all
zero or NaN with an is_blank index column, which EmbeddingWriter carries
into AnnData obs.

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

predict-triplet can now run Katamari foundation checkpoints, restrict
predictions to cells passing the phenotyping eligibility filter, and pass
every experiment channel to the blank-patch QC.

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

Representation export first writes obsm["X_control_mad"], then fits and
writes the pooled marker-specific obsm["X_normalized_pca80"]. Linear
classifiers forward the requested embedding key.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
phenotyping-batch plans one CP-DINO phenotyping job per dataset and
submits only explicitly selected experiments. audit-ai-readiness reports
each collection's storage location and flags stores older than their
latest source assembly. predict_batch gains the stale-source check,
explicit normalization channels, eligibility, and Katamari options.
The CLI also registers filter-embedding-observations and phenotype-cells.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…nto the model matrix

The matrix preflight now fails on stale AI-ready stores, predict skips
datasets whose normalized representation is already written, and an
optional phenotyping stage runs once per dataset ahead of predict.
matrix-eval prepares the pooled representation as a dependent job before
evaluations, and gains a temporal_forecast evaluation type.

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

The organelle_remodeling matrix gains DINOv3 ConvNeXt-Tiny, DINOv3 ViT-B,
and Katamari G10 ViT-L rows with their eval configs. The collection adds
the PCNA DENV/ZIKV plates, and the infectomics classifier recipe adds
uninfected pseudo-controls for the 2025 mock wells.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Planning docs live in the separate ed-planning repo, and .ed_planning is
already git-ignored; this one protocol was committed before the move.

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

run-matrix passes DYNACLR_ELIGIBILITY and positional args 19-24 (YX patch
size, Katamari input size, readout, teacher flag, EvolutionaryScale root,
Pixi executable). predict.sbatch now forwards them to predict-batch and
launches Katamari through Pixi from the EvolutionaryScale checkout.

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

This branch has not been deployed

No deployments
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.

3 participants