Rescale and LOT correction#424
Open
Soorya19Pradeep wants to merge 34 commits into
Open
Conversation
edyoshikun
force-pushed
the
rescale-lot-correction
branch
from
June 29, 2026 17:59
a93d071 to
37b81ff
Compare
…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>
edyoshikun
force-pushed
the
rescale-lot-correction
branch
from
June 29, 2026 20:35
901d5b7 to
dc39073
Compare
Contributor
There was a problem hiding this comment.
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/TripletDatasetto support per-FOVz_rangeresolution viaz_extraction_window+focus_slicemetadata, plus optionalz_reductionand pixel-size aware rescaling. - Add LOT correction core functions + Pydantic config models +
dynaclrCLI 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.
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>
Witness mode produced only the standard metrics/ROC/F1 plots — nothing
showing HOW the pseudo-labels were chosen. build_witness_labels now
attaches the gating diagnostic to the returned AnnData (obs["witness_score"],
obs["witness_ref"] = control_well/perturbed_well/other, uns["witness_gating"]
with threshold, bandwidth, counts, and the pre-gating score/ref arrays).
_save_task_plots leads the {task}_summary.pdf with a per-marker gating page:
the witness-score histogram split by well-of-origin, the dropped dead-zone
band, the sign cut, and labeled/dropped counts. Verified on real SEC61B
embeddings — control/perturbed wells separate weakly, most labeled cells are
"other" near zero, which explains the modest AUROC at a glance.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two related extensions to the witness LC label source:
1. Filter-based references. WitnessLabelSource gains control_filter /
perturbed_filter (arbitrary obs filters) as an alternative to
control_wells / perturbed_wells. obs_filter_mask supports scalar (==),
list (isin), and range dicts ({lt,le,gt,ge}, one bound = half-line, two =
window), with a well/fov_name key routing to the well-prefix match. This
enables contrasts like early-vs-late timepoints or control-vs-perturbed-at-
late as weak-label sources. Validator: wells XOR filter per side, no mixing.
2. Annotation-scored eval. WitnessSettings.eval_annotations lets a witness run
(weak labels for training) be scored against real infection_state joined
from annotation CSVs when the embeddings obs lacks the column, instead of
the perturbation proxy. _join_eval_annotations does the per-experiment join.
Tests: obs_filter_mask forms + filter-based late-window integration (8 green).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tness label The leakage fix recomputed the metrics CSV against ground-truth infection_state but left the ROC + F1-over-time pages drawing from the trainer's witness-label val arrays — so the summary PDF showed a circular AUROC=1.000 while the metrics bar/CSV correctly showed ~0.80. _evaluate_witness_against_annotations now also returns the annotation-scored y_val / y_val_proba / classes / val_hours (aligned to the has-ground-truth subset), and the orchestrator swaps them into the plotting outputs so the whole PDF is consistent. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The witness axis measures how much a MARKER's embedding changes between the
references, so its meaning is marker-dependent: viral_sensor -> infection,
organelle markers (SEC61B/TOMM20/G3BP1) -> remodeling. Scoring every marker
against infection_state was wrong for the organelle markers.
WitnessSettings.marker_eval maps a marker to {eval_against, eval_class_map},
overriding the top-level target for that marker (absent markers fall back).
_resolve_witness_eval applies the override per run; eval_source in the summary
records the actual target used. Lets one run score viral_sensor vs
infection_state and SEC61B vs organelle_state.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the sign+dead-zone witness gate and its circular annotation grading
with a two-stage, annotation-format pipeline (Soorya Pradeep's recipe,
framework-ized). A witness->GMM label IS an annotation: its meaning is named by
the modality it is computed from (viral_sensor -> infection_state; organelle ->
organelle_remodeling_state), so it flows through the existing annotation training
path unchanged.
Stage A (new `dynaclr witness-gmm-labels`): witness score -> per-condition
2-component GMM -> writes an annotation file (named state column + real class
vocabulary). Negatives = all control-well cells; positives = perturbed cells
with GMM posterior >= threshold; unimodal (near-noise) markers skipped.
Stage B: the existing `run-linear-classifiers` (label_source: annotations),
untouched. Teacher/student is free — point Stage B at a different modality's
zarr (e.g. SEC61 labels -> train on phase), joined by cell key.
- New shared `viscy_utils.evaluation.witness_gmm.fit_gmm_labels` (+ tests).
- Rename witness_labels{,_test}.py -> witness_gmm_labels{,_test}.py; recipe ->
witness_gmm_labels_infectomics.yml.
- Remove WitnessSettings, _gate_scores, _evaluate_witness_against_annotations,
_resolve_witness_eval, _join_eval_annotations, the label_source: witness
branch, and the superseded witness_score_classifiers DAG + visuals.
- New DAG witness_gmm_classifiers.md; update evaluation.md + LC README.
- Validated: fit_gmm_labels reproduces Soorya's saved GMM on
2026_03_24_A549_SEC61_ZIKV (bandwidth 238~240, remod weight 0.655~0.649,
confident-pos 59.1%~58.4%).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Remove generated mock visualization PNG/PDF files from version control. Files remain on disk but are no longer tracked. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Collapse _get_device() (~20 lines + module cache) to a one-line _DEVICE constant; torch.cuda.is_available() is sufficient here. - Remove the unused public gaussian_rbf_kernel() NumPy wrapper; nothing imports it and the internal _rbf_kernel does all the work. - Hoist the duplicated _subsample() (identical in compute_mmd.py and witness_gmm_labels.py) into a single shared viscy_utils.evaluation.mmd.subsample. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
_read_pixel_size now raises on a zero/negative/NaN OME-Zarr X scale instead of letting reference/inference division produce inf/NaN that silently corrupts initial_yx_patch_size. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
_collate_norm_meta took timepoint keys from the first sample and indexed every other sample with them; a batch mixing FOVs whose zattrs expose different timepoint_statistics sets raised a cryptic KeyError. Validate the key sets up front and raise an actionable message instead. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The classic TripletDataModule path returned the full per-FOV norm_meta
from zattrs without selecting the sample's timepoint, so
NormalizeSampled(level='timepoint_statistics') received the nested
{tp_idx: {stat: Tensor}} layout and raised KeyError('mean'). Live
DynaCLR configs (DynaCLR-3D-BagOfChannels-v2, DynaCLR-2D-MIP-BagOfChannels,
DINOv3-temporal) use this level, but only via MultiExperimentDataModule,
which already pre-resolves — so the classic path was latently broken.
- Hoist the timepoint resolver out of SlidingWindowDataset into a shared
_utils._resolve_timepoint_norm_meta and reuse it in both datasets.
- Resolve the timepoint in TripletDataset._slice_patch using the row's t.
- Drop the now-dead nested timepoint_statistics branch in
_collate_norm_meta (all collated metas are flat after resolution). This
supersedes the mixed-timepoint guard added in 9b0a114.
- Add a triplet integration test that fails without the resolution and
point the existing sliding-window helper test at the shared function.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
viscy_data eagerly imports triplet.py at package root, so a module-level `from viscy_transforms import ...` made viscy-transforms a hard requirement of `import viscy_data` — forcing the working tree to add it to core dependencies even for HCS-only users. Move the two imports (BatchedZoomd, BatchedChannelWiseZReductiond) into the reference_pixel_size / z_reduction branches where they are used, so they load only when those features are enabled. viscy-transforms stays in the [triplet] extra; `import viscy_data` no longer requires it. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This was referenced Jul 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The following have been implemented in this branch:
TripletDataModulethat rescales the patches to a reference pixel and a uses the Z-MIP projection transform