Temporal straightening

.ed_planning/visreg/temporal-regularization/PLAN.md

@@ -0,0 +1,59 @@
+# Temporal regularization of DynaCLR embeddings — plan
+
+**Status (2026-07-27):** implementation, PR-review hardening, and local integration validation complete; remote CI remains.
+
+## Objective
+
+Keep NT-Xent for state separation, then add two losses on encoder features:
+
+- a shared next-state predictor, `L_pred`, to reward transition structure shared across cells;
+- a three-frame curvature loss, `L_curv`, to make distance along a track better reflect progression.
+
+The losses create temporal structure; they are not gated on straightness already being present in a frozen contrastive embedding.
+
+## Locked design
+
+- Emit fixed `K`-frame, `tau`-spaced sequences as row-major `B*K` tensors.
+- A valid sequence uses one exact `global_track_id` and marker. It stops at divisions and never switches siblings.
+- Apply temporal losses directly to encoder features; do not add a separate temporal head.
+- Stop-gradient the predictor target at `t+1`.
+- Keep temporal weights independently schedulable.
+- Default to track-consistent stochastic augmentation. Retain `independent` and `none` as explicit ablations.
+- Mask incomplete sequences while keeping every configured module in the distributed backward graph.
+
+## Milestones
+
+- [x] Sequence sampling, rectangular batching, and validity masks.
+- [x] Curvature loss and shared predictor.
+- [x] Lightning integration, schedules, and configuration.
+- [x] Local end-to-end smoke coverage.
+- [x] PR review hardening and regressions.
+- [x] Run the final combined local suite.
+- [ ] Confirm remote CI after push.
+- [ ] Compare biology and representation metrics against the NT-Xent baseline.
+
+## PR review resolutions
+
+| Priority | Risk | Resolution |
+| --- | --- | --- |
+| P1 | A lineage stencil could jump between siblings. | Match every frame by exact `global_track_id` and invalidate at division boundaries (`061d6e31`). |
+| P1 | Flattened `B*K` frames received independent random transforms. | Reuse one transform realization per track; expose `consistent`, `independent`, and `none` modes (`061d6e31`). |
+| P1 | An all-invalid rank left predictor parameters unused in DDP. | Forward the empty masked tensor through the predictor and backpropagate a graph-connected zero (`651d693c`). |
+| P2 | `positive_cell_source=self` skipped the lookup needed by sequence emission. | Build the lineage/timepoint lookup whenever sequences are enabled (`061d6e31`). |
+| P2 | Half-open HPI bins omitted a maximum on an exact boundary. | Share an edge builder that always adds a terminal bin (`f9426fab`). |
+| P3 | Zero, negative, or non-finite HPI widths were accepted. | Reject invalid widths during Pydantic validation (`f9426fab`). |
+
+## Acceptance checks
+
+- Exact-track and division-boundary sampler regressions pass.
+- Self-positive sequence emission passes.
+- Track-consistent flip regression passes.
+- Predictor parameters receive non-`None`, zero gradients for all-invalid batches.
+- HPI terminal-boundary and invalid-width regressions pass.
+- Dataset/datamodule, engine, witness-GMM, and final integration suites pass.
+
+Local combined validation passed on 2026-07-27. The two skipped inference-reproducibility tests require external HPC data and CUDA; there were no failures.
+
+## Follow-up experiment
+
+Train matched seeds for NT-Xent-only versus NT-Xent + `L_pred` + `L_curv`. Compare contrastive retrieval, collapse indicators, temporal prediction, per-track progression, and held-out biological separation before promoting the temporal objective to a default recipe.

.ed_planning/visreg/temporal-regularization/straightening-on-dynaclr.md

@@ -0,0 +1,63 @@
+# Straightening on DynaCLR — exploration
+
+**Status (2026-07-27):** promoted to implementation; review findings incorporated.
+
+## Question
+
+Can DynaCLR learn a common biological progression without erasing state organization?
+
+## Working hypothesis
+
+Straightening alone is a per-track geometric prior; it does not directly reward a transition shared across cells. The selected objective therefore combines:
+
+`L = L_NT-Xent + lambda_pred * L_pred + lambda_curv * L_curv`
+
+The shared predictor supplies the common-transition pressure, curvature regularization makes local motion easier to interpret, and NT-Xent resists constant-state collapse.
+
+## Implementation findings
+
+### Identity is stricter than lineage
+
+A lineage may contain a parent and multiple daughters. Sampling any same-lineage row at each timepoint can create a synthetic path that changes physical cells. Temporal stencils must match the anchor's exact `global_track_id`; a division makes the stencil invalid.
+
+### Augmentation is part of the temporal model
+
+Applying random transforms to flattened `B*K` frames independently injects artificial motion. The default must reuse one random realization across the `K` frames of each track while allowing different realizations across tracks. Independent and augmentation-free modes remain useful ablations.
+
+### Masked loss still has distributed semantics
+
+On a rank with no valid sequences, returning a zero connected only to encoder output leaves predictor parameters unused. Passing the empty input through the predictor produces zero gradients for its parameters and keeps DDP iteration state consistent.
+
+### Evaluation bins need explicit boundary semantics
+
+HPI loops use half-open intervals `[lo, hi)`. `np.arange(start, max + width, width)` still omits `max` when it is exactly a boundary. The edge builder must create one additional terminal edge, and widths must be finite and positive before arithmetic.
+
+## Decisions retained
+
+- Operate on encoder features, matching the representation being regularized.
+- Use three or more points for curvature; two points provide smoothing only.
+- Keep fixed-frame spacing initially for a small, auditable implementation.
+- Treat velocity-based phenotype splitting as a possible discovery signal, not automatically as representation damage.
+- Evaluate biology and collapse jointly; straightness by itself is not a success criterion.
+
+## Rejected shortcuts
+
+- Same-lineage sampling without exact track identity.
+- Independent stochastic transforms as the default temporal input.
+- Skipping configured modules on empty masked batches under DDP.
+- Using frozen-embedding straightness as a gate for whether training may induce temporal structure.
+
+## Evidence added by review hardening
+
+- Exact-track selection chooses one sibling consistently and invalidates parent-to-daughter stencils.
+- `self` positives can emit sequences without a missing lookup.
+- Repeated frames remain identical after track-consistent random flips.
+- Every predictor parameter receives a zero, non-`None` gradient on all-invalid batches.
+- Boundary-aligned maximum HPI values fall inside a bin; invalid widths fail at config construction.
+
+## Open empirical questions
+
+- Does the predictor learn biology rather than acquisition-time drift?
+- Which temporal weight schedule preserves contrastive retrieval best?
+- Is `K=3` sufficient, or do longer stencils improve robustness enough to justify their sampling cost?
+- Do consistent spatial transforms improve temporal metrics without weakening useful augmentation diversity?

applications/dynaclr/configs/training/debug/temporal-straightening-smoke.yml

@@ -0,0 +1,82 @@
+# Smoke test: predictor + straightening temporal terms wired into DynaCLR.
+# Fast-dev run on the 2-FOV test parquet to confirm the sequence sampler emits
+# (B*K, C, Z, Y, X), the engine encodes it in one forward, and both temporal
+# losses (loss/curv, loss/pred) are computed without NaN/shape/DDP errors.
+#
+# Run:
+#   uv run --no-sync python -m dynaclr.cli fit -c \
+#     applications/dynaclr/configs/training/debug/temporal-straightening-smoke.yml
+
+seed_everything: 42
+
+trainer:
+  accelerator: gpu
+  devices: 1
+  precision: bf16-mixed
+  fast_dev_run: 4
+  logger: false
+  enable_checkpointing: false
+  enable_model_summary: false
+  use_distributed_sampler: false
+
+model:
+  class_path: dynaclr.engine.ContrastiveModule
+  init_args:
+    encoder:
+      class_path: viscy_models.contrastive.ContrastiveEncoder
+      init_args:
+        backbone: convnext_tiny
+        in_channels: 1
+        embedding_dim: 768
+        in_stack_depth: 1
+        stem_kernel_size: [1, 4, 4]
+        stem_stride: [1, 4, 4]
+        projection_dim: 32
+    loss_function:
+      class_path: viscy_models.contrastive.loss.NTXentLoss
+      init_args:
+        temperature: 0.2
+    lr: 0.00002
+    example_input_array_shape: [1, 1, 1, 160, 160]
+    straightening_loss:
+      class_path: viscy_models.contrastive.loss.TemporalStraighteningLoss
+    predictor:
+      class_path: viscy_models.contrastive.predictor.Predictor
+      init_args:
+        dim: 768
+    lambda_curv: 0.01
+    lambda_pred: 0.1
+
+data:
+  class_path: dynaclr.data.datamodule.MultiExperimentDataModule
+  init_args:
+    cell_index_path: /hpc/projects/organelle_phenotyping/models/collections/DynaCLR-2D-MIP-BagOfChannels-v3.parquet
+    z_window: 1
+    yx_patch_size: [192, 192]
+    final_yx_patch_size: [160, 160]
+    channels_per_sample: 1
+    positive_cell_source: lookup
+    positive_match_columns: [lineage_id]
+    tau_range: [0.5, 2.0]
+    tau_decay_rate: 2.0
+    emit_sequence: true
+    sequence_length: 3
+    sequence_tau_frames: 1
+    stratify_by: [perturbation, marker]
+    split_ratio: 0.8
+    batch_size: 16
+    num_workers: 1
+    seed: 42
+    normalizations:
+      - class_path: viscy_transforms.NormalizeSampled
+        init_args:
+          keys: [channel_0]
+          level: timepoint_statistics
+          subtrahend: mean
+          divisor: std
+    augmentations:
+      - class_path: viscy_transforms.BatchedRandFlipd
+        init_args:
+          keys: [channel_0]
+          spatial_axes: [1, 2]
+          prob: 0.5