Repository navigation
Rescale and LOT correction - #424
Open
Soorya19Pradeep wants to merge 113 commits into
Open
Soorya19Pradeep wants to merge 113 commits into
Soorya19Pradeep wants to merge 113 commits into
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>
edyoshikun
pushed a commit
that referenced
this pull request
Aug 7, 2026
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
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