v4 update

TODO.md

@@ -0,0 +1,23 @@
+# Project TODO
+
+## Paper
+
+- [ ] **Learning curve analysis** — train on 25/50/75/100% of training data, plot AUC vs n.
+  Motivation: empirical evidence that the model is data-starved, which justifies the decision
+  not to pursue attention-gating (transformer) extensions to the NTowerHT bridge.
+  If the curve is still ascending at full data → supports the argument that a more expressive
+  architecture would overfit at this sample size. Generates a figure for the paper.
+
+- [ ] **GradCAM nasal-side analysis** — re-run GradCAM separately for OD and OS eyes rather
+  than aggregated. The current aggregation mirrors the two eyes against each other, washing out
+  any directional bias. Clinically, we would expect GradCAM attention offset from the disc center
+  to trend toward the nasal side (where RNFL loss presents earliest in glaucoma). If the model
+  has learned this, it would only be visible in per-side heatmaps — OD and OS are mirror images
+  so the nasal direction is opposite for each. This could be a strong interpretability result
+  for the paper if the bias is present.
+
+- [ ] **Quantify attention-gating as future work** — use the learning curve result + parameter
+  count ratio (Q/K/V projections over fusion_dim vs training n) to formally justify the choice.
+  Frame in paper as: "we identify cross-attention inside the NTowerHT bridge as a promising
+  extension, but our sample size (N≈400 training patients) is insufficient to avoid overfitting
+  a more expressive interaction layer" — cite the learning curve figure as evidence.

v3/PLAN_v4_hypertower.md

@@ -0,0 +1,179 @@
+# v4 HyperTower Planning Document
+
+## Goals
+Rebuild the orchestrator using `run_ntower_cv` as the architectural foundation, with four key improvements: JSON config, decoupled data sources, declarative tower lists, and a cross-tower communication protocol.
+
+---
+
+## 1. JSON Config (replace argparse)
+
+The orchestrator receives a single JSON config file. It has no hardcoded knowledge of what args individual towers or data modules need — it just forwards the relevant subtrees.
+
+```json
+{
+  "run_name": "v4/ensemble_fused",
+  "eval_mode": "binary",
+  "epochs": 30,
+  "fusion_epochs": 10,
+  "fold_seed": 100,
+  "seed": 1234,
+  "data": {
+    "module": "v4.papila.v4papila",
+    "args": {
+      "iop_corr_method": "ratio",
+      "iop_drop_raw": true,
+      "exclude_cols": ["Axial_Length"]
+    }
+  },
+  "towers": [
+    {
+      "name": "img",
+      "module": "v3.classes.image_towers",
+      "class": "ImageEncoder",
+      "args": { "backbone": "refugelike", "freeze_ratio": 0.5, "augment": true },
+      "warmup_epochs": 0
+    },
+    {
+      "name": "cd",
+      "module": "v3.classes.clinical_towers",
+      "class": "ClinicalEncoder",
+      "args": { "hidden_dim": 128 },
+      "warmup_epochs": 40
+    }
+  ],
+  "bridge": {
+    "mode": "embedding_mlp",
+    "fusion_dim": 256,
+    "hidden_dim": 256
+  },
+  "training": {
+    "lr": 1e-4,
+    "batch_size": 16,
+    "bcd_prob": 0.5,
+    "warmup_tower_epochs": 3,
+    "warmup_fused_epochs": 3
+  }
+}
+```
+
+The orchestrator loads this with `json.load`, then calls `importlib.import_module(cfg["data"]["module"]).build_data(cfg["data"]["args"])` and similarly instantiates towers. No argparse anywhere in the orchestrator.
+
+---
+
+## 2. Decoupled Data Sources
+
+`v3/classes/papila_builders.py` → clone to `v4/papila/v4papila.py`.
+
+Merge in the relevant logic from `papila_data.py` (preprocessing, feature typing, IOP correction, etc.) so `v4papila.py` is self-contained.
+
+Contract: every data module must expose:
+```python
+def build_data(args: dict) -> DataBundle:
+    ...
+```
+The orchestrator calls `build_data` and gets back a `DataBundle`. It knows nothing else about the data source. Future modules (e.g. `v4/eyepacs/eyepacs_data.py`) just implement the same function.
+
+---
+
+## 3. Declarative Tower List
+
+Towers are loaded from the `"towers"` list in the JSON and stored as an ordered dict keyed by `name`. The orchestrator never imports a tower class directly.
+
+```python
+towers = {}
+for t_cfg in cfg["towers"]:
+    mod = importlib.import_module(t_cfg["module"])
+    cls = getattr(mod, t_cfg["class"])
+    # some tower constructors need data (e.g. ClinicalEncoder needs feature_dim)
+    # pass data as an optional kwarg; tower ignores it if not needed
+    towers[t_cfg["name"]] = cls(data=data, **t_cfg["args"])
+```
+
+Tower-specific training metadata (warmup epochs, batch key) lives entirely in the JSON, not in the orchestrator.
+
+---
+
+## 4. Cross-Tower Communication: `early_pass` Protocol
+
+**Problem:** GeometryTower needs to precompute segmentation maps from images, then inject them into other towers' sample dicts before loaders are built. This is currently done imperatively in the orchestrator.
+
+**Proposed solution: `early_pass` connector interface**
+
+Each tower optionally implements:
+```python
+class TowerBase:
+    def early_pass(self, context: EarlyPassContext) -> None:
+        """Called once per fold before loaders are built.
+        Tower can read from / write to shared context."""
+        pass
+```
+
+`EarlyPassContext` is a shared mutable object passed to all towers in order:
+```python
+@dataclass
+class EarlyPassContext:
+    eye_train: list[dict]
+    bilat_train: list[dict]
+    bilat_val: list[dict]
+    bilat_test: list[dict]
+    image_preprocessor: object
+    image_cache: object
+    device: torch.device
+    store: dict = field(default_factory=dict)  # cross-tower key-value store
+```
+
+Example: GeometryTower's `early_pass` computes seg maps and injects them into the sample dicts directly (modifying `eye_train` etc. in place), exactly as it does today — but now the orchestrator just calls:
+```python
+for tower in towers.values():
+    tower.early_pass(context)
+```
+
+The cross-talk case the user described (img_tower outputs geometry → cd_tower reads it) uses `context.store`:
+```python
+# ImageTower.early_pass:
+context.store["geometry_maps"] = self._compute_geometry(context)
+
+# ClinicalTower.early_pass:
+geo = context.store.get("geometry_maps")
+if geo is not None:
+    self._inject_geometry(context, geo)
+```
+
+Tower ordering in the JSON list determines execution order, so dependencies are declared implicitly. If a tower has no `early_pass`, the default no-op in `TowerBase` is used.
+
+**Alternative considered:** explicit dependency graph / DAG execution. Rejected for now — JSON ordering is simpler and sufficient for current needs. Can revisit if cross-tower dependencies become non-linear.
+
+---
+
+## 5. File Layout
+
+```
+v4/
+  hypertower/
+    v4_hypertower.py      # orchestrator (no argparse, no tower imports)
+    split_manager.py      # copy/adapt from v3 (or just import)
+  papila/
+    v4papila.py           # merged papila_builders + papila_data
+  configs/
+    ensemble_fused.json   # example config
+```
+
+Existing `v3/classes/` tower implementations are reused directly — no duplication needed since they're importable by the JSON `"module"` field.
+
+---
+
+## 6. Open Questions / Decisions Needed
+
+- **DataBundle API**: Does `build_data` need to return anything beyond the current `DataBundle`? Or should `DataBundle` grow a `profile` factory method?
+- **Per-tower batch_key convention**: Currently `EYE_KEY_MAP = {"img": "image_1", "cd": "matrix_1"}` is hardcoded. Should this be declared in the tower JSON config or inferred from tower type?
+- **cd_warmup loader**: Slot stripping (`if k != "image_1"`) is currently img-tower-aware. Under the new design, each tower should declare which slots it needs for warmup vs full training, so the orchestrator can build the right loader without knowing about `image_1`.
+- **Geometry injection today vs `early_pass`**: Geometry currently mutates sample dicts; `early_pass` formalizes this. Needs a migration plan for existing GeometryTower.
+
+---
+
+## Implementation Order (once ntower_cv is validated)
+
+1. Write `v4papila.py` (merge papila_builders + papila_data, expose `build_data(args)`)
+2. Add `early_pass(context)` no-op to `TowerBase`; implement in `GeometryTower`
+3. Write `v4_hypertower.py` orchestrator using JSON config + importlib tower loading
+4. Port one config (ensemble_fused) end-to-end and compare outputs against ntower_cv

v3/classes/__init__.py

@@ -21,14 +21,6 @@ from .slot_dataset import SlotDataset, slot_collate
 from .papila_data import PapilaData
 from .papila_builders import build_papila_data
 from .data_bundle import DataBundle
-from .dataset import ClinicalDataset
-from .config_builder import (
-    ConfigAssembly,
-    assemble_config,
-    load_config,
-    resolve_imports,
-)
-from .filters import RegexFilter, ColumnFilter, apply_regex_filters, apply_column_filters
 from .transforms import (
     ImageTransformConfig,
     backbone_transform_config,
@@ -43,11 +35,35 @@ from .transforms import (
     TRANSFORM_REGISTRY,
     build_transform_chain,
 )
-from .model_builder import V2ModelBundle, build_model_bundle
+from .towerbase import TowerBase, build_backbone, train_towers_epoch, collect_probs_towers
-from .towers import ImageTower, ClinicalTower, SiameseImageTower, build_backbone
+from .image_towers import ImageEncoder, SiameseImageTower, ImageTower
-from .bridges import Bridge, VoteBridge
+from .clinical_towers import ClinicalEncoder, ClinicalDataTower
-from .models import SingleEyeHT, BilateralHT
+from .geometry_towers import GeometryTower
-from .v2_hypertower import V2HyperTower, V2ModeComparisonOps, V2ModeComparator
+from .hypertower_models import (
+    SingleEyeHT,
+    BilateralHT,
+    SiameseHT,
+    FusedEnsembleHT,
+    LogitMLPEnsembleHT,
+    EmbeddingMLPEnsembleHT,
+    NTowerHT,
+    NLateralHT,
+    MonoTowerHT,
+    train_single_epoch,
+    train_bilateral_epoch,
+    train_siamese_epoch,
+    train_fusion_epoch,
+    train_ntower_epoch,
+    train_mono_epoch,
+    collect_probs_classic,
+    collect_probs_ensemble,
+    collect_probs_bilateral,
+    collect_probs_siamese,
+    collect_probs_ntower,
+    collect_probs_mono,
+    V2ModeComparisonOps,
+)
+from .bridges import Bridge, HTClassifier, HyperBridge, VoteBridge
 from .hypertower_logger import HypertowerLogger
 
 __all__ = [
@@ -66,18 +82,9 @@ __all__ = [
     "PapilaData",
     "build_papila_data",
     "DataBundle",
-    "ClinicalDataset",
     "SlotLoaderFactory",
     "SlotDataset",
     "slot_collate",
-    "ConfigAssembly",
-    "assemble_config",
-    "load_config",
-    "resolve_imports",
-    "RegexFilter",
-    "ColumnFilter",
-    "apply_regex_filters",
-    "apply_column_filters",
     "ImageTransformConfig",
     "backbone_transform_config",
     "build_backbone_transform",
@@ -90,18 +97,41 @@ __all__ = [
     "UnetMaskProvider",
     "TRANSFORM_REGISTRY",
     "build_transform_chain",
-    "V2ModelBundle",
+    "TowerBase",
-    "build_model_bundle",
-    "ImageTower",
-    "ClinicalTower",
-    "SiameseImageTower",
     "build_backbone",
-    "Bridge",
+    "train_towers_epoch",
-    "VoteBridge",
+    "collect_probs_towers",
+    "ImageEncoder",
+    "SiameseImageTower",
+    "ImageTower",
+    "ClinicalEncoder",
+    "ClinicalDataTower",
+    "GeometryTower",
     "SingleEyeHT",
     "BilateralHT",
-    "V2HyperTower",
+    "SiameseHT",
+    "FusedEnsembleHT",
+    "LogitMLPEnsembleHT",
+    "EmbeddingMLPEnsembleHT",
+    "NTowerHT",
+    "NLateralHT",
+    "MonoTowerHT",
+    "train_single_epoch",
+    "train_bilateral_epoch",
+    "train_siamese_epoch",
+    "train_fusion_epoch",
+    "train_ntower_epoch",
+    "train_mono_epoch",
+    "collect_probs_classic",
+    "collect_probs_ensemble",
+    "collect_probs_bilateral",
+    "collect_probs_siamese",
+    "collect_probs_ntower",
+    "collect_probs_mono",
+    "Bridge",
+    "HTClassifier",
+    "HyperBridge",
+    "VoteBridge",
     "V2ModeComparisonOps",
-    "V2ModeComparator",
     "HypertowerLogger",
 ]

v3/classes/bridges.py

@@ -1,19 +1,64 @@
 from __future__ import annotations
 
+
 import torch
 import torch.nn as nn
 
 from v3.classes.SE_attention import SEBlock, SEGateLogger
 
 
+# ---------------------------------------------------------------------------
+# HTClassifier — standalone classification head
+# ---------------------------------------------------------------------------
+
+class HTClassifier(nn.Module):
+    """Minimal classification head: ReLU → Dropout → Linear(in_dim → num_classes).
+
+    Used as the output stage of Bridge, HyperBridge, and any vehicle that needs
+    a reusable, identifiable classifier type.
+    """
+
+    def __init__(self, in_dim: int, num_classes: int, dropout: float = 0.5):
+        super().__init__()
+        self.head = nn.Sequential(
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(in_dim, num_classes),
+        )
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        return self.head(z)
+
+
+# ---------------------------------------------------------------------------
+# Bridge — N-tower fusion
+# ---------------------------------------------------------------------------
+
 class Bridge(nn.Module):
+    """
+    N-tower fusion bridge.
+
+    Takes a list of tower embeddings, projects each to a common ``fusion_dim``,
+    element-wise multiplies all projections, optionally applies an SE gate, then
+    classifies the fused representation via an HTClassifier.
+
+    Each tower also gets an auxiliary classification head (used for BCD training).
+
+    Construction
+    ------------
+    ``tower_dims`` is an ordered list of embedding dimensionalities — one entry per
+    embedding slot that will be passed to ``fuse()`` or ``forward()``.
+
+    Tower slots are accessed by index: ``W[i]``, ``ln[i]``, ``aux_heads[i]``.
+    The bridge has no knowledge of what modality each slot carries.
+    """
+
     def __init__(
         self,
-        img_dim,
+        tower_dims: list[int],
-        meta_dim,
+        num_classes: int,
-        num_classes,
+        fusion_dim: int = 256,
-        fusion_dim=256,
+        mode: str = "fused",
-        mode="fused",
         dropout: float = 0.5,
         use_se: bool = True,
         se_reduction: int = 16,
@@ -22,29 +67,30 @@ class Bridge(nn.Module):
         super().__init__()
         self.mode = mode
         self.use_se = use_se
+        self.tower_dims = list(tower_dims)
 
-        # project towers to equal width
+        # Per-tower projection heads: each projects dim_i → fusion_dim
-        self.W_img = nn.Linear(img_dim, fusion_dim)
+        self.W = nn.ModuleList([nn.Linear(d, fusion_dim) for d in tower_dims])
-        self.W_md = nn.Linear(meta_dim, fusion_dim)
+        self.ln = nn.ModuleList(
+            [nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
+             for _ in tower_dims]
+        )
 
-        # optional: layernorm before SE
+        # Per-tower auxiliary classifiers (for BCD training)
-        self.ln_img = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
+        self.aux_heads = nn.ModuleList([nn.Linear(d, num_classes) for d in tower_dims])
-        self.ln_md = nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
 
         # SE gate on the fused vector
         self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
         self.se_log = SEGateLogger(enabled=use_se, track_channels=False, dim=fusion_dim)
 
-        # heads
+        # Fused classifier head
-        self.classifier_fused = nn.Sequential(
+        self.classifier_fused = HTClassifier(fusion_dim, num_classes, dropout)
-            nn.ReLU(),
-            nn.Dropout(dropout),
-            nn.Linear(fusion_dim, num_classes),
-        )
-        self.classifier_img = nn.Linear(img_dim, num_classes)
-        self.classifier_cd = nn.Linear(meta_dim, num_classes)
 
-    def reset_se_stats(self):
+    # ------------------------------------------------------------------
+    # SE helpers
+    # ------------------------------------------------------------------
+
+    def reset_se_stats(self) -> None:
         """Call at epoch start."""
         if getattr(self, "se_log", None):
             self.se_log.reset()
@@ -55,41 +101,205 @@ class Bridge(nn.Module):
             return self.se_log.get(reset=reset)
         return None
 
-    def _compute_fused(self, img_feats, md_feats):
+    # ------------------------------------------------------------------
-        """Return z_fused embedding (before classifier_fused). Used by encode() and forward()."""
+    # Core fusion
-        hi = self.ln_img(self.W_img(img_feats))
+    # ------------------------------------------------------------------
-        hm = self.ln_md(self.W_md(md_feats))
+
-        fused = hi * hm
+    def _compute_fused(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
+        """Return z_fused embedding (before classifier_fused)."""
+        assert len(embeddings) == len(self.W), (
+            f"Bridge expects {len(self.W)} embeddings, got {len(embeddings)}"
+        )
+        h = self.ln[0](self.W[0](embeddings[0]))
+        for i in range(1, len(embeddings)):
+            h = h * self.ln[i](self.W[i](embeddings[i]))
         if self.se is not None:
-            fused, gates = self.se(fused)
+            h, gates = self.se(h)
             if self.se_log.enabled:
                 self.se_log.accumulate(gates)
-        return fused
+        return h
 
-    def encode(self, img_feats, md_feats) -> torch.Tensor:
+    # ------------------------------------------------------------------
-        """Return z_fused embedding without applying the classifier head."""
+    # N-tower API
-        assert self.mode == "fused", "encode() only valid in fused mode"
+    # ------------------------------------------------------------------
-        return self._compute_fused(img_feats, md_feats)
 
-    def forward(self, img_feats, md_feats):
+    def fuse(
-        out_img = None if self.mode == "clinical_only" else self.classifier_img(img_feats)
+        self, embeddings: list[torch.Tensor]
-        out_md = None if self.mode == "image_only" else self.classifier_cd(md_feats)
+    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        """
+        N-tower forward pass.
 
-        if self.mode == "fused":
+        Parameters
-            fused = self._compute_fused(img_feats, md_feats)
+        ----------
-            out_f = self.classifier_fused(fused)
+        embeddings : list of Tensor — one per tower slot (same order as tower_dims).
-            return out_f, out_img, out_md
-        # if ablation modes:
-        if self.mode == "image_only":
-            return out_img, out_img, None
-        if self.mode == "clinical_only":
-            return out_md, None, out_md
 
+        Returns
+        -------
+        logits_fused : Tensor  [B, num_classes]
+        aux_logits   : list of Tensor  — one per tower slot, each [B, num_classes]
+        """
+        z_fused = self._compute_fused(embeddings)
+        logits_fused = self.classifier_fused(z_fused)
+        aux = [head(e) for head, e in zip(self.aux_heads, embeddings)]
+        return logits_fused, aux
+
+    def encode(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
+        """Return z_fused without applying the classifier head."""
+        return self._compute_fused(embeddings)
+
+    def set_phase(self, phase: str) -> None:
+        """
+        Set requires_grad on bridge sub-modules according to training phase.
+
+          - ``cd_warmup``    — freeze everything in the bridge
+          - ``tower_warmup`` — aux_heads trainable, projections + fused head frozen
+          - ``fused_warmup`` — projections + fused head trainable, aux_heads frozen
+          - ``main`` / other — everything trainable
+        """
+        def _rg(module, enabled):
+            for p in module.parameters():
+                p.requires_grad = enabled
+
+        if phase == "cd_warmup":
+            _rg(self, False)
+            return
+        if phase == "tower_warmup":
+            for head in self.aux_heads:
+                _rg(head, True)
+            for W_i in self.W:
+                _rg(W_i, False)
+            for ln_i in self.ln:
+                _rg(ln_i, False)
+            _rg(self.classifier_fused, False)
+            if self.se is not None:
+                _rg(self.se, False)
+            return
+        if phase == "fused_warmup":
+            for head in self.aux_heads:
+                _rg(head, False)
+            for W_i in self.W:
+                _rg(W_i, True)
+            for ln_i in self.ln:
+                _rg(ln_i, True)
+            _rg(self.classifier_fused, True)
+            if self.se is not None:
+                _rg(self.se, True)
+            return
+        _rg(self, True)
+
+    def forward(
+        self, embeddings: list[torch.Tensor]
+    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        """N-tower forward. Delegates to fuse()."""
+        return self.fuse(embeddings)
+
+
+# ---------------------------------------------------------------------------
+# HyperBridge — higher-order bridge over HT module outputs
+# ---------------------------------------------------------------------------
+
+class HyperBridge(nn.Module):
+    """Higher-order bridge that fuses z_fused embeddings from multiple HT modules.
+
+    Operates at the HT output level (z_fused from each HT's Bridge.encode())
+    rather than raw tower embedding level.
+
+    Modes
+    -----
+    embedding_mlp (default)
+        Concatenate all z_fused inputs → MLP → logits.
+        Analogous to EmbeddingMLPEnsembleHT, generalised to N inputs.
+        ``Linear(N*fusion_dim → hidden_dim) → ReLU → Dropout → Linear(hidden_dim → num_classes)``
+
+    classic_bridge
+        Project each input to ``hidden_dim``, Hadamard product, HTClassifier.
+        Analogous to Bridge operating at the HT level — handles inputs of
+        differing dims via per-input projection layers.
+        ``W[i](z_i) → LayerNorm → Hadamard → ReLU → Dropout → Linear(hidden_dim → num_classes)``
+
+    Both modes expose per-input auxiliary HTClassifier heads for BCD-style training.
+
+    Parameters
+    ----------
+    input_dims  : ordered dict {name: dim} for each HT input.
+                  In embedding_mlp mode, dims may differ.
+                  In classic_bridge mode, all dims must be equal (shared space).
+    num_classes : output classes
+    hidden_dim  : hidden dim for the embedding_mlp MLP head
+    mode        : "embedding_mlp" | "classic_bridge"
+    dropout     : dropout throughout
+    """
+
+    def __init__(
+        self,
+        input_dims: dict[str, int],
+        num_classes: int,
+        hidden_dim: int = 256,
+        mode: str = "embedding_mlp",
+        dropout: float = 0.3,
+    ):
+        super().__init__()
+        self.input_names = list(input_dims.keys())
+        self.mode = mode
+        dims = list(input_dims.values())
+
+        if mode == "embedding_mlp":
+            total_dim = sum(dims)
+            self.head = nn.Sequential(
+                nn.Linear(total_dim, hidden_dim), nn.ReLU(),
+                nn.Dropout(dropout), nn.Linear(hidden_dim, num_classes),
+            )
+        elif mode == "classic_bridge":
+            # Project each input to shared fusion_dim space, then Hadamard
+            self.W  = nn.ModuleList([nn.Linear(d, hidden_dim) for d in dims])
+            self.ln = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in dims])
+            self.head = HTClassifier(hidden_dim, num_classes, dropout)
+        else:
+            raise ValueError(f"Unknown HyperBridge mode: {mode!r}")
+
+        # Per-input auxiliary classifiers (both modes)
+        self.aux_heads = nn.ModuleList([
+            HTClassifier(d, num_classes, dropout) for d in dims
+        ])
+
+    def forward(
+        self,
+        inputs: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """Fuse HT-level embeddings.
+
+        Parameters
+        ----------
+        inputs : {name: z_fused [B, dim]} — z_fused from each HT's encode()
+
+        Returns
+        -------
+        logits   : [B, num_classes]
+        aux_dict : {name: [B, num_classes]} — per-input aux head logits
+        """
+        ordered = [inputs[name] for name in self.input_names]
+
+        if self.mode == "embedding_mlp":
+            logits = self.head(torch.cat(ordered, dim=1))
+        else:  # classic_bridge: project → Hadamard → classify
+            h = self.ln[0](self.W[0](ordered[0]))
+            for i in range(1, len(ordered)):
+                h = h * self.ln[i](self.W[i](ordered[i]))
+            logits = self.head(h)
+
+        aux = {name: head(z) for name, head, z
+               in zip(self.input_names, self.aux_heads, ordered)}
+        return logits, aux
+
+
+# ---------------------------------------------------------------------------
+# VoteBridge
+# ---------------------------------------------------------------------------
 
 class VoteBridge(nn.Module):
     def __init__(self, num_classes):
         super().__init__()
-        self.vote_combiner = nn.Linear(num_classes * 2, num_classes)  # two sets of logits
+        self.vote_combiner = nn.Linear(num_classes * 2, num_classes)
 
     def forward(self, out_img, out_md):
         votes = torch.cat([out_img, out_md], dim=1)

v3/classes/clinical_towers.py

@@ -0,0 +1,162 @@
+"""clinical_towers — ClinicalEncoder and ClinicalDataTower.
+
+Self-contained: defines ClinicalEncoder directly (does not import it from
+towers.py).  Imports only TowerBase from towerbase plus infrastructure
+(SEBlock, DataBundle).
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import nn
+
+from v3.classes.towerbase import TowerBase
+from v3.classes.SE_attention import SEBlock
+from v3.classes.data_bundle import DataBundle
+
+
+# ---------------------------------------------------------------------------
+# ClinicalEncoder — MLP over tabular features
+# ---------------------------------------------------------------------------
+
+class ClinicalEncoder(nn.Module):
+    """MLP over DataBundle.vectorize_row outputs (converts to torch inside tower)."""
+
+    def __init__(
+        self,
+        clinical_data: DataBundle,
+        hidden_dim: int = 128,
+        dropout: float = 0.1,
+        use_se: bool = False,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+    ):
+        super().__init__()
+        self.feature_dim = clinical_data.feature_dim
+        self.out_dim = hidden_dim
+        # Two-block MLP so we can optionally freeze/thaw per block.
+        self.block0 = nn.Sequential(
+            nn.Linear(self.feature_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+        )
+        self.block1 = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+        )
+        self.net = nn.Sequential(self.block0, self.block1)
+        self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = (
+            SEBlock(hidden_dim, reduction=se_reduction, residual=True)
+            if use_se
+            else None
+        )
+
+    def forward(self, meta_np_or_torch) -> torch.Tensor:
+        if isinstance(meta_np_or_torch, torch.Tensor):
+            x = meta_np_or_torch
+        else:
+            x = torch.as_tensor(meta_np_or_torch, dtype=torch.float32)
+        h = self.net(x)
+        if self.tower_se is not None:
+            h, _ = self.tower_se(self.tower_ln(h))
+        return h
+
+    def set_freeze_ratio(self, ratio: float):
+        """Optionally freeze earliest blocks of the MLP."""
+        r = max(0.0, min(1.0, float(ratio)))
+        for p in self.block0.parameters():
+            p.requires_grad = True
+        for p in self.block1.parameters():
+            p.requires_grad = True
+        if r >= 0.5:
+            for p in self.block0.parameters():
+                p.requires_grad = False
+        if r >= 1.0:
+            for p in self.block1.parameters():
+                p.requires_grad = False
+
+
+# ---------------------------------------------------------------------------
+# ClinicalDataTower — TowerBase implementation
+# ---------------------------------------------------------------------------
+
+class ClinicalDataTower(TowerBase, nn.Module):
+    """
+    TowerBase implementation for the clinical metadata modality.
+
+    Wraps ClinicalEncoder (MLP over tabular features).
+    Contributes one embedding per eye slot: [z_cd].
+
+    Implements ``cd_warmup_embedding`` so train_towers_epoch can identify
+    this tower for cd_warmup phase via duck typing rather than isinstance checks.
+    """
+
+    def __init__(
+        self,
+        *,
+        clinical_data,
+        cd_hidden_dim: int = 128,
+        cd_dropout: float = 0.1,
+        use_se: bool = False,
+    ):
+        nn.Module.__init__(self)
+        self._encoder = ClinicalEncoder(
+            clinical_data=clinical_data,
+            hidden_dim=cd_hidden_dim,
+            dropout=cd_dropout,
+            use_se=use_se,
+        )
+
+    @property
+    def out_dim(self) -> int:
+        return self._encoder.out_dim
+
+    @property
+    def embed_dims(self) -> list[int]:
+        return [self._encoder.out_dim]
+
+    def set_phase(self, phase: str) -> None:
+        enabled = phase not in ("fused_warmup",)
+        for p in self._encoder.parameters():
+            p.requires_grad = enabled
+
+    def cd_warmup_embedding(
+        self,
+        batch: dict,
+        *,
+        device: torch.device,
+    ) -> Optional[torch.Tensor]:
+        """Return clinical embedding for slot 1, or None if matrix_1 is absent."""
+        m = batch.get("matrix_1")
+        if not torch.is_tensor(m):
+            return None
+        return self._encoder(m.to(device))
+
+    def embed_batch(
+        self,
+        batch: dict,
+        *,
+        device: torch.device,
+        slot: int = 1,
+    ) -> list[torch.Tensor]:
+        m = batch.get(f"matrix_{slot}")
+        if not torch.is_tensor(m):
+            raise ValueError(f"ClinicalDataTower.embed_batch: matrix_{slot} is missing or not a tensor")
+        return [self._encoder(m.to(device))]
+
+    def prepare_fold(
+        self,
+        *,
+        eye_train,
+        bilat_train,
+        bilat_val,
+        bilat_test,
+        image_preprocessor,
+        image_cache,
+        device,
+        args,
+    ) -> None:
+        pass  # shares loader with ImageTower; no per-fold setup needed

v3/classes/config_builder.py

@@ -1,276 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional
-
-import json
-
-from v3.classes.papila_data import PapilaData
-
-
-@dataclass
-class ImportSpec:
-    id: str
-    class_name: str
-    params: Dict[str, Any]
-
-
-@dataclass
-class DataSourceSpec:
-    node_id: str
-    label: str
-    output_type: str
-    source: Optional[Dict[str, Any]]
-    source_ref: Optional[Dict[str, Any]]
-
-
-@dataclass
-class TransformSpec:
-    node_id: str
-    label: str
-    transform_type: str
-    params: Dict[str, Any]
-
-
-@dataclass
-class LoaderSpec:
-    node_id: str
-    label: str
-    input_type: str
-    input_index: str
-    input_key: str
-    output_key: str
-    transforms: List[TransformSpec]
-    data_source: Optional[DataSourceSpec]
-
-
-@dataclass
-class TowerSpec:
-    node_id: str
-    label: str
-    tower_type: str
-    params: Dict[str, Any]
-
-
-@dataclass
-class BridgeSpec:
-    node_id: str
-    label: str
-    method: str
-    params: Dict[str, Any]
-
-
-@dataclass
-class ClassifierSpec:
-    node_id: str
-    label: str
-
-
-@dataclass
-class ConfigAssembly:
-    raw: Dict[str, Any]
-    imports: Dict[str, ImportSpec]
-    data_sources: Dict[str, DataSourceSpec]
-    transforms: Dict[str, TransformSpec]
-    loaders: Dict[str, LoaderSpec]
-    towers: Dict[str, TowerSpec]
-    bridges: Dict[str, BridgeSpec]
-    classifiers: Dict[str, ClassifierSpec]
-
-
-def load_config(path: Path) -> Dict[str, Any]:
-    payload = json.loads(Path(path).read_text())
-    if not isinstance(payload, dict):
-        raise ValueError("Config JSON must be an object.")
-    return payload
-
-
-def assemble_config(path: Path) -> ConfigAssembly:
-    config = load_config(path)
-    meta = config.get("meta", {})
-    imports = _build_imports(meta.get("imports", []))
-    nodes = {node["id"]: node for node in config.get("nodes", [])}
-    edges = config.get("edges", [])
-
-    data_sources: Dict[str, DataSourceSpec] = {}
-    transforms: Dict[str, TransformSpec] = {}
-    loaders: Dict[str, LoaderSpec] = {}
-    towers: Dict[str, TowerSpec] = {}
-    bridges: Dict[str, BridgeSpec] = {}
-    classifiers: Dict[str, ClassifierSpec] = {}
-
-    for node in nodes.values():
-        ntype = node.get("type")
-        if ntype == "data":
-            data_sources[node["id"]] = DataSourceSpec(
-                node_id=node["id"],
-                label=node.get("label", ""),
-                output_type=node.get("outputType", ""),
-                source=node.get("source"),
-                source_ref=node.get("sourceRef"),
-            )
-        elif ntype == "transform":
-            transforms[node["id"]] = TransformSpec(
-                node_id=node["id"],
-                label=node.get("label", ""),
-                transform_type=node.get("transformType", ""),
-                params=_extract_transform_params(node),
-            )
-        elif ntype == "loader":
-            loaders[node["id"]] = LoaderSpec(
-                node_id=node["id"],
-                label=node.get("label", ""),
-                input_type=node.get("inputType", ""),
-                input_index=node.get("inputIndex", ""),
-                input_key=node.get("inputKey", ""),
-                output_key=node.get("outputKey", ""),
-                transforms=[],
-                data_source=None,
-            )
-        elif ntype in ("image_tower", "metadata_tower"):
-            towers[node["id"]] = TowerSpec(
-                node_id=node["id"],
-                label=node.get("label", ""),
-                tower_type=node.get("towerType", "image" if ntype == "image_tower" else "clinical data"),
-                params=_extract_tower_params(node),
-            )
-        elif ntype == "bridge":
-            bridges[node["id"]] = BridgeSpec(
-                node_id=node["id"],
-                label=node.get("label", ""),
-                method=node.get("bridgeMethod", "fusion"),
-                params=_extract_bridge_params(node),
-            )
-        elif ntype == "classifier":
-            classifiers[node["id"]] = ClassifierSpec(
-                node_id=node["id"],
-                label=node.get("label", ""),
-            )
-
-    # attach transforms + data sources to loaders by walking upstream
-    for loader_id, loader in loaders.items():
-        chain = _upstream_chain(loader_id, nodes, edges)
-        for node_id in reversed(chain):
-            if node_id in transforms:
-                loader.transforms.append(transforms[node_id])
-            if node_id in data_sources:
-                loader.data_source = data_sources[node_id]
-
-    return ConfigAssembly(
-        raw=config,
-        imports=imports,
-        data_sources=data_sources,
-        transforms=transforms,
-        loaders=loaders,
-        towers=towers,
-        bridges=bridges,
-        classifiers=classifiers,
-    )
-
-
-def resolve_imports(assembly: ConfigAssembly) -> Dict[str, Any]:
-    resolved: Dict[str, Any] = {}
-    for import_id, spec in assembly.imports.items():
-        if spec.class_name == "PapilaData":
-            params = spec.params
-            resolved[import_id] = PapilaData.from_dirs(
-                image_dir=params.get("image_dir", "Papila/FundusImages"),
-                clinical_dir=params.get("clinical_dir", "Papila/ClinicalData"),
-                label_col=params.get("label_col", "Diagnosis"),
-                cat_cols=params.get("cat_cols", ["Gender", "Phakic/Pseudophakic"]),
-            )
-        else:
-            raise ValueError(f"Unsupported import class {spec.class_name!r}")
-    return resolved
-
-
-def _build_imports(entries: Iterable[Dict[str, Any]]) -> Dict[str, ImportSpec]:
-    specs: Dict[str, ImportSpec] = {}
-    for entry in entries or []:
-        import_id = entry.get("id")
-        if not import_id:
-            continue
-        specs[import_id] = ImportSpec(
-            id=import_id,
-            class_name=entry.get("className", ""),
-            params=entry.get("params", {}) or {},
-        )
-    return specs
-
-
-def _extract_transform_params(node: Dict[str, Any]) -> Dict[str, Any]:
-    return {
-        "transformType": node.get("transformType"),
-        "roiMaskSource": node.get("roiMaskSource"),
-        "roiScale": node.get("roiScale"),
-        "roiTargetSize": node.get("roiTargetSize"),
-        "roiFallback": node.get("roiFallback"),
-        "centerCropSize": node.get("centerCropSize"),
-        "jitterHFlip": node.get("jitterHFlip"),
-        "jitterVFlip": node.get("jitterVFlip"),
-        "jitterRotation": node.get("jitterRotation"),
-        "jitterColorEnabled": node.get("jitterColorEnabled"),
-        "jitterColor": node.get("jitterColor"),
-        "resizeSize": node.get("resizeSize"),
-    }
-
-
-def _extract_tower_params(node: Dict[str, Any]) -> Dict[str, Any]:
-    if node.get("towerType") == "clinical data":
-        return {
-            "hidden_dim": node.get("mdHiddenDim"),
-            "dropout": node.get("mdDropout"),
-            "use_se": node.get("mdUseSe"),
-            "se_reduction": node.get("mdSeReduction"),
-            "se_pre_norm": node.get("mdSePreNorm"),
-            "freeze_ratio": node.get("mdFreezeRatio"),
-        }
-    return {
-        "backbone": node.get("imageBackbone"),
-        "freeze_ratio": node.get("imageFreezeRatio"),
-        "augment": node.get("imageAugment"),
-        "geometry_dim": node.get("imageGeometryDim"),
-        "use_se": node.get("imageUseSe"),
-        "se_reduction": node.get("imageSeReduction"),
-        "se_pre_norm": node.get("imageSePreNorm"),
-    }
-
-
-def _extract_bridge_params(node: Dict[str, Any]) -> Dict[str, Any]:
-    return {
-        "fusion_dim": node.get("bridgeFusionDim"),
-        "use_se": node.get("bridgeUseSe"),
-        "se_reduction": node.get("bridgeSeReduction"),
-        "se_pre_norm": node.get("bridgeSePreNorm"),
-    }
-
-
-def _edge_from(edge: Dict[str, Any]) -> Optional[str]:
-    return edge.get("from") or edge.get("source")
-
-
-def _edge_to(edge: Dict[str, Any]) -> Optional[str]:
-    return edge.get("to") or edge.get("target")
-
-
-def _upstream_chain(start_id: str, nodes: Dict[str, Dict[str, Any]], edges: List[Dict[str, Any]]) -> List[str]:
-    chain: List[str] = []
-    visited = set()
-    current = start_id
-    while True:
-        if current in visited:
-            break
-        visited.add(current)
-        incoming = [edge for edge in edges if _edge_to(edge) == current]
-        if not incoming:
-            break
-        # prefer first incoming edge for now
-        current = _edge_from(incoming[0])
-        if not current:
-            break
-        chain.append(current)
-        node = nodes.get(current)
-        if node and node.get("type") == "data":
-            break
-    return chain

v3/classes/dataset.py

@@ -1,115 +0,0 @@
-from torch.utils.data import Dataset
-from PIL import Image
-import numpy as np
-import torch
-
-
-class ClinicalDataset(Dataset):
-    """Generic dataset wrapping a DataBundle-like instance.
-    Returns (img_tensor, meta_tensor, label)."""
-
-    def __init__(
-        self,
-        clinical_data,
-        img_transform,
-        meta_transform=None,
-        image_preprocessor=None,
-        geometry_provider=None,
-        geometry_dim: int = 0,
-        image_cache: "dict | None" = None,
-    ):
-        self.clinical = clinical_data
-        self.transform_image = img_transform
-        self.meta_transform = meta_transform or (lambda x: x)
-        self.image_preprocessor = image_preprocessor
-        self.geometry_provider = geometry_provider
-        self.geometry_dim = geometry_dim if geometry_provider is not None else 0
-        self.image_cache = image_cache
-
-    def __len__(self):
-        return len(self.clinical.df)
-
-    def __getitem__(self, idx: int):
-        row = self.clinical.df.iloc[idx]
-        # load & transform image
-        img_path = self.clinical.get_image_path(row)
-        cache_key = str(img_path)
-        if self.image_cache is not None and cache_key in self.image_cache:
-            orig_img = Image.fromarray(self.image_cache[cache_key])
-        else:
-            orig_img = Image.open(img_path).convert("RGB")
-            if self.image_cache is not None:
-                self.image_cache[cache_key] = np.asarray(orig_img, dtype=np.uint8)
-        img = orig_img
-        if self.image_preprocessor is not None:
-            img = self.image_preprocessor(img, img_path)
-        img_t = self.transform_image(img)
-        # encode & transform metadata
-        meta = self.clinical.encode_metadata(row)
-        meta_t = self.meta_transform(meta)
-        # label
-        label = self.clinical.get_label(row)
-        if self.geometry_dim > 0:
-            features = None
-            if self.geometry_provider is not None and hasattr(self.geometry_provider, "geometry_features"):
-                features = self.geometry_provider.geometry_features(orig_img, img_path)
-            if features is None:
-                geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32)
-            else:
-                features = np.asarray(features, dtype=np.float32)
-                if features.shape[0] != self.geometry_dim:
-                    geom_vec = torch.zeros(self.geometry_dim, dtype=torch.float32)
-                else:
-                    geom_vec = torch.from_numpy(features)
-            return img_t, meta_t, geom_vec, label
-        return img_t, meta_t, label
-
-
-# ---------------------------------------------------------------------------
-# _ClinicalView — shim used by V2HyperTower._run_fold
-# ---------------------------------------------------------------------------
-
-from .data_bundle import DataBundle  # noqa: E402
-
-
-class _ClinicalView:
-    """Minimal shim so ClinicalDataset can iterate an epoch-specific DataFrame
-    while still delegating encoding/paths/labels to the DataBundle object."""
-
-    def __init__(self, base: DataBundle, df):
-        self.base = base
-        self.df = df
-
-    @property
-    def image_dir(self):
-        return self.base.image_dir
-
-    @property
-    def clinical_dir(self):
-        return self.base.clinical_dir
-
-    @property
-    def id_cols(self):
-        return ("Patient ID", "eyeID")
-
-    @property
-    def label_col(self):
-        return self.base.label_col
-
-    @property
-    def filename_template(self):
-        return getattr(self.base, "filename_template", "RET{pid:03d}{eye}.jpg")
-
-    @property
-    def dim(self):
-        return self.base.feature_dim
-
-    def encode_metadata(self, row):
-        vec = self.base.vectorize_row(row)
-        return torch.as_tensor(vec, dtype=torch.float32)
-
-    def get_image_path(self, row):
-        return self.base.get_image_path(row)
-
-    def get_label(self, row):
-        return int(row[self.base.label_col])

v3/classes/filters.py

@@ -1,119 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Iterable, List, Sequence, Tuple, Union
-import re
-
-import pandas as pd
-
-
-@dataclass
-class RegexFilter:
-    pattern: str
-    flags: int = 0
-
-    def apply_paths(self, paths: Sequence[str]) -> Tuple[List[str], List[str]]:
-        if not self.pattern:
-            return list(paths), []
-        try:
-            regex = re.compile(self.pattern, self.flags)
-        except re.error as err:
-            return list(paths), [f'Invalid regex "{self.pattern}": {err}']
-        filtered = [p for p in paths if regex.search(p)]
-        return filtered, []
-
-
-@dataclass
-class ColumnFilter:
-    column: str
-    operator: str
-    value: str
-    case_insensitive: bool = True
-
-    def apply_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, List[str]]:
-        warnings: List[str] = []
-        if not self.column:
-            return df, ["Column filter missing column name."]
-        columns = list(df.columns)
-        col_index = _resolve_column_index(columns, self.column, warnings)
-        if col_index is None:
-            return df, warnings
-        col_name = columns[col_index]
-        if self.value is None or self.value == "":
-            return df, [f'Column filter "{self.column}" missing value.']
-        series = df[col_name]
-        mask = series.apply(
-            lambda cell: compare_cell(
-                cell, self.value, self.operator, case_insensitive=self.case_insensitive
-            )
-        )
-        return df[mask], warnings
-
-
-FilterSpec = Union[RegexFilter, ColumnFilter]
-
-
-def apply_regex_filters(paths: Sequence[str], filters: Iterable[RegexFilter]) -> Tuple[List[str], List[str]]:
-    filtered = list(paths)
-    warnings: List[str] = []
-    for filt in filters:
-        filtered, warn = filt.apply_paths(filtered)
-        warnings.extend(warn)
-    return filtered, warnings
-
-
-def apply_column_filters(df: pd.DataFrame, filters: Iterable[ColumnFilter]) -> Tuple[pd.DataFrame, List[str]]:
-    filtered = df
-    warnings: List[str] = []
-    for filt in filters:
-        filtered, warn = filt.apply_df(filtered)
-        warnings.extend(warn)
-    return filtered, warnings
-
-
-def compare_cell(cell, raw_value: str, operator: str, case_insensitive: bool = True) -> bool:
-    cell_str = "" if cell is None else str(cell).strip()
-    value_str = "" if raw_value is None else str(raw_value).strip()
-    if case_insensitive:
-        cell_str = cell_str.lower()
-        value_str = value_str.lower()
-    if operator == "=":
-        return cell_str == value_str
-    if operator == "!=":
-        return cell_str != value_str
-    cell_num = _to_float(cell_str)
-    value_num = _to_float(value_str)
-    if cell_num is None or value_num is None:
-        return False
-    if operator == ">":
-        return cell_num > value_num
-    if operator == ">=":
-        return cell_num >= value_num
-    if operator == "<":
-        return cell_num < value_num
-    if operator == "<=":
-        return cell_num <= value_num
-    return False
-
-
-def _resolve_column_index(columns: Sequence[str], column: str, warnings: List[str]) -> int | None:
-    try:
-        return columns.index(column)
-    except ValueError:
-        lower = column.lower()
-        matches = [idx for idx, col in enumerate(columns) if str(col).lower() == lower]
-        if matches:
-            if len(matches) > 1:
-                warnings.append(
-                    f'Column "{column}" matched multiple headers; using "{columns[matches[0]]}".'
-                )
-            return matches[0]
-    warnings.append(f'Column "{column}" not found.')
-    return None
-
-
-def _to_float(value: str) -> float | None:
-    try:
-        return float(value)
-    except (TypeError, ValueError):
-        return None

v3/classes/geometry_towers.py

@@ -1,19 +1,22 @@
-"""Segmentation-map CNN for glaucoma grading.
+"""geometry_towers — GeometryTower and all segmentation-map infrastructure.
 
-Trains a CNN on combined disc/cup segmentation maps — pixel values
+Self-contained: absorbs everything that was in seg_cnn.py so that file can
-  0 = background,  1 = disc (rim only),  2 = cup
+eventually be removed.  Does not import from seg_cnn.py or any other tower file.
-— instead of raw RGB fundus images, forcing the model to learn purely
+Imports only TowerBase from towerbase plus standard infrastructure.
-from optic nerve head geometry (CDR, rim width, cup location, etc.).
 
-Two segmentation sources are supported:
+Contents
-  gt    – rasterise expert contour/mask annotations directly (pure NumPy/PIL,
+--------
-          no CUDA — safe in DataLoader worker processes)
+SegMapRecord           — labelled-eye data record
-  unet  – run a trained UNetSegmenter on the raw fundus image
+_combine_masks         — merge disc/cup binary masks → 3-class label map
-
+crop_to_disc           — tight bounding-box crop
-Usage (import from training script):
+seg_map_to_tensor      — (H,W) uint8 → (C,H,W) float32 tensor
-    from v3.classes.seg_cnn import SegMapRecord, SegMapDataset, SegCNN, seg_map_to_tensor
+load_gt_masks          — load GT disc+cup masks from contour/mask files
+UNetFineTuneDataset    — Dataset for fine-tuning the UNet on GT annotations
+precompute_unet_seg_maps — batch UNet inference helper
+SegMapDataset          — Dataset yielding (seg_tensor, label) pairs
+SegCNN                 — pretrained CNN adapted for segmentation-map input
+GeometryTower          — TowerBase implementation (the main class to use)
 """
-
 from __future__ import annotations
 
 from dataclasses import dataclass
@@ -29,6 +32,10 @@ from torch.utils.data import Dataset
 from torchvision import models, transforms
 from tqdm import tqdm
 
+import pandas as pd
+
+from v3.classes.towerbase import TowerBase
+
 
 # ---------------------------------------------------------------------------
 # Data record
@@ -53,8 +60,7 @@ class SegMapRecord:
 # ---------------------------------------------------------------------------
 
 def _combine_masks(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
-    """
+    """Combine binary disc and cup masks into a 3-class label map.
-    Combine binary disc and cup masks into a 3-class label map.
 
     Returns a uint8 array with values:
       0 — background
@@ -69,8 +75,7 @@ def _combine_masks(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np.ndarray:
 
 
 def crop_to_disc(seg_map: np.ndarray) -> np.ndarray:
-    """
+    """Crop a seg map tightly to the disc bounding box.
-    Crop a seg map tightly to the disc bounding box.
 
     The disc is anywhere seg_map > 0 (i.e. rim or cup).
     Returns the original array unchanged if no disc is found.
@@ -89,8 +94,7 @@ def seg_map_to_tensor(
     channels: int,
     target_size: int,
 ) -> torch.Tensor:
-    """
+    """Convert an (H, W) seg map with values {0, 1, 2} to a float tensor.
-    Convert an (H, W) seg map with values {0, 1, 2} to a float tensor.
 
     channels=1  →  (1, H, W) float in [0, 1]  (values 0/0.5/1.0)
     channels=3  →  (3, H, W) one-hot binary channels [bg, disc_rim, cup]
@@ -118,13 +122,15 @@ def seg_map_to_tensor(
 
 def _load_contour(path: Path) -> np.ndarray:
     """Load x,y contour pairs from a whitespace- or comma-delimited text file."""
+    arr = np.zeros((0, 2), dtype=np.float32)
     for delimiter in (",", None):
         try:
-            arr = np.loadtxt(str(path), delimiter=delimiter, comments="#", dtype=np.float32)
+            candidate = np.loadtxt(str(path), delimiter=delimiter, comments="#", dtype=np.float32)
-            if arr.size > 0:
+            if candidate.size > 0:
+                arr = candidate
                 break
         except Exception:
-            arr = np.zeros((0, 2), dtype=np.float32)
+            pass
     if arr.size == 0 or arr.ndim == 1:
         return np.zeros((0, 2), dtype=np.float32)
     if arr.shape[1] < 2:
@@ -135,13 +141,10 @@ def _load_contour(path: Path) -> np.ndarray:
 def _contour_to_mask(
     coords: np.ndarray, image_size: Tuple[int, int], target_size: int
 ) -> np.ndarray:
-    """
+    """Rasterise a polygon defined by (x, y) coords into a binary mask.
-    Rasterise a polygon defined by (x, y) coords into a binary mask.
 
     image_size is the (width, height) of the original fundus image — the
-    coordinate space the contour was annotated in.  The mask is drawn at
+    coordinate space the contour was annotated in.
-    that resolution then resized to target_size, matching UNetSegmenter's
-    behaviour and avoiding off-canvas clipping.
     """
     if coords is None or len(coords) < 3:
         return np.zeros((target_size, target_size), dtype=np.uint8)
@@ -155,8 +158,7 @@ def _contour_to_mask(
 def _extract_masks_from_image(
     mask_path: Path, target_size: int
 ) -> Tuple[np.ndarray, np.ndarray]:
-    """
+    """Extract disc and cup binary masks from a segmentation image file.
-    Extract disc and cup binary masks from a segmentation image file.
 
     Handles both grayscale label images (e.g. REFUGE .bmp) and
     RGB colour-coded masks.  Returns (disc_mask, cup_mask) both at
@@ -166,7 +168,6 @@ def _extract_masks_from_image(
     arr = np.array(raw)
 
     if arr.ndim == 2:
-        # Grayscale: identify background from edge statistics
         edges = np.concatenate([arr[0], arr[-1], arr[:, 0], arr[:, -1]])
         bg_val = int(np.argmax(np.bincount(edges.astype(np.int64).clip(0, 255), minlength=256)))
         disc_arr = (arr != bg_val).astype(np.uint8)
@@ -182,9 +183,7 @@ def _extract_masks_from_image(
         img_rgb = raw.convert("RGB")
         arr = np.array(img_rgb)
         h, w, c = arr.shape
-        edges_rgb = np.concatenate(
+        edges_rgb = np.concatenate([arr[0], arr[-1], arr[:, 0], arr[:, -1]], axis=0)
-            [arr[0], arr[-1], arr[:, 0], arr[:, -1]], axis=0
-        )
         edge_colors, edge_counts = np.unique(edges_rgb.reshape(-1, c), axis=0, return_counts=True)
         bg_color = edge_colors[int(np.argmax(edge_counts))]
         colors, counts = np.unique(arr.reshape(-1, c), axis=0, return_counts=True)
@@ -200,7 +199,6 @@ def _extract_masks_from_image(
                 cup_color = colors[order[1]]
                 cup_arr[np.all(arr == cup_color, axis=-1)] = 1
 
-    # Resize to target_size with nearest-neighbour to preserve binary values
     def _resize(m: np.ndarray) -> np.ndarray:
         pil = Image.fromarray((m > 0).astype(np.uint8) * 255)
         pil = pil.resize((target_size, target_size), Resampling.NEAREST)
@@ -209,9 +207,8 @@ def _extract_masks_from_image(
     return _resize(disc_arr), _resize(cup_arr)
 
 
-def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np.ndarray]:
+def load_gt_masks(rec: SegMapRecord, target_size: int) -> Tuple[np.ndarray, np.ndarray]:
-    """
+    """Load GT disc + cup masks for one record.
-    Load GT disc + cup masks for one record.
 
     Handles annotation_type "contour" (x,y text file) and "mask" (image file).
     Returns (disc_mask, cup_mask) as uint8 arrays of shape (target_size, target_size).
@@ -219,7 +216,6 @@ def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np
     disc_mask: Optional[np.ndarray] = None
     cup_mask:  Optional[np.ndarray] = None
 
-    # Get original image size so contour coordinates are drawn in the right space
     with Image.open(rec.image_path) as _img:
         image_size = _img.size  # (width, height)
 
@@ -246,7 +242,6 @@ def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np
     if cup_mask is None:
         cup_mask = np.zeros((target_size, target_size), dtype=np.uint8)
 
-    # Structural prior: cup must lie within disc
     cup_mask = (cup_mask > 0) & (disc_mask > 0)
     return disc_mask.astype(np.uint8), cup_mask.astype(np.uint8)
 
@@ -256,11 +251,7 @@ def load_gt_masks(rec: "SegMapRecord", target_size: int) -> Tuple[np.ndarray, np
 # ---------------------------------------------------------------------------
 
 class UNetFineTuneDataset(Dataset):
-    """
+    """Loads (image_tensor, mask_tensor) pairs for fine-tuning the U-Net."""
-    Loads (image_tensor, mask_tensor) pairs for fine-tuning the U-Net on
-    PAPILA GT annotations.  Uses the same preprocessing as UNetSegmenter
-    so the fine-tuned weights are compatible with inference.
-    """
 
     def __init__(
         self,
@@ -292,7 +283,6 @@ class UNetFineTuneDataset(Dataset):
         image = Image.open(rec.image_path).convert("RGB")
         image = image.resize((self.target_size, self.target_size), Resampling.BILINEAR)
         img_tensor = self._normalize(self.to_tensor(image))
-
         disc_mask, cup_mask = load_gt_masks(rec, self.target_size)
         mask_tensor = torch.from_numpy(
             np.stack([disc_mask, cup_mask], axis=0).astype(np.float32)
@@ -301,20 +291,15 @@ class UNetFineTuneDataset(Dataset):
 
 
 # ---------------------------------------------------------------------------
-# U-Net precomputation (run once per full record list, not per fold)
+# U-Net precomputation
 # ---------------------------------------------------------------------------
 
 def precompute_unet_seg_maps(
-    records: List["SegMapRecord"],
+    records: List[SegMapRecord],
     segmenter,
     threshold: float = 0.5,
 ) -> List[np.ndarray]:
-    """
+    """Run the U-Net on every record and return a list of combined seg maps."""
-    Run the U-Net on every record and return a list of combined seg maps.
-
-    Call this once before the CV loop and pass the results to each fold's
-    SegMapDataset via precomputed_seg_maps, so the U-Net isn't re-run per fold.
-    """
     to_tensor = transforms.ToTensor()
     seg_maps = []
     for rec in tqdm(records, desc="U-Net inference", unit="img", leave=False):
@@ -334,28 +319,11 @@ def precompute_unet_seg_maps(
 
 
 # ---------------------------------------------------------------------------
-# Dataset
+# SegMapDataset
 # ---------------------------------------------------------------------------
 
 class SegMapDataset(Dataset):
-    """
+    """PyTorch Dataset that yields (seg_tensor, label) pairs."""
-    PyTorch Dataset that yields (seg_tensor, label) pairs.
-
-    Parameters
-    ----------
-    records      : list of SegMapRecord
-    target_size  : CNN input spatial size (images are resized to this)
-    channels     : 1 = single-channel label map;  3 = one-hot three channels
-    augment      : apply random flips + rotation (for training set)
-    unet_segmenter : if provided, use U-Net predictions instead of GT masks;
-                     must be a loaded UNetSegmenter with model weights set
-    unet_threshold : threshold for U-Net logit → binary mask
-    seg_target_size: resolution at which GT masks are rasterised (or U-Net
-                     output size).  Default 512 matches UNetSegmenter default.
-    crop_to_disc   : crop the seg map tightly to the disc bounding box before
-                     resizing to target_size (default True — eliminates the
-                     background zeros that make up most of the full image)
-    """
 
     def __init__(
         self,
@@ -366,7 +334,7 @@ class SegMapDataset(Dataset):
         unet_segmenter=None,
         unet_threshold: float = 0.5,
         seg_target_size: int = 512,
-        crop_to_disc: bool = True,
+        crop_to_disc_flag: bool = True,
         precomputed_seg_maps: Optional[List[np.ndarray]] = None,
     ) -> None:
         self.records = records
@@ -374,7 +342,7 @@ class SegMapDataset(Dataset):
         self.channels = channels
         self.augment = augment
         self.seg_target_size = seg_target_size
-        self.crop_to_disc = crop_to_disc
+        self.crop_to_disc_flag = crop_to_disc_flag
 
         if precomputed_seg_maps is not None:
             self._seg_maps = precomputed_seg_maps
@@ -385,13 +353,10 @@ class SegMapDataset(Dataset):
         else:
             self._seg_maps = None
 
-    # ------------------------------------------------------------------
     def __len__(self) -> int:
         return len(self.records)
 
-    # ------------------------------------------------------------------
     def _augment(self, seg_map: np.ndarray) -> np.ndarray:
-        """Random flips + 90° rotations (label-safe since NEAREST resize)."""
         if np.random.rand() < 0.5:
             seg_map = np.fliplr(seg_map)
         if np.random.rand() < 0.5:
@@ -401,19 +366,16 @@ class SegMapDataset(Dataset):
             seg_map = np.rot90(seg_map, k=k)
         return np.ascontiguousarray(seg_map)
 
-    # ------------------------------------------------------------------
     def __getitem__(self, idx: int):
         rec = self.records[idx]
-
         if self._seg_maps is not None:
             seg_map = self._seg_maps[idx]
         else:
             disc_mask, cup_mask = load_gt_masks(rec, self.seg_target_size)
             seg_map = _combine_masks(disc_mask, cup_mask)
 
-        if self.crop_to_disc:
+        if self.crop_to_disc_flag:
             seg_map = crop_to_disc(seg_map)
-
         if self.augment:
             seg_map = self._augment(seg_map)
 
@@ -422,19 +384,24 @@ class SegMapDataset(Dataset):
 
 
 # ---------------------------------------------------------------------------
-# Model
+# SegCNN
 # ---------------------------------------------------------------------------
 
+_SEGCNN_FEAT_DIM = {
+    "resnet18": 512,
+    "resnet50": 2048,
+    "efficientnet_b0": 1280,
+}
+
+
 class SegCNN(nn.Module):
-    """
+    """Pretrained CNN backbone adapted for segmentation-map input.
-    Pretrained CNN backbone adapted for segmentation-map input.
 
     Parameters
     ----------
     num_classes : output classes (2 for binary glaucoma grading)
     backbone    : "resnet18" | "resnet50" | "efficientnet_b0"
-    pretrained  : initialise with ImageNet weights (recommended even for
+    pretrained  : initialise with ImageNet weights
-                  non-RGB input — transfer generalises across domains)
     in_channels : 1 (single label map) or 3 (one-hot channels)
     dropout     : dropout rate before the final classifier head
     """
@@ -448,7 +415,6 @@ class SegCNN(nn.Module):
         dropout: float = 0.3,
     ) -> None:
         super().__init__()
-
         weights_arg = "DEFAULT" if pretrained else None
 
         if backbone == "resnet18":
@@ -466,7 +432,6 @@ class SegCNN(nn.Module):
         else:
             raise ValueError(f"Unknown backbone: {backbone!r}")
 
-        # Adapt first conv layer if in_channels ≠ 3
         if in_channels != 3:
             first_conv = self._find_first_conv(base)
             new_conv = nn.Conv2d(
@@ -478,7 +443,6 @@ class SegCNN(nn.Module):
                 bias=first_conv.bias is not None,
             )
             if pretrained:
-                # Average pretrained RGB weights across channel dim
                 with torch.no_grad():
                     new_conv.weight.copy_(
                         first_conv.weight.mean(dim=1, keepdim=True).expand_as(new_conv.weight)
@@ -491,7 +455,6 @@ class SegCNN(nn.Module):
             nn.Linear(feat_dim, num_classes),
         )
 
-    # ------------------------------------------------------------------
     @staticmethod
     def _find_first_conv(module: nn.Module) -> nn.Conv2d:
         for m in module.modules():
@@ -501,7 +464,6 @@ class SegCNN(nn.Module):
 
     @staticmethod
     def _replace_first_conv(module: nn.Module, new_conv: nn.Conv2d) -> None:
-        """Replace the first Conv2d in-place (handles resnet and efficientnet)."""
         for name, child in module.named_children():
             if isinstance(child, nn.Conv2d):
                 setattr(module, name, new_conv)
@@ -513,9 +475,304 @@ class SegCNN(nn.Module):
                 pass
         raise RuntimeError("Could not replace first Conv2d")
 
-    # ------------------------------------------------------------------
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         feats = self.backbone(x)
         if feats.dim() > 2:
             feats = feats.flatten(1)
         return self.head(feats)
+
+
+# ---------------------------------------------------------------------------
+# GeometryTower — TowerBase implementation
+# ---------------------------------------------------------------------------
+
+class GeometryTower(TowerBase, nn.Module):
+    """TowerBase implementation for the optic-disc/cup segmentation modality.
+
+    Encodes a 3-class disc/cup segmentation map (bg=0, rim=1, cup=2) through a
+    CNN backbone, contributing one spatial embedding to the bridge.
+
+    The seg map is produced from GT annotations (manifest-based) or from a
+    trained U-Net, depending on ``geometry_source``.
+
+    ``prepare_fold`` builds a seg-map generator, pre-computes all maps for the
+    fold, and caches them keyed by image path.  ``augment_samples`` then
+    injects ``seg_map_1`` / ``seg_map_2`` float32 numpy arrays (shape C×H×W)
+    into each sample dict so the DataLoader delivers them as tensors to
+    ``embed_batch``.
+
+    Parameters
+    ----------
+    backbone              : CNN backbone — "resnet18" | "resnet50" | "efficientnet_b0"
+    in_channels           : 1 (label map) or 3 (one-hot disc/rim/cup channels)
+    pretrained            : initialise backbone with ImageNet weights
+    frozen                : if True, backbone is always frozen
+    target_size           : spatial size the seg map tensor is resized to
+    seg_target_size       : resolution at which GT masks are rasterised / U-Net runs
+    crop_to_disc          : crop seg map tightly to disc bounding box before resizing
+    geometry_source       : "gt" (manifest annotations) or "unet" (U-Net predictions)
+    manifest_path         : path to the geometry manifest CSV (required)
+    weights_path          : path to UNet checkpoint (required when source="unet")
+    unet_normalize        : UNet normalisation mode (default "per_image")
+    unet_threshold        : UNet mask threshold (default 0.5)
+    finetune_unet_epochs  : epochs to fine-tune U-Net per fold (0 = disabled)
+    finetune_unet_lr      : learning rate for U-Net fine-tuning
+    """
+
+    def __init__(
+        self,
+        *,
+        backbone: str = "resnet18",
+        in_channels: int = 3,
+        pretrained: bool = True,
+        frozen: bool = False,
+        target_size: int = 224,
+        seg_target_size: int = 512,
+        crop_to_disc: bool = True,
+        geometry_source: str = "gt",
+        manifest_path=None,
+        weights_path=None,
+        unet_normalize: str = "per_image",
+        unet_threshold: float = 0.5,
+        finetune_unet_epochs: int = 0,
+        finetune_unet_lr: float = 1e-5,
+    ):
+        nn.Module.__init__(self)
+        self._backbone_name        = backbone
+        self._in_channels          = in_channels
+        self._frozen               = frozen
+        self._target_size          = target_size
+        self._seg_target_size      = seg_target_size
+        self._crop_to_disc         = crop_to_disc
+        self._geometry_source      = geometry_source
+        self._manifest_path        = Path(manifest_path) if manifest_path is not None else None
+        self._weights_path         = Path(weights_path)  if weights_path  is not None else None
+        self._unet_normalize       = unet_normalize
+        self._unet_threshold       = unet_threshold
+        self._finetune_unet_epochs = finetune_unet_epochs
+        self._finetune_unet_lr     = finetune_unet_lr
+
+        self._out_dim  = _SEGCNN_FEAT_DIM.get(backbone, 512)
+        self._seg_cnn  = SegCNN(
+            num_classes=2,
+            backbone=backbone,
+            pretrained=pretrained,
+            in_channels=in_channels,
+        )
+        self._seg_cache: dict = {}  # image_path_str → float32 (C, H, W) numpy array
+
+    # ------------------------------------------------------------------
+    # TowerBase interface
+    # ------------------------------------------------------------------
+
+    @property
+    def embed_dims(self) -> list[int]:
+        return [self._out_dim]
+
+    @property
+    def total_epochs(self) -> int:
+        return 0 if self._frozen else 1
+
+    def set_phase(self, phase: str) -> None:
+        trainable = not self._frozen and phase in ("tower_warmup", "main")
+        for p in self._seg_cnn.parameters():
+            p.requires_grad = trainable
+
+    def prepare_fold(
+        self,
+        *,
+        eye_train,
+        bilat_train,
+        bilat_val,
+        bilat_test,
+        image_preprocessor,
+        image_cache,
+        device,
+        args,
+    ) -> None:
+        """Build seg-map generator and pre-compute maps for all fold images."""
+        if self._manifest_path is None:
+            raise ValueError("GeometryTower requires manifest_path")
+
+        all_paths: dict = {}
+        for split in (eye_train, bilat_train, bilat_val, bilat_test):
+            for s in split:
+                for slot in ("image_1", "image_2"):
+                    p = s.get(slot)
+                    if p is not None:
+                        all_paths[str(Path(p).resolve())] = None
+
+        if self._geometry_source == "unet":
+            self._prepare_fold_unet(list(all_paths.keys()), eye_train, device)
+        else:
+            self._prepare_fold_gt(list(all_paths.keys()))
+
+    def augment_samples(self, samples: list) -> list:
+        """Inject ``seg_map_1`` / ``seg_map_2`` float32 arrays into each sample dict.
+
+        Arrays have shape (C, H, W) and are collated by the DataLoader into
+        (B, C, H, W) tensors delivered to ``embed_batch``.
+        """
+        blank = np.zeros(
+            (self._in_channels, self._target_size, self._target_size), dtype=np.float32
+        )
+        for s in samples:
+            for img_slot, seg_slot in (("image_1", "seg_map_1"), ("image_2", "seg_map_2")):
+                img_path = s.get(img_slot)
+                if img_path is None:
+                    continue
+                key = str(Path(img_path).resolve())
+                s[seg_slot] = self._seg_cache.get(key, blank)
+        return samples
+
+    def embed_batch(
+        self,
+        batch: dict,
+        *,
+        device: torch.device,
+        slot: int = 1,
+    ) -> list[torch.Tensor]:
+        seg = batch.get(f"seg_map_{slot}")
+        if seg is None or not torch.is_tensor(seg):
+            ref = batch.get(f"image_{slot}")
+            bs = ref.shape[0] if torch.is_tensor(ref) else 1
+            return [torch.zeros(bs, self._out_dim, device=device)]
+        return [self._seg_cnn.backbone(seg.float().to(device))]
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _seg_map_to_array(self, seg_map: np.ndarray) -> np.ndarray:
+        """Apply crop + resize and return a (C, H, W) float32 numpy array."""
+        if self._crop_to_disc:
+            seg_map = crop_to_disc(seg_map)
+        return seg_map_to_tensor(seg_map, self._in_channels, self._target_size).numpy()
+
+    def _prepare_fold_gt(self, image_paths: list) -> None:
+        """Pre-compute GT seg maps from manifest annotations."""
+        manifest_df = pd.read_csv(self._manifest_path)
+        manifest_df["_img_key"] = manifest_df["image_path"].apply(
+            lambda p: str(Path(p).resolve())
+        )
+        manifest_index = manifest_df.set_index("_img_key").to_dict("index")
+
+        print(
+            f"[GeometryTower] pre-computing GT seg maps for {len(image_paths)} images...",
+            flush=True,
+        )
+        n_ok = 0
+        blank = np.zeros((self._seg_target_size, self._seg_target_size), dtype=np.uint8)
+        for img_path in image_paths:
+            entry = manifest_index.get(img_path)
+            if entry is None:
+                self._seg_cache[img_path] = self._seg_map_to_array(blank)
+                continue
+            rec = SegMapRecord(
+                sample_id="",
+                image_path=Path(img_path),
+                annotation_disc=Path(entry["annotation_disc"]),
+                annotation_cup=Path(entry["annotation_cup"]),
+                annotation_type_disc=entry["annotation_type_disc"],
+                annotation_type_cup=entry["annotation_type_cup"],
+                patient_id=0,
+                eye="",
+                label=0,
+            )
+            try:
+                disc_mask, cup_mask = load_gt_masks(rec, self._seg_target_size)
+                self._seg_cache[img_path] = self._seg_map_to_array(
+                    _combine_masks(disc_mask, cup_mask)
+                )
+                n_ok += 1
+            except Exception:
+                self._seg_cache[img_path] = self._seg_map_to_array(blank)
+        print(f"[GeometryTower] {n_ok}/{len(image_paths)} GT seg maps computed", flush=True)
+
+    def _prepare_fold_unet(self, image_paths: list, eye_train: list, device) -> None:
+        """Pre-compute U-Net seg maps, with optional per-fold fine-tuning."""
+        from v3.classes.unet_segmenter import UNetSegmenter
+        from torch.utils.data import DataLoader as _DL
+
+        if self._weights_path is None:
+            raise ValueError("GeometryTower(source='unet') requires weights_path")
+
+        segmenter = UNetSegmenter(
+            manifest_path=self._manifest_path,
+            normalize=self._unet_normalize,
+        )
+        state = torch.load(self._weights_path, map_location=segmenter.device)
+        segmenter.model.load_state_dict(state.get("model", state))
+        segmenter.model.to(segmenter.device).eval()
+
+        if self._finetune_unet_epochs > 0:
+            print(
+                f"[GeometryTower] fine-tuning U-Net for {self._finetune_unet_epochs} epochs...",
+                flush=True,
+            )
+            ft_loader = _DL(
+                UNetFineTuneDataset(
+                    self._build_records_from_samples(eye_train),
+                    target_size=segmenter.target_size,
+                    normalize=self._unet_normalize,
+                ),
+                batch_size=4, shuffle=True, num_workers=0,
+            )
+            optimizer = torch.optim.Adam(segmenter.model.parameters(), lr=self._finetune_unet_lr)
+            criterion = torch.nn.BCEWithLogitsLoss()
+            segmenter.model.train()
+            for _ in range(self._finetune_unet_epochs):
+                for images, masks in ft_loader:
+                    images, masks = images.to(segmenter.device), masks.to(segmenter.device)
+                    optimizer.zero_grad()
+                    criterion(segmenter.model(images), masks).backward()
+                    optimizer.step()
+            segmenter.model.eval()
+
+        print(
+            f"[GeometryTower] running U-Net inference on {len(image_paths)} images...",
+            flush=True,
+        )
+        records = [
+            SegMapRecord(
+                sample_id="", image_path=Path(p),
+                annotation_disc=Path(p), annotation_cup=Path(p),
+                annotation_type_disc="", annotation_type_cup="",
+                patient_id=0, eye="", label=0,
+            )
+            for p in image_paths
+        ]
+        seg_maps = precompute_unet_seg_maps(records, segmenter, self._unet_threshold)
+        for img_path, seg_map in zip(image_paths, seg_maps):
+            self._seg_cache[img_path] = self._seg_map_to_array(seg_map)
+        print(f"[GeometryTower] {len(seg_maps)} U-Net seg maps cached", flush=True)
+
+    def _build_records_from_samples(self, samples: list) -> list:
+        """Build SegMapRecord list from HyperTower sample dicts (for U-Net fine-tuning)."""
+        manifest_df = pd.read_csv(self._manifest_path)
+        manifest_df["_img_key"] = manifest_df["image_path"].apply(
+            lambda p: str(Path(p).resolve())
+        )
+        manifest_index = manifest_df.set_index("_img_key").to_dict("index")
+        records = []
+        for s in samples:
+            for slot in ("image_1", "image_2"):
+                p = s.get(slot)
+                if p is None:
+                    continue
+                key = str(Path(p).resolve())
+                entry = manifest_index.get(key)
+                if entry is None:
+                    continue
+                records.append(SegMapRecord(
+                    sample_id="",
+                    image_path=Path(p),
+                    annotation_disc=Path(entry["annotation_disc"]),
+                    annotation_cup=Path(entry["annotation_cup"]),
+                    annotation_type_disc=entry["annotation_type_disc"],
+                    annotation_type_cup=entry["annotation_type_cup"],
+                    patient_id=int(s.get("patient_id", 0)),
+                    eye=str(s.get("eye", "")),
+                    label=int(s.get("label", 0)),
+                ))
+        return records

v3/classes/image_towers.py

@@ -0,0 +1,252 @@
+"""image_towers — ImageEncoder, SiameseImageTower, and ImageTower (TowerBase).
+
+Self-contained image-modality tower layer.  No dependencies on other tower
+files, bridge, or model classes.  Clear contract: accepts an image batch,
+returns a fixed-size embedding vector.
+"""
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from v3.classes.towerbase import TowerBase, build_backbone
+from v3.classes.backbones import BACKBONES
+from v3.classes.SE_attention import SEBlock
+
+
+# ---------------------------------------------------------------------------
+# ImageEncoder — vision backbone → pooled feature vector
+# ---------------------------------------------------------------------------
+
+class ImageEncoder(nn.Module):
+    """Vision backbone → pooled feature vector.
+
+    Wraps a torchvision backbone (default weights), strips the classifier,
+    and optionally appends an SE attention block and/or a geometry vector.
+
+    Parameters
+    ----------
+    backbone     : backbone key (see backbones.py)
+    freeze_ratio : fraction of early blocks to freeze in [0, 1]
+    use_se       : apply SE attention over the pooled feature vector
+    augment      : include random flip/rotation/jitter in the transform
+    geometry_dim : if > 0, concatenate a geometry vector of this length
+    """
+
+    def __init__(
+        self,
+        backbone: str = "efficientnet_b0",
+        freeze_ratio: float = 0.0,
+        use_se: bool = False,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+        augment: bool = True,
+        geometry_dim: int = 0,
+    ):
+        super().__init__()
+        self.backbone, base_dim, self.transform = build_backbone(
+            backbone, freeze_ratio, augment=augment
+        )
+        self._name   = backbone
+        key          = (self._name or "").lower()
+        self._spec   = BACKBONES[key]
+        self._blocks = self._spec.blocks(self.backbone)
+        self.base_dim     = base_dim
+        self.geometry_dim = max(0, int(geometry_dim))
+        self.out_dim      = self.base_dim + self.geometry_dim
+        self.tower_ln = nn.LayerNorm(self.base_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = (
+            SEBlock(self.base_dim, reduction=se_reduction, residual=True)
+            if use_se else None
+        )
+
+    def forward(
+        self, x: torch.Tensor, geometry: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        y = self.backbone(x)
+        assert y.dim() == 2 and y.size(1) == self.base_dim, (
+            f"Expected features [N,{self.base_dim}], got {tuple(y.shape)}"
+        )
+        if self.tower_se is not None:
+            y, _ = self.tower_se(self.tower_ln(y))
+        if self.geometry_dim > 0:
+            if geometry is None or geometry.numel() == 0:
+                geom = torch.zeros(y.size(0), self.geometry_dim, device=y.device, dtype=y.dtype)
+            else:
+                geom = geometry.unsqueeze(0) if geometry.dim() == 1 else geometry
+                geom = geom.to(device=y.device, dtype=y.dtype)
+                if geom.size(0) != y.size(0):
+                    raise ValueError(f"Geometry batch size mismatch: {geom.size(0)} vs {y.size(0)}")
+                if geom.size(1) != self.geometry_dim:
+                    raise ValueError(f"Expected geometry dim {self.geometry_dim}, got {geom.size(1)}")
+            y = torch.cat([y, geom], dim=1)
+        return y
+
+    def set_freeze_ratio(self, ratio: float) -> None:
+        """Dynamically freeze earliest floor(N*ratio) backbone blocks."""
+        r = max(0.0, min(1.0, float(ratio)))
+        freeze_n = int(math.floor(len(self._blocks) * r))
+        for b in self._blocks:
+            for p in b.parameters():
+                p.requires_grad = True
+        for b in self._blocks[:freeze_n]:
+            for p in b.parameters():
+                p.requires_grad = False
+
+
+# ---------------------------------------------------------------------------
+# SiameseImageTower — shared-weight bilateral image encoder
+# ---------------------------------------------------------------------------
+
+class SiameseImageTower(nn.Module):
+    """Shared-weight bilateral image tower.
+
+    Runs OD and OS images through a single shared backbone and returns
+    cat([f_mean, f_delta]) where:
+        f_mean  = (f_od + f_os) / 2   — shared bilateral representation
+        f_delta = f_od - f_os          — signed asymmetry (OD-relative)
+
+    out_dim = 2 × backbone_out_dim.  When x_os is None the tower degrades
+    gracefully: f_mean = f_od, f_delta = zeros.
+    """
+
+    def __init__(
+        self,
+        backbone: str = "efficientnet_b0",
+        freeze_ratio: float = 0.0,
+        use_se: bool = False,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+        augment: bool = True,
+    ):
+        super().__init__()
+        self._tower = ImageEncoder(
+            backbone=backbone,
+            freeze_ratio=freeze_ratio,
+            use_se=use_se,
+            se_reduction=se_reduction,
+            se_pre_norm=se_pre_norm,
+            augment=augment,
+        )
+        self.out_dim   = self._tower.out_dim * 2
+        self.transform = self._tower.transform
+
+    def forward(
+        self,
+        x_od: torch.Tensor,
+        x_os: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        f_od = self._tower(x_od)
+        if x_os is None:
+            return torch.cat([f_od, torch.zeros_like(f_od)], dim=1)
+        f_os = self._tower(x_os)
+        return torch.cat([(f_od + f_os) * 0.5, f_od - f_os], dim=1)
+
+    def set_freeze_ratio(self, ratio: float) -> None:
+        self._tower.set_freeze_ratio(ratio)
+
+
+# ---------------------------------------------------------------------------
+# ImageTower — TowerBase implementation
+# ---------------------------------------------------------------------------
+
+class ImageTower(TowerBase, nn.Module):
+    """TowerBase implementation for the fundus image modality.
+
+    Wraps ImageEncoder (backbone → pooled features).
+    Contributes one embedding per eye slot: [z_img].
+    """
+
+    def __init__(
+        self,
+        *,
+        backbone: str,
+        freeze_ratio: float = 0.0,
+        augment: bool = True,
+        use_se: bool = False,
+    ):
+        nn.Module.__init__(self)
+        self._encoder      = ImageEncoder(backbone=backbone, freeze_ratio=freeze_ratio,
+                                          augment=augment, use_se=use_se)
+        self._train_loader = None
+
+    @property
+    def transform(self):
+        return self._encoder.transform
+
+    @property
+    def out_dim(self) -> int:
+        return self._encoder.out_dim
+
+    @property
+    def embed_dims(self) -> list[int]:
+        return [self._encoder.out_dim]
+
+    @property
+    def total_epochs(self) -> int:
+        return 0
+
+    @property
+    def train_loader(self) -> Optional[DataLoader]:
+        return self._train_loader
+
+    def set_phase(self, phase: str) -> None:
+        enabled = phase not in ("cd_warmup", "fused_warmup")
+        for p in self._encoder.parameters():
+            p.requires_grad = enabled
+
+    def embed_batch(
+        self,
+        batch: dict,
+        *,
+        device: torch.device,
+        slot: int = 1,
+    ) -> list[torch.Tensor]:
+        x = batch.get(f"image_{slot}")
+        if not torch.is_tensor(x):
+            raise ValueError(f"ImageTower.embed_batch: image_{slot} missing or not a tensor")
+        return [self._encoder(x.to(device))]
+
+    def prepare_fold(
+        self,
+        *,
+        eye_train,
+        bilat_train,
+        bilat_val,
+        bilat_test,
+        image_preprocessor,
+        image_cache,
+        device,
+        args,
+    ) -> None:
+        from v3.classes.loader_factory import (
+            build_balanced_sampler,
+            filter_eye_samples,
+            make_loader,
+        )
+        from v3.classes.profiles import build_papila_profile
+
+        profile_eye = build_papila_profile(
+            patient_col="Patient ID", label_col=args.label_col, sample_mode="eye"
+        )
+        slots_eye    = profile_eye.slot_descriptors()
+        use_balanced = bool(getattr(args, "balanced_sampling", False))
+        _persistent  = args.num_workers > 0
+        loader_kw    = dict(
+            batch_size=args.batch_size,
+            num_workers=args.num_workers,
+            image_cache=image_cache,
+            persistent_workers=_persistent,
+        )
+        eye_samples = filter_eye_samples(eye_train)
+        sampler     = build_balanced_sampler(eye_samples) if use_balanced else None
+        self._train_loader = make_loader(
+            eye_samples, slots_eye,
+            image_transform=self.transform,
+            image_preprocessor=image_preprocessor,
+            shuffle=True, sampler=sampler, **loader_kw,
+        )

v3/classes/loader_factory.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from typing import Any, Callable, Optional
 
 import torch
-from torch.utils.data import DataLoader, WeightedRandomSampler
+from torch.utils.data import DataLoader, Sampler, WeightedRandomSampler
 
 from .network_manager import LoaderBundle, PatientSplit
 from .slot_dataset import SlotDataset, slot_collate
@@ -210,7 +210,7 @@ def make_loader(
     batch_size: int,
     shuffle: bool,
     num_workers: int,
-    sampler: Optional[WeightedRandomSampler] = None,
+    sampler: Optional[Sampler] = None,
     persistent_workers: bool = False,
 ) -> DataLoader:
     ds = SlotDataset(

v3/classes/model_builder.py

@@ -1,148 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Any, Callable, Optional
-
-import torch
-from torch import nn
-
-from v3.classes.bridges import Bridge, VoteBridge
-from v3.classes.towers import ImageTower, ClinicalTower
-
-from .config_builder import ConfigAssembly
-from .transforms import build_transform_chain
-
-
-@dataclass
-class V2ModelBundle:
-    image_tower: Optional[ImageTower]
-    metadata_tower: Optional[ClinicalTower]
-    bridge: Optional[nn.Module]
-    classifier: Optional[nn.Module]
-    image_transform: Optional[Callable]
-    matrix_transform: Optional[Callable]
-
-
-def build_model_bundle(
-    assembly: ConfigAssembly,
-    clinical: Any,
-    *,
-    device: Optional[torch.device] = None,
-    strict: bool = True,
-) -> V2ModelBundle:
-    """
-    Build torch modules and input transforms from a V2 config assembly.
-    """
-    image_tower_spec = _pick_tower(assembly, "image")
-    cd_tower_spec = _pick_tower(assembly, "clinical data")
-    bridge_spec = _pick_bridge(assembly)
-    image_loader = _pick_loader(assembly, input_type="image")
-
-    clinical_core = getattr(clinical, "clinical", clinical)
-    num_classes = _infer_num_classes(clinical)
-
-    img_tower = None
-    if image_tower_spec is not None:
-        img_tower = ImageTower(
-            backbone=image_tower_spec.params.get("backbone", "efficientnet_b0"),
-            freeze_ratio=float(image_tower_spec.params.get("freeze_ratio", 0.0) or 0.0),
-            use_se=bool(image_tower_spec.params.get("use_se", False)),
-            se_reduction=int(image_tower_spec.params.get("se_reduction", 16) or 16),
-            se_pre_norm=bool(image_tower_spec.params.get("se_pre_norm", True)),
-            augment=bool(image_tower_spec.params.get("augment", True)),
-            geometry_dim=int(image_tower_spec.params.get("geometry_dim", 0) or 0),
-        )
-        if device is not None:
-            img_tower = img_tower.to(device)
-
-    cd_tower = None
-    if cd_tower_spec is not None:
-        cd_tower = ClinicalTower(
-            clinical_core,
-            hidden_dim=int(cd_tower_spec.params.get("hidden_dim", 128) or 128),
-            dropout=float(cd_tower_spec.params.get("dropout", 0.1) or 0.1),
-            use_se=bool(cd_tower_spec.params.get("use_se", False)),
-            se_reduction=int(cd_tower_spec.params.get("se_reduction", 16) or 16),
-            se_pre_norm=bool(cd_tower_spec.params.get("se_pre_norm", True)),
-        )
-        if device is not None:
-            cd_tower = cd_tower.to(device)
-
-    bridge = None
-    if bridge_spec is not None and img_tower is not None and cd_tower is not None:
-        if bridge_spec.method == "consensus":
-            bridge = VoteBridge(num_classes=num_classes)
-        else:
-            bridge = Bridge(
-                img_dim=img_tower.out_dim,
-                meta_dim=cd_tower.out_dim,
-                num_classes=num_classes,
-                fusion_dim=int(bridge_spec.params.get("fusion_dim", 256) or 256),
-                mode="fused",
-                use_se=bool(bridge_spec.params.get("use_se", True)),
-                se_reduction=int(bridge_spec.params.get("se_reduction", 16) or 16),
-                se_pre_norm=bool(bridge_spec.params.get("se_pre_norm", True)),
-            )
-        if device is not None:
-            bridge = bridge.to(device)
-
-    classifier = None
-    if assembly.classifiers:
-        classifier = nn.Identity()
-        if device is not None:
-            classifier = classifier.to(device)
-
-    image_transform = None
-    if image_loader is not None and image_tower_spec is not None:
-        image_transform = build_transform_chain(
-            image_loader.transforms,
-            backbone_name=image_tower_spec.params.get("backbone", "efficientnet_b0"),
-            augment=bool(image_tower_spec.params.get("augment", True)),
-            strict=strict,
-        )
-
-    return V2ModelBundle(
-        image_tower=img_tower,
-        metadata_tower=cd_tower,
-        bridge=bridge,
-        classifier=classifier,
-        image_transform=image_transform,
-        matrix_transform=None,
-    )
-
-
-def _pick_tower(assembly: ConfigAssembly, tower_type: str):
-    matches = [tower for tower in assembly.towers.values() if tower.tower_type == tower_type]
-    if not matches:
-        return None
-    if len(matches) > 1:
-        raise ValueError(f"Multiple {tower_type} towers found; only one is supported for now.")
-    return matches[0]
-
-
-def _pick_bridge(assembly: ConfigAssembly):
-    if not assembly.bridges:
-        return None
-    if len(assembly.bridges) > 1:
-        raise ValueError("Multiple bridges found; only one is supported for now.")
-    return next(iter(assembly.bridges.values()))
-
-
-def _pick_loader(assembly: ConfigAssembly, input_type: str):
-    matches = [loader for loader in assembly.loaders.values() if loader.input_type == input_type]
-    if not matches:
-        return None
-    if len(matches) > 1:
-        raise ValueError(f"Multiple loaders with input_type={input_type!r} found.")
-    return matches[0]
-
-
-def _infer_num_classes(clinical: Any) -> int:
-    df = getattr(clinical, "df", None)
-    label_col = getattr(clinical, "label_col", None)
-    if df is None and hasattr(clinical, "clinical"):
-        df = clinical.clinical.df
-        label_col = clinical.clinical.label_col
-    if df is None or label_col is None or label_col not in df.columns:
-        return 2
-    return int(df[label_col].dropna().nunique())

v3/classes/towerbase.py

@@ -0,0 +1,455 @@
+"""towerbase — unified TowerBase ABC, backbone factory, and modular training utilities.
+
+This module is the structural backbone of the HyperTower v3 architecture.
+It provides the abstract tower interface plus the training/eval helpers that
+operate on any list of TowerBase instances.
+
+Design principles
+-----------------
+* No concrete tower classes are defined here (ImageEncoder, ClinicalEncoder, etc.
+  live in their respective tower files).
+* No imports from any tower file — this module is self-contained with respect to
+  the tower layer.  Tower files import from here; this file does not import from them.
+* Removing any tower file leaves this module fully intact.
+* Duck-typing via optional TowerBase methods (e.g. cd_warmup_embedding) replaces
+  isinstance checks so new tower types never require changes here.
+"""
+
+from __future__ import annotations
+
+import math
+from abc import ABC, abstractmethod
+from random import random
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+from v3.classes.backbones import BACKBONES, list_names, load_backbone_weights
+from v3.classes.bridges import Bridge
+from random import random
+
+
+# ---------------------------------------------------------------------------
+# Cross-tower communication context
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class EarlyPassContext:
+    eye_train: list[dict]
+    bilat_train: list[dict]
+    bilat_val: list[dict]
+    bilat_test: list[dict]
+    image_preprocessor: object
+    image_cache: object
+    device: torch.device
+    store: dict = field(default_factory=dict)  # cross-tower key-value store
+
+
+# ---------------------------------------------------------------------------
+# Backbone factory
+# ---------------------------------------------------------------------------
+
+
+def build_backbone(name: str, freeze_ratio: float = 0.0, augment: bool = True):
+    """
+    Operational builder:
+      - instantiate with DEFAULT weights
+      - strip classifier → features
+      - apply ratio-based freezing over coarse blocks
+      - return (model, out_dim, transform)
+    """
+    key = (name or "").lower()
+    if key not in BACKBONES:
+        raise ValueError(
+            f"Unsupported backbone '{name}'. Valid options: {list_names()}"
+        )
+
+    spec = BACKBONES[key]
+    m = spec.ctor(weights=spec.weights_default)
+    out_dim, m = spec.strip(m)
+    load_backbone_weights(key, m)
+
+    mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406))
+    std = getattr(spec.weights_default, "meta", {}).get("std", (0.229, 0.224, 0.225))
+    crop = 299 if key == "inception_v3" else 224
+
+    if augment:
+        transform = transforms.Compose(
+            [
+                transforms.Resize(256),
+                transforms.CenterCrop(crop),
+                transforms.RandomHorizontalFlip(),
+                transforms.RandomVerticalFlip(),
+                transforms.RandomRotation(15),
+                transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=mean, std=std),
+            ]
+        )
+    else:
+        transform = transforms.Compose(
+            [
+                transforms.Resize(256),
+                transforms.CenterCrop(crop),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=mean, std=std),
+            ]
+        )
+
+    fr = max(0.0, min(1.0, float(freeze_ratio)))
+    blocks = spec.blocks(m)
+    n = len(blocks)
+    freeze_n = int(math.floor(n * fr))
+    for b in blocks[:freeze_n]:
+        for p in b.parameters():
+            p.requires_grad = False
+
+    return m, out_dim, transform
+
+
+# ---------------------------------------------------------------------------
+# TowerBase ABC
+# ---------------------------------------------------------------------------
+
+
+class TowerBase(ABC):
+    """Abstract base class for a HyperTower tower.
+
+    Concrete sub-classes must implement ``embed_dims``, ``embed_batch``, and
+    ``prepare_fold``.  Everything else has a sensible default no-op.
+    """
+
+    # ------------------------------------------------------------------
+    # Required interface
+    # ------------------------------------------------------------------
+
+    @property
+    @abstractmethod
+    def embed_dims(self) -> list[int]:
+        """Ordered list of embedding dimensionalities contributed to the bridge.
+
+        Most towers contribute one embedding (e.g. GeometryTower → [geom_hidden]).
+        ImageClinicalTower contributes two (image + clinical → [img_dim, cd_dim]).
+        """
+
+    @abstractmethod
+    def embed_batch(
+        self,
+        batch: dict,
+        *,
+        device: torch.device,
+        slot: int = 1,
+    ) -> list[torch.Tensor]:
+        """Return a list of embeddings for one eye slot in *batch*.
+
+        Parameters
+        ----------
+        batch  : dict  — batch produced by a SlotDataset loader
+        device : torch.device
+        slot   : 1 (OD / image_1 / matrix_1) or 2 (OS / image_2 / matrix_2)
+
+        Returns
+        -------
+        list of Tensor — same length and order as ``embed_dims``
+        """
+
+    @abstractmethod
+    def prepare_fold(
+        self,
+        *,
+        eye_train: list,
+        bilat_train: list,
+        bilat_val: list,
+        bilat_test: list,
+        image_preprocessor,
+        image_cache,
+        device: torch.device,
+        args,
+    ) -> None:
+        """Called once per fold before the main epoch loop."""
+
+    def early_pass(self, context: EarlyPassContext) -> None:
+        """Optional: called once per fold before loaders are built."""
+        pass
+
+    def get_sample(self, entry) -> "torch.Tensor | dict":  # noqa: ARG002
+        """Optional: called by HTDataset to retrieve one sample for this tower.
+
+        entry : ShellEntry — carries entity_id, label, side, paired.
+
+        Single mode (entry.paired=False): return one Tensor.
+        Paired mode (entry.paired=True):  return {"a": Tensor, "b": Tensor}.
+
+        Default raises NotImplementedError.  Towers that participate in the
+        v4 HTDataset pipeline must implement this.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__}.get_sample() is not implemented. "
+            "Implement it to use this tower with HTDataset."
+        )
+
+    # ------------------------------------------------------------------
+    # Optional interface
+    # ------------------------------------------------------------------
+
+    def augment_samples(self, samples: list) -> list:
+        """Optional: add modality-specific keys to sample dicts before loaders are built.
+
+        Called by the orchestrator on each sample list (eye_train, bilat_train,
+        bilat_val, bilat_test) *after* ``prepare_fold`` completes.
+
+        The default implementation is a no-op.  GeometryTower overrides this to
+        inject ``seg_map_1`` / ``seg_map_2`` numpy arrays so the loader can deliver
+        them as tensors alongside the image and clinical slots.
+        """
+        return samples
+
+    def cd_warmup_embedding(
+        self,
+        batch: dict,
+        *,
+        device: torch.device,
+    ) -> Optional[torch.Tensor]:
+        """Return the clinical embedding for cd_warmup phase, or None if not applicable.
+
+        ClinicalDataTower overrides this to return its encoder output.
+        All other towers return None (the default).
+
+        This replaces isinstance(t, ClinicalDataTower) checks in train_towers_epoch,
+        so new tower types never require changes to towerbase.py.
+        """
+        return None
+
+    def set_phase(self, phase: str) -> None:
+        """Control requires_grad on this tower's parameters for *phase*.
+
+        Phases: ``cd_warmup``, ``tower_warmup``, ``fused_warmup``, ``main``.
+        Default: no-op (tower parameters always trainable unless overridden).
+        """
+
+    @property
+    def total_epochs(self) -> int:
+        """How many epochs this tower participates in the main loop."""
+        return 0
+
+    @property
+    def train_loader(self) -> Optional[DataLoader]:
+        """Single-eye training loader, or None if not applicable."""
+        return None
+
+    @property
+    def cd_only_loader(self) -> Optional[DataLoader]:
+        """Clinical-data-only loader for cd_warmup phase, or None."""
+        return None
+
+    def finalize_fold(
+        self,
+        *,
+        bridge,
+        bilat_train_loader: Optional[DataLoader] = None,
+        val_loader: Optional[DataLoader] = None,
+        device: torch.device,
+        args,
+    ) -> None:
+        """Optional post-epoch-loop operations (e.g. fused head training)."""
+
+
+# ---------------------------------------------------------------------------
+# Shared training utilities
+# ---------------------------------------------------------------------------
+
+
+def _set_requires_grad(module: nn.Module, enabled: bool) -> None:
+    for p in module.parameters():
+        p.requires_grad = enabled
+
+
+def _to_label_tensor(labels, device: torch.device) -> torch.Tensor:
+    if torch.is_tensor(labels):
+        return labels.to(device=device, dtype=torch.long)
+    return torch.as_tensor(labels, dtype=torch.long, device=device)
+
+
+# ---------------------------------------------------------------------------
+# Modular tower training / evaluation (multi-tower interface)
+# ---------------------------------------------------------------------------
+
+
+def train_towers_epoch(
+    towers: "list[TowerBase]",
+    bridge: Bridge,
+    loader: DataLoader,
+    optimizer,
+    device: torch.device,
+    *,
+    phase: str,
+    bcd_prob: float = 0.5,
+    tower_loss_mode: str = "bcd",
+) -> tuple[float, float]:
+    """Train one epoch using the modular tower interface.
+
+    All towers and the bridge are set to the given phase via their
+    ``set_phase`` methods.  BCD / fused loss semantics mirror the
+    existing ``train_single_epoch`` logic:
+
+      cd_warmup    — clinical encoder aux head only (slot index 1)
+      tower_warmup — img + cd aux heads equally
+      fused_warmup — fused bridge output only
+      main         — BCD (randomly img or cd aux) vs fused, per tower_loss_mode
+
+    Duck typing: towers that implement ``cd_warmup_embedding`` participate in
+    cd_warmup; all others are skipped for that phase.  No isinstance checks.
+    """
+    for t in towers:
+        t.set_phase(phase)
+    bridge.set_phase(phase)
+
+    total_loss = total_correct = total_n = 0
+
+    for batch in loader:
+        y = batch.get("label_1")
+        if y is None:
+            continue
+
+        # ---- cd_warmup: train clinical encoder via its aux head ----
+        if phase == "cd_warmup":
+            z_cd = None
+            for t in towers:
+                z = t.cd_warmup_embedding(batch, device=device)
+                if z is not None:
+                    z_cd = z
+                    break
+            if z_cd is None:
+                continue
+            logits = bridge.aux_heads[1](z_cd)
+            y_t = _to_label_tensor(y, device)
+            loss = F.cross_entropy(logits, y_t)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            bs = y_t.shape[0]
+            total_loss += float(loss.item()) * bs
+            total_correct += int((logits.argmax(1) == y_t).sum())
+            total_n += bs
+            continue
+
+        # ---- collect all slot-1 embeddings from every tower ----
+        all_embs = []
+        for t in towers:
+            all_embs.extend(t.embed_batch(batch, device=device, slot=1))
+
+        if any(e is None for e in all_embs):
+            continue
+
+        y_t = _to_label_tensor(y, device)
+
+        if phase == "tower_warmup":
+            # Average loss across all available aux heads
+            aux_logits = [bridge.aux_heads[i](emb) for i, emb in enumerate(all_embs)]
+            loss = sum(F.cross_entropy(l, y_t) for l in aux_logits) / len(aux_logits)
+            # Softmax average for metrics
+            logits = sum(F.softmax(l, dim=1) for l in aux_logits) / len(aux_logits)
+
+        elif phase == "fused_warmup":
+            logits_fused, _ = bridge.fuse(all_embs)
+            loss = F.cross_entropy(logits_fused, y_t)
+            logits = logits_fused
+
+        else:  # main
+            if tower_loss_mode == "all":
+                logits_fused, aux = bridge.fuse(all_embs)
+                loss = F.cross_entropy(logits_fused, y_t)
+                for aux_l in aux:
+                    loss = loss + F.cross_entropy(aux_l, y_t)
+                logits = logits_fused
+            elif random() < bcd_prob:
+                # Randomly pick ONE tower to train (Generalized BCD)
+                idx = int(random() * len(all_embs))
+                logits = bridge.aux_heads[idx](all_embs[idx])
+                loss = F.cross_entropy(logits, y_t)
+            else:
+                logits_fused, _ = bridge.fuse(all_embs)
+                loss = F.cross_entropy(logits_fused, y_t)
+                logits = logits_fused
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        bs = y_t.shape[0]
+        total_loss += float(loss.item()) * bs
+        total_correct += int((logits.argmax(1) == y_t).sum())
+        total_n += bs
+
+    return (
+        total_loss / total_n if total_n else float("nan"),
+        total_correct / total_n if total_n else float("nan"),
+    )
+
+
+def collect_probs_towers(
+    towers: "list[TowerBase]",
+    bridge: Bridge,
+    loader: DataLoader,
+    device: torch.device,
+    *,
+    tower_mode: str = "ensemble",
+) -> tuple[np.ndarray, np.ndarray]:
+    """Evaluate using the modular tower interface.
+
+    Returns ``(y_true, probs)`` with patient-level probabilities
+    (OD + OS averaged for ensemble mode, OD-only for single mode).
+    """
+    for t in towers:
+        if isinstance(t, nn.Module):
+            t.eval()
+    bridge.eval()
+
+    y_chunks, p_chunks = [], []
+    with torch.no_grad():
+        for batch in loader:
+            y = batch.get("label_1")
+            if y is None:
+                continue
+            if not (
+                torch.is_tensor(batch.get("image_1"))
+                and torch.is_tensor(batch.get("image_2"))
+            ):
+                continue
+
+            all_embs_od = []
+            all_embs_os = []
+            for t in towers:
+                all_embs_od.extend(t.embed_batch(batch, device=device, slot=1))
+                all_embs_os.extend(t.embed_batch(batch, device=device, slot=2))
+
+            if any(e is None for e in all_embs_od + all_embs_os):
+                continue
+
+            logits_od, _ = bridge.fuse(all_embs_od)
+            logits_os, _ = bridge.fuse(all_embs_os)
+
+            if tower_mode == "ensemble":
+                probs = 0.5 * (
+                    F.softmax(logits_od, dim=1) + F.softmax(logits_os, dim=1)
+                )
+            else:
+                probs = F.softmax(logits_od, dim=1)
+
+            y_chunks.append(_to_label_tensor(y, device).cpu().numpy())
+            p_chunks.append(probs.cpu().numpy())
+
+    for t in towers:
+        if isinstance(t, nn.Module):
+            t.train()
+    bridge.train()
+
+    if not y_chunks:
+        return np.array([], dtype=np.int64), np.zeros((0, 0), dtype=np.float32)
+    return np.concatenate(y_chunks), np.concatenate(p_chunks, axis=0)

v3/classes/towers.py

@@ -1,279 +0,0 @@
-from __future__ import annotations
-
-import math
-from typing import Optional
-
-import torch
-from torch import nn
-from torchvision import transforms
-
-from v3.classes.backbones import BACKBONES, list_names, load_backbone_weights
-from v3.classes.SE_attention import SEBlock
-from v3.classes.data_bundle import DataBundle
-
-
-def build_backbone(name: str, freeze_ratio: float = 0.0, augment: bool = True):
-    """
-    Operational builder:
-      - instantiate with DEFAULT weights
-      - strip classifier → features
-      - apply ratio-based freezing over coarse blocks
-      - return (model, out_dim, transform)
-    """
-    key = (name or "").lower()
-    if key not in BACKBONES:
-        raise ValueError(f"Unsupported backbone '{name}'. Valid options: {list_names()}")
-
-    spec = BACKBONES[key]
-    m = spec.ctor(weights=spec.weights_default)
-    out_dim, m = spec.strip(m)
-    load_backbone_weights(key, m)
-
-    # transforms: use the weights’ mean/std, but keep your augmentation pipeline
-    mean = getattr(spec.weights_default, "meta", {}).get("mean", (0.485, 0.456, 0.406))
-    std = getattr(spec.weights_default, "meta", {}).get("std", (0.229, 0.224, 0.225))
-    crop = 299 if key == "inception_v3" else 224
-
-    if augment:
-        transform = transforms.Compose(
-            [
-                transforms.Resize(256),
-                transforms.CenterCrop(crop),
-                transforms.RandomHorizontalFlip(),
-                transforms.RandomVerticalFlip(),
-                transforms.RandomRotation(15),
-                transforms.ColorJitter(0.1, 0.1, 0.1, 0.05),
-                transforms.ToTensor(),
-                transforms.Normalize(mean=mean, std=std),
-            ]
-        )
-    else:
-        transform = transforms.Compose(
-            [
-                transforms.Resize(256),
-                transforms.CenterCrop(crop),
-                transforms.ToTensor(),
-                transforms.Normalize(mean=mean, std=std),
-            ]
-        )
-
-    # ratio-based freezing: freeze earliest floor(N * freeze_ratio) blocks
-    fr = max(0.0, min(1.0, float(freeze_ratio)))
-    blocks = spec.blocks(m)
-    n = len(blocks)
-    freeze_n = int(math.floor(n * fr))
-    for b in blocks[:freeze_n]:
-        for p in b.parameters():
-            p.requires_grad = False
-
-    return m, out_dim, transform
-
-
-class ImageTower(nn.Module):
-    """
-    Vision backbone → pooled features.
-    - backbone: one of list_names() (default 'efficientnet_b0')
-    - always DEFAULT torchvision weights
-    - freeze_ratio ∈ [0,1] freezes earliest floor(N*freeze_ratio) blocks
-    - returns [N, out_dim] features from backbone forward
-    """
-
-    def __init__(
-        self,
-        backbone: str = "efficientnet_b0",
-        freeze_ratio: float = 0.0,
-        use_se: bool = False,
-        se_reduction: int = 16,
-        se_pre_norm: bool = True,
-        augment: bool = True,
-        geometry_dim: int = 0,
-    ):
-        super().__init__()
-        self.backbone, base_dim, self.transform = build_backbone(
-            backbone, freeze_ratio, augment=augment
-        )
-        self._name = backbone
-        # Keep ordered blocks for dynamic freezing/thawing
-        key = (self._name or "").lower()
-        self._spec = BACKBONES[key]
-        self._blocks = self._spec.blocks(self.backbone)
-        # Optional tower-level SE over the final feature vector
-        self.base_dim = base_dim
-        self.geometry_dim = max(0, int(geometry_dim))
-        self.out_dim = self.base_dim + self.geometry_dim
-        self.tower_ln = nn.LayerNorm(self.base_dim) if se_pre_norm else nn.Identity()
-        self.tower_se = (
-            SEBlock(self.base_dim, reduction=se_reduction, residual=True)
-            if use_se
-            else None
-        )
-
-    def forward(
-        self, x: torch.Tensor, geometry: Optional[torch.Tensor] = None
-    ) -> torch.Tensor:
-        y = self.backbone(x)
-        # sanity: pooled features, not logits
-        assert y.dim() == 2 and y.size(1) == self.base_dim, (
-            f"Expected features [N,{self.base_dim}], got {tuple(y.shape)}"
-        )
-        if self.tower_se is not None:
-            y, _ = self.tower_se(self.tower_ln(y))
-        if self.geometry_dim > 0:
-            if geometry is None or geometry.numel() == 0:
-                geom = torch.zeros(
-                    y.size(0), self.geometry_dim, device=y.device, dtype=y.dtype
-                )
-            else:
-                if geometry.dim() == 1:
-                    geom = geometry.unsqueeze(0)
-                else:
-                    geom = geometry
-                geom = geom.to(device=y.device, dtype=y.dtype)
-                if geom.size(0) != y.size(0):
-                    raise ValueError(
-                        f"Geometry batch size mismatch: {geom.size(0)} vs {y.size(0)}"
-                    )
-                if geom.size(1) != self.geometry_dim:
-                    raise ValueError(
-                        f"Expected geometry dim {self.geometry_dim}, got {geom.size(1)}"
-                    )
-            y = torch.cat([y, geom], dim=1)
-        return y
-
-    def set_freeze_ratio(self, ratio: float):
-        """Dynamically freeze earliest floor(N*ratio) backbone blocks."""
-        r = max(0.0, min(1.0, float(ratio)))
-        n = len(self._blocks)
-        freeze_n = int(math.floor(n * r))
-        # Unfreeze all first
-        for b in self._blocks:
-            for p in b.parameters():
-                p.requires_grad = True
-        # Freeze earliest blocks
-        for b in self._blocks[:freeze_n]:
-            for p in b.parameters():
-                p.requires_grad = False
-
-
-class SiameseImageTower(nn.Module):
-    """
-    Shared-weight bilateral image tower.
-
-    Runs OD and OS images through a single shared backbone, then returns
-    cat([f_mean, f_delta]) where:
-        f_mean  = (f_od + f_os) / 2   -- shared bilateral representation
-        f_delta = f_od - f_os          -- asymmetry, signed OD-relative
-
-    out_dim = 2 * backbone_out_dim
-
-    When x_os is None (single-eye fallback):
-        f_mean  = f_od
-        f_delta = zeros
-    so the module degrades gracefully when only one eye is available.
-
-    The shared backbone means both eyes contribute to every gradient update,
-    effectively doubling the training signal for the visual pathway without
-    doubling parameters.
-    """
-
-    def __init__(
-        self,
-        backbone: str = "efficientnet_b0",
-        freeze_ratio: float = 0.0,
-        use_se: bool = False,
-        se_reduction: int = 16,
-        se_pre_norm: bool = True,
-        augment: bool = True,
-    ):
-        super().__init__()
-        self._tower = ImageTower(
-            backbone=backbone,
-            freeze_ratio=freeze_ratio,
-            use_se=use_se,
-            se_reduction=se_reduction,
-            se_pre_norm=se_pre_norm,
-            augment=augment,
-            geometry_dim=0,
-        )
-        self.out_dim = self._tower.out_dim * 2
-        self.transform = self._tower.transform
-
-    def forward(
-        self,
-        x_od: torch.Tensor,
-        x_os: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        f_od = self._tower(x_od)
-        if x_os is None:
-            f_mean = f_od
-            f_delta = torch.zeros_like(f_od)
-        else:
-            f_os = self._tower(x_os)
-            f_mean = (f_od + f_os) * 0.5
-            f_delta = f_od - f_os
-        return torch.cat([f_mean, f_delta], dim=1)
-
-    def set_freeze_ratio(self, ratio: float) -> None:
-        """Delegates to the shared inner tower."""
-        self._tower.set_freeze_ratio(ratio)
-
-
-class ClinicalTower(nn.Module):
-    """MLP over DataBundle.vectorize_row outputs (convert to torch inside tower)."""
-
-    def __init__(
-        self,
-        clinical_data: DataBundle,
-        hidden_dim: int = 128,
-        dropout: float = 0.1,
-        use_se: bool = False,
-        se_reduction: int = 16,
-        se_pre_norm: bool = True,
-    ):
-        super().__init__()
-        self.feature_dim = clinical_data.feature_dim
-        self.out_dim = hidden_dim
-        # two-block MLP so we can optionally freeze/thaw per block
-        self.block0 = nn.Sequential(
-            nn.Linear(self.feature_dim, hidden_dim),
-            nn.LayerNorm(hidden_dim),
-            nn.ReLU(inplace=True),
-            nn.Dropout(dropout),
-        )
-        self.block1 = nn.Sequential(
-            nn.Linear(hidden_dim, hidden_dim),
-            nn.ReLU(inplace=True),
-        )
-        self.net = nn.Sequential(self.block0, self.block1)
-        self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity()
-        self.tower_se = (
-            SEBlock(hidden_dim, reduction=se_reduction, residual=True)
-            if use_se
-            else None
-        )
-
-    def forward(self, meta_np_or_torch) -> torch.Tensor:
-        if isinstance(meta_np_or_torch, torch.Tensor):
-            x = meta_np_or_torch
-        else:
-            x = torch.as_tensor(meta_np_or_torch, dtype=torch.float32)
-        h = self.net(x)
-        if self.tower_se is not None:
-            h, _ = self.tower_se(self.tower_ln(h))
-        return h
-
-    def set_freeze_ratio(self, ratio: float):
-        """Optionally freeze earliest blocks of the MLP."""
-        r = max(0.0, min(1.0, float(ratio)))
-        # Unfreeze all
-        for p in self.block0.parameters():
-            p.requires_grad = True
-        for p in self.block1.parameters():
-            p.requires_grad = True
-        # Freeze earliest blocks based on ratio threshold
-        if r >= 0.5:
-            for p in self.block0.parameters():
-                p.requires_grad = False
-        if r >= 1.0:
-            for p in self.block1.parameters():
-                p.requires_grad = False

v3/classes/v3_hypertower.py

@@ -29,7 +29,6 @@ from v3.classes.croppers import (
     build_image_preprocessor_from_args,
 )
 from v3.classes.image_loader import CachedImageLoader
-from v3.classes.dataset import _ClinicalView  # noqa: F401
 from v3.classes.loader_factory import (
     build_balanced_sampler,
     filter_bilateral_samples,
@@ -37,7 +36,12 @@ from v3.classes.loader_factory import (
     make_loader,
 )
 from v3.classes.metrics import _score_arrays, _svf, _tune_and_snap
-from v3.classes.models import (
+from v3.classes.bridges import Bridge
+from v3.classes.towerbase import train_towers_epoch, collect_probs_towers
+from v3.classes.image_towers import ImageTower
+from v3.classes.clinical_towers import ClinicalDataTower
+from v3.classes.geometry_towers import GeometryTower
+from v3.classes.hypertower_models import (
     BilateralHT,
     EmbeddingMLPEnsembleHT,
     FusedEnsembleHT,
@@ -154,7 +158,7 @@ class V3HyperTower:
         ap.add_argument("--exclude-cols", nargs="*", default=[])
         ap.add_argument("--eval-mode",    choices=["binary", "multiclass"], default="binary")
         ap.add_argument(
-            "--tower-mode", choices=["single", "ensemble", "bilateral", "siamese", "classic"],
+            "--hypertower-mode", choices=["single", "ensemble", "bilateral", "siamese", "classic"],
             default="ensemble",
         )
         ap.add_argument("--n-splits",  type=int, default=5)
@@ -252,6 +256,20 @@ class V3HyperTower:
         ap.add_argument("--geometry-source", default="gt", choices=["gt", "unet"],
                         help="Source for geometry features: gt (GT contour annotations) or "
                              "unet (U-Net segmentation). unet also requires --img-crop-weights.")
+        ap.add_argument("--geometry-tower", action="store_true",
+                        help="Add a dedicated GeometryTower (disc/cup seg-map CNN) fused via the "
+                             "bridge alongside ImageTower and ClinicalDataTower. Requires "
+                             "--img-crop-manifest.")
+        ap.add_argument("--geometry-tower-backbone", default="resnet18",
+                        choices=["resnet18", "resnet50", "efficientnet_b0"],
+                        help="SegCNN backbone for GeometryTower (default: resnet18).")
+        ap.add_argument("--geometry-tower-in-channels", type=int, default=3, choices=[1, 3],
+                        help="1 = single label map; 3 = one-hot disc/rim/cup (default: 3).")
+        ap.add_argument("--geometry-tower-frozen", action="store_true",
+                        help="Freeze GeometryTower backbone throughout training.")
+        ap.add_argument("--geometry-tower-finetune-unet-epochs", type=int, default=0,
+                        help="Epochs to fine-tune the U-Net per fold before seg-map extraction "
+                             "(0 = disabled; only applies when --geometry-source unet).")
         return ap
 
     def __init__(self, args) -> None:
@@ -333,7 +351,7 @@ class V3HyperTower:
         out_dir.mkdir(parents=True, exist_ok=True)
 
         mode       = args.eval_mode
-        tower_mode = "single" if args.tower_mode == "classic" else args.tower_mode
+        tower_mode = "single" if args.hypertower_mode == "classic" else args.hypertower_mode
         df_mode    = self.data.df.copy()
 
         if args.exclude_mixed_patients:
@@ -482,6 +500,300 @@ class V3HyperTower:
 
         return out_dir
 
+    def _run_fold_towers(
+        self,
+        *,
+        fold: int,
+        split,
+        mode: str,
+        data,
+        num_classes: int,
+        profile_eye,
+        profile_patient,
+        fold_dir: Path,
+        pred_store,
+        image_cache,
+    ):
+        """Modular TowerBase training path (used when --geometry-tower is set).
+
+        Builds [ImageTower, ClinicalDataTower, GeometryTower], runs the full
+        fold lifecycle (prepare_fold → augment_samples → loader build → epoch
+        loop → eval), and returns (FoldResult, FoldArtifacts) with metrics in
+        the ensemble_val_* slots.
+        """
+        args   = self.args
+        device = self.device
+        nan    = float("nan")
+
+        # ------------------------------------------------------------------ samples
+        eye_train   = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data))
+        bilat_train = filter_bilateral_samples(profile_patient.build_samples(df=split.train, clinical=data))
+        bilat_val   = filter_bilateral_samples(profile_patient.build_samples(df=split.val,   clinical=data))
+        bilat_test  = filter_bilateral_samples(profile_patient.build_samples(
+            df=split.test, clinical=data)) if split.test is not None else []
+
+        # Old --geometry-dim path still applies (injects geometry into clinical stream)
+        if self.geometry_provider is not None:
+            eye_train   = self._augment_geometry(eye_train)
+            bilat_train = self._augment_geometry(bilat_train)
+            bilat_val   = self._augment_geometry(bilat_val)
+            bilat_test  = self._augment_geometry(bilat_test)
+
+        if len(bilat_val) == 0:
+            empty = FoldResult(
+                mode=mode, fold=fold,
+                best_epoch_single=0, best_epoch_bilat=0,
+                classic_val_auc=nan,  classic_val_acc=nan,  classic_val_kappa=nan,
+                classic_val_mcc=nan,  classic_val_f1=nan,   classic_val_recall=None,
+                classic_val_ece=nan,  classic_val_threshold=nan, classic_val_bias=None,
+                classic_val_n=0,
+                ensemble_val_auc=nan, ensemble_val_acc=nan, ensemble_val_kappa=nan,
+                ensemble_val_mcc=nan, ensemble_val_f1=nan,  ensemble_val_recall=None,
+                ensemble_val_ece=nan, ensemble_val_threshold=nan, ensemble_val_bias=None,
+                ensemble_val_n=0,
+                bilat_val_auc=nan,  bilat_val_acc=nan,  bilat_val_kappa=nan,
+                bilat_val_mcc=nan,  bilat_val_f1=nan,   bilat_val_recall=None,
+                bilat_val_ece=nan,  bilat_val_threshold=nan, bilat_val_bias=None,
+                bilat_val_n=0,
+                single_train_n=len(eye_train), bilat_train_n=len(bilat_train),
+            )
+            return empty, FoldArtifacts(
+                y_true_classic=None, probs_classic=None,
+                y_true_ensemble=None, probs_ensemble=None,
+                y_true_bilat=None, probs_bilat=None,
+            )
+
+        # ------------------------------------------------------------------ towers
+        img_tower = ImageTower(
+            backbone=args.backbone,
+            freeze_ratio=args.freeze_ratio,
+            augment=args.augment,
+            use_se=getattr(args, "se_img_tower", False),
+        )
+        cd_tower = ClinicalDataTower(
+            clinical_data=data,
+            cd_hidden_dim=args.cd_hidden_dim,
+            cd_dropout=getattr(args, "cd_dropout", 0.1),
+            use_se=getattr(args, "se_cd_tower", False),
+        )
+
+        geom_tower = GeometryTower(
+            backbone=getattr(args, "geometry_tower_backbone", "resnet18"),
+            in_channels=getattr(args, "geometry_tower_in_channels", 3),
+            pretrained=not getattr(args, "no_pretrained", False),
+            frozen=getattr(args, "geometry_tower_frozen", False),
+            geometry_source=getattr(args, "geometry_source", "gt"),
+            manifest_path=getattr(args, "img_crop_manifest", None),
+            weights_path=getattr(args, "img_crop_weights", None),
+            unet_normalize=getattr(args, "img_crop_normalize", "per_image"),
+            unet_threshold=getattr(args, "img_crop_threshold", 0.5),
+            finetune_unet_epochs=getattr(args, "geometry_tower_finetune_unet_epochs", 0),
+        )
+
+        # GeometryTower.prepare_fold must run before augment_samples (precomputes seg maps)
+        geom_tower.prepare_fold(
+            eye_train=eye_train, bilat_train=bilat_train,
+            bilat_val=bilat_val, bilat_test=bilat_test,
+            image_preprocessor=self.image_preprocessor,
+            image_cache=image_cache, device=device, args=args,
+        )
+        # Inject seg_map_1/seg_map_2 into all sample lists before loaders are built
+        for sample_list in (eye_train, bilat_train, bilat_val, bilat_test):
+            geom_tower.augment_samples(sample_list)
+
+        # Now ImageTower.prepare_fold sees augmented samples → loader includes seg maps
+        img_tower.prepare_fold(
+            eye_train=eye_train, bilat_train=bilat_train,
+            bilat_val=bilat_val, bilat_test=bilat_test,
+            image_preprocessor=self.image_preprocessor,
+            image_cache=image_cache, device=device, args=args,
+        )
+        cd_tower.prepare_fold(
+            eye_train=eye_train, bilat_train=bilat_train,
+            bilat_val=bilat_val, bilat_test=bilat_test,
+            image_preprocessor=self.image_preprocessor,
+            image_cache=image_cache, device=device, args=args,
+        )
+
+        towers = [img_tower, cd_tower, geom_tower]
+
+        # ------------------------------------------------------------------ bridge
+        tower_dims = []
+        for t in towers:
+            tower_dims.extend(t.embed_dims)
+        bridge = Bridge(
+            tower_dims=tower_dims,
+            num_classes=num_classes,
+            fusion_dim=args.fusion_dim,
+            mode=getattr(args, "bridge_mode", "fused"),
+            dropout=getattr(args, "bridge_dropout", 0.5),
+        )
+
+        # Move all nn.Modules to device
+        for t in towers:
+            if isinstance(t, torch.nn.Module):
+                t.to(device)
+        bridge.to(device)
+
+        # ------------------------------------------------------------------ loaders
+        slots_patient = profile_patient.slot_descriptors()
+        _persistent = args.num_workers > 0
+        loader_kw = dict(
+            batch_size=args.batch_size, num_workers=args.num_workers,
+            image_cache=image_cache, persistent_workers=_persistent,
+        )
+        eval_transform = build_eval_transform(args.backbone)
+        val_loader = make_loader(
+            bilat_val, slots_patient,
+            image_transform=eval_transform,
+            image_preprocessor=self.image_preprocessor,
+            shuffle=False, **loader_kw,
+        )
+        test_loader = None
+        if bilat_test:
+            test_loader = make_loader(
+                bilat_test, slots_patient,
+                image_transform=eval_transform,
+                image_preprocessor=self.image_preprocessor,
+                shuffle=False, **loader_kw,
+            )
+
+        train_loader = img_tower.train_loader
+        val_loader.dataset.prebuild_image_cache()
+        if test_loader is not None:
+            test_loader.dataset.prebuild_image_cache()
+
+        # ------------------------------------------------------------------ optimizer
+        all_params = list(bridge.parameters())
+        for t in towers:
+            if isinstance(t, torch.nn.Module):
+                all_params.extend(t.parameters())
+        optimizer = torch.optim.AdamW(
+            [p for p in all_params if p.requires_grad],
+            lr=args.lr,
+            weight_decay=getattr(args, "weight_decay", 1e-4),
+        )
+
+        # ------------------------------------------------------------------ epoch loop
+        global_warmup_tower = getattr(args, "warmup_tower_epochs", None)
+        global_warmup_fused = getattr(args, "warmup_fused_epochs", None)
+        warmup_cd    = int(getattr(args, "warmup_cd_epochs", 0))
+        warmup_tower = int(getattr(args, "single_warmup_tower_epochs", None) or global_warmup_tower or 2)
+        warmup_fused = int(getattr(args, "single_warmup_fused_epochs", None) or global_warmup_fused or 2)
+        main_epochs  = int(args.epochs)
+
+        schedule = []
+        if warmup_cd    > 0: schedule.append(("cd_warmup",    warmup_cd))
+        if warmup_tower > 0: schedule.append(("tower_warmup", warmup_tower))
+        if warmup_fused > 0: schedule.append(("fused_warmup", warmup_fused))
+        schedule.append(("main", main_epochs))
+
+        best_val_auc      = float("-inf")
+        best_epoch        = 0
+        best_tower_states = None
+        best_bridge_state = None
+        epoch_idx         = 0
+
+        for phase, n_epochs in schedule:
+            for _ in range(n_epochs):
+                for t in towers:
+                    if isinstance(t, torch.nn.Module):
+                        t.train()
+                train_towers_epoch(
+                    towers, bridge, train_loader, optimizer, device,
+                    phase=phase,
+                    bcd_prob=getattr(args, "bcd_prob", 0.5),
+                    tower_loss_mode=getattr(args, "tower_loss_mode", "bcd"),
+                )
+                y_v, p_v = collect_probs_towers(towers, bridge, val_loader, device,
+                                                tower_mode="ensemble")
+                _, val_auc, _ = _score_arrays(y_v, p_v, num_classes)
+                if not np.isnan(val_auc) and val_auc > best_val_auc:
+                    best_val_auc      = val_auc
+                    best_epoch        = epoch_idx
+                    best_tower_states = [
+                        t.state_dict() if isinstance(t, torch.nn.Module) else None
+                        for t in towers
+                    ]
+                    best_bridge_state = bridge.state_dict()
+                epoch_idx += 1
+
+        # Restore best
+        if best_bridge_state is not None:
+            bridge.load_state_dict(best_bridge_state)
+        if best_tower_states is not None:
+            for t, st in zip(towers, best_tower_states):
+                if isinstance(t, torch.nn.Module) and st is not None:
+                    t.load_state_dict(st)
+
+        # ------------------------------------------------------------------ eval
+        y_val, p_val = collect_probs_towers(towers, bridge, val_loader, device, tower_mode="ensemble")
+        acc_val, auc_val, n_val = _score_arrays(y_val, p_val, num_classes)
+        snap_val, _, thr_val, bias_val = _tune_and_snap(
+            y_val, p_val, acc_val, num_classes, args, n_bins=10
+        )
+
+        y_test = p_test = None
+        test_auc = test_acc = nan
+        test_n = 0
+        if test_loader is not None:
+            y_test, p_test = collect_probs_towers(towers, bridge, test_loader, device,
+                                                  tower_mode="ensemble")
+            test_acc, test_auc, test_n = _score_arrays(y_test, p_test, num_classes)
+
+        result = FoldResult(
+            mode=mode, fold=fold,
+            best_epoch_single=best_epoch, best_epoch_bilat=0,
+            classic_val_auc=nan,  classic_val_acc=nan,  classic_val_kappa=nan,
+            classic_val_mcc=nan,  classic_val_f1=nan,   classic_val_recall=None,
+            classic_val_ece=nan,  classic_val_threshold=nan, classic_val_bias=None,
+            classic_val_n=0,
+            ensemble_val_auc=snap_val["auc"],  ensemble_val_acc=snap_val["acc"],
+            ensemble_val_kappa=snap_val["kappa"], ensemble_val_mcc=snap_val["mcc"],
+            ensemble_val_f1=snap_val["macro_f1"],
+            ensemble_val_recall=_sv(snap_val["per_class_recall"]),
+            ensemble_val_ece=snap_val["ece"],
+            ensemble_val_threshold=snap_val["threshold"],
+            ensemble_val_bias=_svf(bias_val),
+            ensemble_val_n=snap_val["n"],
+            bilat_val_auc=nan,  bilat_val_acc=nan,  bilat_val_kappa=nan,
+            bilat_val_mcc=nan,  bilat_val_f1=nan,   bilat_val_recall=None,
+            bilat_val_ece=nan,  bilat_val_threshold=nan, bilat_val_bias=None,
+            bilat_val_n=0,
+            ensemble_test_auc=test_auc, ensemble_test_acc=test_acc,
+            test_n=test_n,
+            single_train_n=len(eye_train), bilat_train_n=len(bilat_train),
+        )
+        artifacts = FoldArtifacts(
+            y_true_classic=None,   probs_classic=None,
+            y_true_ensemble=y_val, probs_ensemble=p_val,
+            y_true_bilat=None,     probs_bilat=None,
+            y_true_test=y_test,    probs_test=p_test,
+        )
+        return result, artifacts
+
+    def _augment_geometry_slot(self, samples: list) -> list:
+        """Add geom_1/geom_2 keys to each sample dict (geometry tower mode).
+
+        Unlike _augment_geometry, this does NOT touch matrix_1/matrix_2 — the
+        geometry vector lives in its own slot so ImageTower and ClinicalDataTower
+        each receive only their own modality.
+        """
+        if self.geometry_provider is None:
+            return samples
+        geom_dim = int(getattr(self.args, "geometry_dim", 0)) or 5
+        for s in samples:
+            for img_slot, geom_slot in (("image_1", "geom_1"), ("image_2", "geom_2")):
+                img_path = s.get(img_slot)
+                if img_path is None:
+                    continue
+                vec = self.geometry_provider.geometry_for_image(img_path)
+                if vec is not None and len(vec) >= geom_dim:
+                    s[geom_slot] = vec[:geom_dim].astype(np.float32)
+                else:
+                    s[geom_slot] = np.zeros(geom_dim, dtype=np.float32)
+        return samples
+
     def _augment_geometry(self, samples: list) -> list:
         """Append geometry features to matrix_1/matrix_2 in each sample dict."""
         if self.geometry_provider is None:
@@ -517,6 +829,16 @@ class V3HyperTower:
         image_cache,
     ):
         args = self.args
+
+        # Modular tower path — bypasses the legacy single/bilat/siamese code entirely
+        if getattr(args, "geometry_tower", False):
+            return self._run_fold_towers(
+                fold=fold, split=split, mode=mode, data=data,
+                num_classes=num_classes, profile_eye=profile_eye,
+                profile_patient=profile_patient, fold_dir=fold_dir,
+                pred_store=pred_store, image_cache=image_cache,
+            )
+
         device = self.device
         image_preprocessor = self.image_preprocessor
         nan = float("nan")

v3/distributed/server.py

@@ -10,6 +10,7 @@ Usage:
 Environment:
     HT_TOKEN  — fallback if --token is not passed
 """
+
 from __future__ import annotations
 
 import argparse
@@ -44,8 +45,8 @@ from .protocol import (
 
 _TOKEN: str = ""
 _DB_PATH: Path = Path("v3/distributed/jobs.db")
-_CLIENT_TTL: int = 120   # seconds before a client is considered gone
+_CLIENT_TTL: int = 120  # seconds before a client is considered gone
-_MAX_ATTEMPTS: int = 3   # max times a job is retried before being left as failed
+_MAX_ATTEMPTS: int = 3  # max times a job is retried before being left as failed
 
 _clients: dict[str, ClientInfo] = {}
 _clients_lock = threading.Lock()
@@ -60,11 +61,15 @@ def _reap_stale_clients():
         # Step 1: evict timed-out clients from registry
         with _clients_lock:
             stale = [
-                cid for cid, c in _clients.items()
+                cid
+                for cid, c in _clients.items()
                 if datetime.fromisoformat(c.last_seen).timestamp() < cutoff
             ]
             for cid in stale:
-                print(f"[server] reaped stale client {cid} ({_clients[cid].hostname})", flush=True)
+                print(
+                    f"[server] reaped stale client {cid} ({_clients[cid].hostname})",
+                    flush=True,
+                )
                 del _clients[cid]
             known_ids = set(_clients.keys())
 
@@ -78,15 +83,20 @@ def _reap_stale_clients():
                     conn.execute(
                         "UPDATE jobs SET state='pending', assigned_to=NULL, started_at=NULL "
                         "WHERE job_id=?",
-                        (row["job_id"],)
+                        (row["job_id"],),
                     )
-                    print(f"[server] re-queued job {row['job_id']} "
+                    print(
-                          f"(client {row['assigned_to']} unknown)", flush=True)
+                        f"[server] re-queued job {row['job_id']} "
+                        f"(client {row['assigned_to']} unknown)",
+                        flush=True,
+                    )
+
 
 # ──────────────────────────────────────────────────────────────
 # Database helpers
 # ──────────────────────────────────────────────────────────────
 
+
 @contextmanager
 def _db():
     conn = sqlite3.connect(str(_DB_PATH))
@@ -101,7 +111,8 @@ def _db():
 def _init_db():
     _DB_PATH.parent.mkdir(parents=True, exist_ok=True)
     with _db() as conn:
-        conn.execute("""
+        conn.execute(
+            """
             CREATE TABLE IF NOT EXISTS jobs (
                 job_id       TEXT PRIMARY KEY,
                 run_name     TEXT NOT NULL,
@@ -117,10 +128,16 @@ def _init_db():
                 error_msg    TEXT,
                 attempts     INTEGER NOT NULL DEFAULT 0
             )
-        """)
+        """
+        )
+        conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_state ON jobs(state)")
+        conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_priority ON jobs(priority)")
+        conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_run_name ON jobs(run_name)")
         # Add attempts column to existing DBs that predate this field
         try:
-            conn.execute("ALTER TABLE jobs ADD COLUMN attempts INTEGER NOT NULL DEFAULT 0")
+            conn.execute(
+                "ALTER TABLE jobs ADD COLUMN attempts INTEGER NOT NULL DEFAULT 0"
+            )
         except Exception:
             pass  # column already exists
         # Note: running jobs are NOT reset on startup — active clients will re-register
@@ -147,6 +164,7 @@ def _ensure_client(client_id: str, hostname: str = "", gpu_info: str = "") -> bo
         return True
     return False
 
+
 # ──────────────────────────────────────────────────────────────
 # FastAPI app
 # ──────────────────────────────────────────────────────────────
@@ -161,12 +179,20 @@ def _check_token(x_token: str = Header(...)):
 
 # ── Registration ──────────────────────────────────────────────
 
-@app.post("/register", response_model=RegisterResponse,
+
-          dependencies=[Depends(_check_token)])
+@app.post(
+    "/register", response_model=RegisterResponse, dependencies=[Depends(_check_token)]
+)
 def register(req: RegisterRequest, reuse_id: Optional[str] = None):
     with _clients_lock:
-        client_id = reuse_id if (reuse_id and reuse_id in _clients) else str(uuid.uuid4())[:8]
+        client_id = (
-        existing_status = _clients[client_id].status if client_id in _clients else StatusPush(state="idle")
+            reuse_id if (reuse_id and reuse_id in _clients) else str(uuid.uuid4())[:8]
+        )
+        existing_status = (
+            _clients[client_id].status
+            if client_id in _clients
+            else StatusPush(state="idle")
+        )
         _clients[client_id] = ClientInfo(
             client_id=client_id,
             hostname=req.hostname,
@@ -175,14 +201,16 @@ def register(req: RegisterRequest, reuse_id: Optional[str] = None):
             last_seen=_now(),
         )
     action = "re-registered" if reuse_id else "registered"
-    print(f"[server] {action} {client_id} ({req.hostname} | {req.gpu_info})", flush=True)
+    print(
+        f"[server] {action} {client_id} ({req.hostname} | {req.gpu_info})", flush=True
+    )
     return RegisterResponse(client_id=client_id)
 
 
 # ── Job polling ───────────────────────────────────────────────
 
-@app.post("/poll", response_model=PollResponse,
+
-          dependencies=[Depends(_check_token)])
+@app.post("/poll", response_model=PollResponse, dependencies=[Depends(_check_token)])
 def poll(client_id: str):
     with _clients_lock:
         needs_reregister = _ensure_client(client_id)
@@ -206,7 +234,10 @@ def poll(client_id: str):
             (client_id, job_id),
         )
         if cur.rowcount:
-            print(f"[server] reset {cur.rowcount} orphaned running job(s) for {client_id}", flush=True)
+            print(
+                f"[server] reset {cur.rowcount} orphaned running job(s) for {client_id}",
+                flush=True,
+            )
         conn.execute(
             "UPDATE jobs SET state='running', assigned_to=?, started_at=? WHERE job_id=?",
             (client_id, _now(), job_id),
@@ -231,6 +262,7 @@ def poll(client_id: str):
 
 # ── Status ────────────────────────────────────────────────────
 
+
 @app.post("/status/{client_id}", dependencies=[Depends(_check_token)])
 def push_status(client_id: str, status: StatusPush):
     with _clients_lock:
@@ -256,6 +288,7 @@ def get_client(client_id: str):
 
 # ── Job completion ────────────────────────────────────────────
 
+
 @app.post("/complete", dependencies=[Depends(_check_token)])
 def complete(result: JobResult):
     with _db() as conn:
@@ -274,14 +307,17 @@ def complete(result: JobResult):
                 run_name = run_row["run_name"]
                 remaining = conn.execute(
                     "SELECT COUNT(*) FROM jobs WHERE run_name=? AND state != 'done'",
-                    (run_name,)
+                    (run_name,),
                 ).fetchone()[0]
                 if remaining == 0:
                     total = conn.execute(
                         "SELECT COUNT(*) FROM jobs WHERE run_name=?", (run_name,)
                     ).fetchone()[0]
                     conn.execute("DELETE FROM jobs WHERE run_name=?", (run_name,))
-                    print(f"[server] run '{run_name}' complete ({total} jobs) — cleared", flush=True)
+                    print(
+                        f"[server] run '{run_name}' complete ({total} jobs) — cleared",
+                        flush=True,
+                    )
         else:
             row = conn.execute(
                 "SELECT attempts FROM jobs WHERE job_id=?", (result.job_id,)
@@ -293,21 +329,28 @@ def complete(result: JobResult):
                     "attempts=?, error_msg=? WHERE job_id=?",
                     (attempts, result.error_msg, result.job_id),
                 )
-                print(f"[server] job {result.job_id} failed (attempt {attempts}/{_MAX_ATTEMPTS}), "
+                print(
-                      f"re-queuing", flush=True)
+                    f"[server] job {result.job_id} failed (attempt {attempts}/{_MAX_ATTEMPTS}), "
+                    f"re-queuing",
+                    flush=True,
+                )
             else:
                 conn.execute(
                     "UPDATE jobs SET state='failed', completed_at=?, attempts=?, error_msg=? "
                     "WHERE job_id=?",
                     (_now(), attempts, result.error_msg, result.job_id),
                 )
-                print(f"[server] job {result.job_id} failed permanently after "
+                print(
-                      f"{attempts} attempts", flush=True)
+                    f"[server] job {result.job_id} failed permanently after "
+                    f"{attempts} attempts",
+                    flush=True,
+                )
     return {"ok": True}
 
 
 # ── Job queue management ──────────────────────────────────────
 
+
 @app.post("/jobs", dependencies=[Depends(_check_token)])
 def submit_job(job: JobSubmit):
     job_id = str(uuid.uuid4())[:12]
@@ -316,8 +359,15 @@ def submit_job(job: JobSubmit):
             "INSERT INTO jobs "
             "(job_id, run_name, module, args, output_dir, priority, created_at) "
             "VALUES (?,?,?,?,?,?,?)",
-            (job_id, job.run_name, job.module, json.dumps(job.args),
+            (
-             job.output_dir, job.priority, _now()),
+                job_id,
+                job.run_name,
+                job.module,
+                json.dumps(job.args),
+                job.output_dir,
+                job.priority,
+                _now(),
+            ),
         )
     print(f"[server] queued {job_id} ({job.run_name})", flush=True)
     return {"job_id": job_id}
@@ -347,7 +397,9 @@ def clear_jobs(body: dict):
         else:
             states = body.get("states", ["done", "failed", "cancelled"])
             placeholders = ",".join("?" * len(states))
-            cur = conn.execute(f"DELETE FROM jobs WHERE state IN ({placeholders})", states)
+            cur = conn.execute(
+                f"DELETE FROM jobs WHERE state IN ({placeholders})", states
+            )
     print(f"[server] cleared {cur.rowcount} jobs", flush=True)
     return {"cleared": cur.rowcount}
 
@@ -366,19 +418,33 @@ def cancel_job(job_id: str):
 # Entry point
 # ──────────────────────────────────────────────────────────────
 
+
 def main():
-    ap = argparse.ArgumentParser(description=__doc__,
+    ap = argparse.ArgumentParser(
-                                  formatter_class=argparse.RawDescriptionHelpFormatter)
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
     ap.add_argument("--port", type=int, default=8765)
     ap.add_argument("--host", default="0.0.0.0")
-    ap.add_argument("--token", default=os.environ.get("HT_TOKEN", ""),
+    ap.add_argument(
-                    help="Shared secret (or set HT_TOKEN env var)")
+        "--token",
-    ap.add_argument("--db", default="v3/distributed/jobs.db",
+        default=os.environ.get("HT_TOKEN", ""),
-                    help="Path to SQLite job database")
+        help="Shared secret (or set HT_TOKEN env var)",
-    ap.add_argument("--client-ttl", type=int, default=120,
+    )
-                    help="Seconds of silence before a client is reaped (default: 120)")
+    ap.add_argument(
-    ap.add_argument("--max-attempts", type=int, default=3,
+        "--db", default="v3/distributed/jobs.db", help="Path to SQLite job database"
-                    help="Max times a failed job is retried before being left as failed (default: 3)")
+    )
+    ap.add_argument(
+        "--client-ttl",
+        type=int,
+        default=120,
+        help="Seconds of silence before a client is reaped (default: 120)",
+    )
+    ap.add_argument(
+        "--max-attempts",
+        type=int,
+        default=3,
+        help="Max times a failed job is retried before being left as failed (default: 3)",
+    )
     args = ap.parse_args()
 
     if not args.token:
@@ -394,7 +460,10 @@ def main():
     reaper = threading.Thread(target=_reap_stale_clients, daemon=True)
     reaper.start()
 
-    print(f"[server] listening on {args.host}:{args.port}  client_ttl={_CLIENT_TTL}s", flush=True)
+    print(
+        f"[server] listening on {args.host}:{args.port}  client_ttl={_CLIENT_TTL}s",
+        flush=True,
+    )
     uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
 
 

v3/scripts/development/analyze_emergent_features.py

@@ -0,0 +1,256 @@
+import sys
+from pathlib import Path
+
+# Make sure we can import from v3 classes when running directly
+sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+
+from v3.classes.hypertower_models import SingleEyeHT
+from v3.classes.papila_builders import build_papila_data
+from v3.classes.profiles import build_papila_profile
+from v3.classes.loader_factory import make_loader, filter_eye_samples
+from v3.classes.transforms import build_eval_transform
+from v3.classes.utils import choose_device
+
+
+def extract_emergent_features(model, loader, device, num_classes=2):
+    """
+    Finds samples where the fused bridge is correct, but both individual
+    towers are wrong, and extracts the driving features from the fusion_dim.
+
+    Assumes `model` is a SingleEyeHT. Can be adapted for NTowerHT.
+    """
+    model.eval()
+
+    emergent_samples = []
+
+    # Access the final linear layer weights in the HTClassifier
+    # HTClassifier head is: Sequential(ReLU(), Dropout(), Linear())
+    # Index 2 is the Linear layer
+    final_linear = model.bridge.classifier_fused.head[2]
+    final_weights = (
+        final_linear.weight.detach().cpu()
+    )  # Shape: [num_classes, fusion_dim]
+    final_bias = final_linear.bias.detach().cpu()
+
+    with torch.no_grad():
+        for batch in loader:
+            x = batch.get("image_1")
+            m = batch.get("matrix_1")
+            y = batch.get("label_1")
+
+            if not (torch.is_tensor(x) and torch.is_tensor(m) and torch.is_tensor(y)):
+                continue
+
+            x, m, y = x.to(device), m.to(device), y.to(device)
+
+            # 1. Get Tower Embeddings
+            img_feats = model.img_tower(x)
+            md_feats = model.cd_tower(m)
+
+            # 2. Get Independent Tower Predictions
+            logits_i = model.bridge.aux_heads[0](img_feats)
+            logits_m = model.bridge.aux_heads[1](md_feats)
+            pi = logits_i.argmax(dim=1)
+            pm = logits_m.argmax(dim=1)
+
+            # 3. Get Fused Representation & Prediction
+            # _compute_fused applies the Hadamard product and optional SE gate
+            z_fused = model.bridge._compute_fused([img_feats, md_feats])
+            logits_fused = model.bridge.classifier_fused(z_fused)
+            pf = logits_fused.argmax(dim=1)
+
+            # 4. Find the "Aha!" Moments and "Corrections"
+            for i in range(len(y)):
+                yi = y[i].cpu().item()
+                is_correct = pf[i] == yi
+                img_wrong = pi[i] != yi
+                md_wrong = pm[i] != yi
+
+                if not is_correct:
+                    continue
+
+                is_aha = img_wrong and md_wrong
+                is_correction = img_wrong or md_wrong
+
+                if is_correction:
+                    category = (
+                        "Aha!"
+                        if is_aha
+                        else ("Corrected Image" if img_wrong else "Corrected Clinical")
+                    )
+                    # Apply the ReLU that happens inside HTClassifier before the Linear layer
+                    z_act = F.relu(z_fused[i]).cpu()
+
+                    # Calculate how much each feature contributed to the correct class logit
+                    feature_contributions = z_act * final_weights[yi]
+
+                    emergent_samples.append(
+                        {
+                            "patient_id": (
+                                batch.get("id_1", ["Unknown"])[i]
+                                if "id_1" in batch
+                                else "Unknown"
+                            ),
+                            "target_class": yi,
+                            "category": category,
+                            "z_activated": z_act.numpy(),
+                            "contributions": feature_contributions.numpy(),
+                            "total_logit": logits_fused[i, yi].cpu().item(),
+                        }
+                    )
+
+    return emergent_samples
+
+
+def plot_top_emergent_features(emergent_samples, top_k=10):
+    """
+    Plots the top K contributing dimensions across all emergent success samples.
+    """
+    if not emergent_samples:
+        print("No emergent success or correction samples found in this pass.")
+        return
+
+    ahas = [s for s in emergent_samples if s["category"] == "Aha!"]
+    corrected_img = [s for s in emergent_samples if s["category"] == "Corrected Image"]
+    corrected_clin = [
+        s for s in emergent_samples if s["category"] == "Corrected Clinical"
+    ]
+
+    print(f"\nFound {len(emergent_samples)} total events of interest:")
+    print(f"  - 'Aha!' Moments (Both wrong, Fused right): {len(ahas)}")
+    print(f"  - Corrected Image (Image wrong, Fused right): {len(corrected_img)}")
+    print(
+        f"  - Corrected Clinical (Clinical wrong, Fused right): {len(corrected_clin)}"
+    )
+
+    # Prioritize true Aha moments if they exist, otherwise use corrections
+    plot_samples = ahas if ahas else emergent_samples
+    plot_title = (
+        "Aha! Moments (Both Wrong, Fused Right)"
+        if ahas
+        else "Fusion Corrections (At least one wrong)"
+    )
+
+    # Average the feature contributions across samples
+    all_contribs = np.stack([s["contributions"] for s in plot_samples])
+    mean_contribs = all_contribs.mean(axis=0)
+
+    # Get indices of the top K features with the highest absolute contribution
+    top_indices = np.argsort(np.abs(mean_contribs))[-top_k:][::-1]
+
+    top_values = mean_contribs[top_indices]
+    labels = [f"Dim {idx}" for idx in top_indices]
+
+    plt.figure(figsize=(10, 6))
+    colors = ["green" if v > 0 else "red" for v in top_values]
+
+    plt.barh(np.arange(top_k), top_values[::-1], color=colors[::-1])
+    plt.yticks(np.arange(top_k), labels[::-1])
+    plt.xlabel("Mean Contribution to Correct Logit")
+    plt.title(f"Top {top_k} Features Driving {plot_title}")
+    plt.tight_layout()
+    plt.show()
+
+    print(f"\nAnalyzed {len(plot_samples)} samples for this plot.")
+    print("\nTop Feature Breakdown:")
+    for idx, val in zip(top_indices, top_values):
+        print(f"Dimension {idx:3d}: {val:+.4f} average logit push")
+
+
+def main():
+    device = choose_device("auto")
+    print(f"Using device: {device}")
+
+    repo_root = Path(__file__).resolve().parents[3]
+    image_dir = repo_root / "Papila" / "FundusImages"
+    clinical_dir = repo_root / "Papila" / "ClinicalData"
+
+    print("Loading PAPILA data with Phase 5 settings...")
+    data = build_papila_data(
+        image_dir=str(image_dir),
+        clinical_dir=str(clinical_dir),
+        label_col="Diagnosis",
+        cat_cols=["Gender", "Phakic/Pseudophakic"],
+        iop_corr_method="ratio",
+        iop_drop_raw=True,
+        exclude_cols=["Axial_Length"],
+    )
+    # Filter to binary
+    data.df = data.df[data.df["Diagnosis"].isin([0, 1])].reset_index(drop=True)
+
+    print("Building dataloader (All samples)...")
+    profile_eye = build_papila_profile(
+        patient_col="Patient ID", label_col="Diagnosis", sample_mode="eye"
+    )
+    eye_samples = filter_eye_samples(
+        profile_eye.build_samples(df=data.df, clinical=data)
+    )
+
+    loader = make_loader(
+        eye_samples,
+        profile_eye.slot_descriptors(),
+        image_transform=build_eval_transform("refugelike"),
+        batch_size=16,
+        shuffle=False,
+        num_workers=4,
+    )
+
+    print("Building SingleEyeHT model...")
+    model = SingleEyeHT(
+        backbone="refugelike",
+        freeze_ratio=0.0,
+        augment=False,
+        clinical_data=data,
+        num_classes=2,
+        cd_hidden_dim=128,
+        fusion_dim=256,
+        bridge_mode="fused",
+    ).to(device)
+
+    base_ckpt_dir = repo_root / "v3" / "results" / "phase5" / "logit_mlp_head_ckpt"
+    checkpoints = sorted(list(base_ckpt_dir.rglob("best_single.pt")))
+
+    if not checkpoints:
+        print(f"\n[!] No checkpoints found in {base_ckpt_dir}")
+        print("Please ensure you ran the jobs with the --save-checkpoints flag.")
+        return
+
+    all_emergent_data = []
+
+    for checkpoint_path in checkpoints:
+        print(f"\nProcessing {checkpoint_path.relative_to(repo_root)}...")
+        state_dict = torch.load(checkpoint_path, map_location=device)
+
+        # Backward compatibility for checkpoints saved before the N-tower bridge refactor
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            k = k.replace("bridge.W_img.", "bridge.W.0.")
+            k = k.replace("bridge.W_md.", "bridge.W.1.")
+            k = k.replace("bridge.ln_img.", "bridge.ln.0.")
+            k = k.replace("bridge.ln_md.", "bridge.ln.1.")
+            k = k.replace("bridge.classifier_img.", "bridge.aux_heads.0.")
+            k = k.replace("bridge.classifier_cd.", "bridge.aux_heads.1.")
+            k = k.replace(
+                "bridge.classifier_fused.2.", "bridge.classifier_fused.head.2."
+            )
+            new_state_dict[k] = v
+
+        model.load_state_dict(new_state_dict)
+
+        emergent_data = extract_emergent_features(model, loader, device)
+        all_emergent_data.extend(emergent_data)
+        print(f"  -> Found {len(emergent_data)} events of interest.")
+
+    print(
+        f"\nTotal aggregated events across {len(checkpoints)} checkpoints: {len(all_emergent_data)}"
+    )
+    plot_top_emergent_features(all_emergent_data, top_k=15)
+
+
+if __name__ == "__main__":
+    main()

v3/scripts/development/stage_seg_cnn/run_seg_cnn.py

@@ -61,7 +61,7 @@ from tqdm import tqdm
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[4]))
 
-from v3.classes.seg_cnn import (
+from v3.classes.geometry_towers import (
     SegCNN, SegMapDataset, SegMapRecord,
     UNetFineTuneDataset, precompute_unet_seg_maps,
 )

v3/scripts/main/phase3/epoch_grid.json

@@ -7,7 +7,7 @@
 
   "common_args": [
     "--eval-mode", "binary",
-    "--tower-mode", "single",
+    "--hypertower-mode", "single",
     "--in-memory-cache",
     "--augment",
     "--tune-binary-threshold",

v3/scripts/main/phase3/experiment_grid.json

@@ -12,7 +12,7 @@
   "common_args": [
     "--eval-mode",
     "binary",
-    "--tower-mode",
+    "--hypertower-mode",
     "single",
     "--epochs",
     "30",

v3/scripts/main/phase3/phase35_grid.json

@@ -8,7 +8,7 @@
 
   "common_args": [
     "--eval-mode", "binary",
-    "--tower-mode", "single",
+    "--hypertower-mode", "single",
     "--epochs", "30",
     "--in-memory-cache",
     "--augment",

v3/scripts/main/phase4/experiment_grid.json

@@ -31,7 +31,7 @@
   "baseline": {
     "run_name": "phase4/single",
     "description": "Single-eye baseline with best phase 3 image settings. Direct comparison point for bilateral modes.",
-    "extra_args": ["--tower-mode", "single"]
+    "extra_args": ["--hypertower-mode", "single"]
   },
 
   "groups": [
@@ -42,17 +42,17 @@
         {
           "run_name": "phase4/ensemble",
           "description": "Ensemble: two independent single-eye forward passes, patient-level average of OD+OS scores.",
-          "extra_args": ["--tower-mode", "ensemble"]
+          "extra_args": ["--hypertower-mode", "ensemble"]
         },
         {
           "run_name": "phase4/bilateral",
           "description": "BilateralHT: shared towers, concat OD+OS → learned joint projection MLP → classifier.",
-          "extra_args": ["--tower-mode", "bilateral"]
+          "extra_args": ["--hypertower-mode", "bilateral"]
         },
         {
           "run_name": "phase4/siamese",
           "description": "SiameseHT: shared backbone, mean+delta (asymmetry) representation → classifier.",
-          "extra_args": ["--tower-mode", "siamese"]
+          "extra_args": ["--hypertower-mode", "siamese"]
         }
       ]
     },
@@ -64,12 +64,12 @@
         {
           "run_name": "phase4/bilateral_loss_all",
           "description": "BilateralHT with all-losses mode — joint training dynamics may differ from single-eye.",
-          "extra_args": ["--tower-mode", "bilateral", "--tower-loss-mode", "all"]
+          "extra_args": ["--hypertower-mode", "bilateral", "--tower-loss-mode", "all"]
         },
         {
           "run_name": "phase4/siamese_loss_all",
           "description": "SiameseHT with all-losses mode.",
-          "extra_args": ["--tower-mode", "siamese", "--tower-loss-mode", "all"]
+          "extra_args": ["--hypertower-mode", "siamese", "--tower-loss-mode", "all"]
         }
       ]
     }

v3/scripts/main/phase5/dispatch_logit_mlp_ckpt.py

@@ -38,7 +38,7 @@ N_REPS         = 10
 RUN_ARGS = [
     "--eval-mode",            "binary",
     "--bridge-mode",          "fused",
-    "--tower-mode",           "ensemble",
+    "--hypertower-mode",           "ensemble",
     "--fused-head",
     "--head-type",            "logit_mlp",
     "--epochs",               "30",

v3/scripts/main/phase5/experiment_grid.json

@@ -36,7 +36,7 @@
   "baseline": {
     "run_name": "phase5/single_fused",
     "description": "Single-eye + clinical data — phase 3 best config, re-run as direct comparison baseline for phase 5.",
-    "extra_args": ["--tower-mode", "single"]
+    "extra_args": ["--hypertower-mode", "single"]
   },
 
   "groups": [
@@ -47,17 +47,17 @@
         {
           "run_name": "phase5/ensemble_fused",
           "description": "Ensemble (independent OD+OS) + clinical data via fused bridge.",
-          "extra_args": ["--tower-mode", "ensemble"]
+          "extra_args": ["--hypertower-mode", "ensemble"]
         },
         {
           "run_name": "phase5/bilateral_fused",
           "description": "BilateralHT + clinical data — full canonical HyperTower.",
-          "extra_args": ["--tower-mode", "bilateral"]
+          "extra_args": ["--hypertower-mode", "bilateral"]
         },
         {
           "run_name": "phase5/siamese_fused",
           "description": "SiameseHT + clinical data — siamese mean+delta with fused clinical bridge.",
-          "extra_args": ["--tower-mode", "siamese"]
+          "extra_args": ["--hypertower-mode", "siamese"]
         }
       ]
     },
@@ -69,17 +69,17 @@
         {
           "run_name": "phase5/ensemble_fused_head",
           "description": "Ensemble + clinical data + attention scorer head (Linear(C→1) per eye, softmax-weighted average).",
-          "extra_args": ["--tower-mode", "ensemble", "--fused-head"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--fused-head"]
         },
         {
           "run_name": "phase5/logit_mlp_head",
           "description": "Ensemble + clinical data + logit-level MLP head (cat([logit_od, logit_os]) → FC(64) → FC(C)).",
-          "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--head-type", "logit_mlp"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--head-type", "logit_mlp"]
         },
         {
           "run_name": "phase5/embedding_mlp_head",
           "description": "Ensemble + clinical data + embedding-level MLP head (cat([z_od, z_os]) → FC(256) → FC(C)).",
-          "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--head-type", "embedding_mlp"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--head-type", "embedding_mlp"]
         }
       ]
     }

v3/scripts/main/phase6/experiment_grid.json

@@ -47,7 +47,7 @@
   "baseline": {
     "run_name": "phase6/single_no_geom",
     "description": "Single-eye + clinical data, no geometry — needed as Phase 6 baseline since no prior phase ran single with all tuned hyperparameters (iop_ratio_drop_raw, bcd_p05, refugelike).",
-    "extra_args": ["--tower-mode", "single"]
+    "extra_args": ["--hypertower-mode", "single"]
   },
 
   "groups": [
@@ -58,17 +58,17 @@
         {
           "run_name": "phase6/vec_gt_single",
           "description": "Single-eye + clinical + GT geometry vector appended to clinical stream.",
-          "extra_args": ["--tower-mode", "single", "--geometry-dim", "5", "--geometry-source", "gt"]
+          "extra_args": ["--hypertower-mode", "single", "--geometry-dim", "5", "--geometry-source", "gt"]
         },
         {
           "run_name": "phase6/vec_gt_ensemble",
           "description": "Ensemble + clinical + GT geometry vector (per-eye geometry, independent OD+OS mean).",
-          "extra_args": ["--tower-mode", "ensemble", "--geometry-dim", "5", "--geometry-source", "gt"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--geometry-dim", "5", "--geometry-source", "gt"]
         },
         {
           "run_name": "phase6/vec_gt_fused_head",
           "description": "Ensemble + fused head + clinical + GT geometry vector.",
-          "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--geometry-dim", "5", "--geometry-source", "gt"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--geometry-dim", "5", "--geometry-source", "gt"]
         }
       ]
     },
@@ -80,17 +80,17 @@
         {
           "run_name": "phase6/vec_unet_single",
           "description": "Single-eye + clinical + U-Net geometry vector.",
-          "extra_args": ["--tower-mode", "single", "--geometry-dim", "5", "--geometry-source", "unet"]
+          "extra_args": ["--hypertower-mode", "single", "--geometry-dim", "5", "--geometry-source", "unet"]
         },
         {
           "run_name": "phase6/vec_unet_ensemble",
           "description": "Ensemble + clinical + U-Net geometry vector.",
-          "extra_args": ["--tower-mode", "ensemble", "--geometry-dim", "5", "--geometry-source", "unet"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--geometry-dim", "5", "--geometry-source", "unet"]
         },
         {
           "run_name": "phase6/vec_unet_fused_head",
           "description": "Ensemble + fused head + clinical + U-Net geometry vector.",
-          "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--geometry-dim", "5", "--geometry-source", "unet"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--geometry-dim", "5", "--geometry-source", "unet"]
         }
       ]
     },
@@ -103,17 +103,17 @@
         {
           "run_name": "phase6/tower_gt_single",
           "description": "Single-eye + clinical tower + dedicated GT geometry tower.",
-          "extra_args": ["--tower-mode", "single", "--geometry-tower", "--geometry-source", "gt"]
+          "extra_args": ["--hypertower-mode", "single", "--geometry-tower", "--geometry-source", "gt"]
         },
         {
           "run_name": "phase6/tower_gt_ensemble",
           "description": "Ensemble + clinical tower + dedicated GT geometry tower.",
-          "extra_args": ["--tower-mode", "ensemble", "--geometry-tower", "--geometry-source", "gt"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--geometry-tower", "--geometry-source", "gt"]
         },
         {
           "run_name": "phase6/tower_gt_fused_head",
           "description": "Ensemble + fused head + clinical tower + dedicated GT geometry tower.",
-          "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--geometry-tower", "--geometry-source", "gt"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--geometry-tower", "--geometry-source", "gt"]
         }
       ]
     },
@@ -126,17 +126,17 @@
         {
           "run_name": "phase6/tower_unet_single",
           "description": "Single-eye + clinical tower + dedicated U-Net geometry tower.",
-          "extra_args": ["--tower-mode", "single", "--geometry-tower", "--geometry-source", "unet"]
+          "extra_args": ["--hypertower-mode", "single", "--geometry-tower", "--geometry-source", "unet"]
         },
         {
           "run_name": "phase6/tower_unet_ensemble",
           "description": "Ensemble + clinical tower + dedicated U-Net geometry tower.",
-          "extra_args": ["--tower-mode", "ensemble", "--geometry-tower", "--geometry-source", "unet"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--geometry-tower", "--geometry-source", "unet"]
         },
         {
           "run_name": "phase6/tower_unet_fused_head",
           "description": "Ensemble + fused head + clinical tower + dedicated U-Net geometry tower.",
-          "extra_args": ["--tower-mode", "ensemble", "--fused-head", "--geometry-tower", "--geometry-source", "unet"]
+          "extra_args": ["--hypertower-mode", "ensemble", "--fused-head", "--geometry-tower", "--geometry-source", "unet"]
         }
       ]
     }

v3/scripts/main/run_cv.py

@@ -9,7 +9,7 @@ Usage (single fold-seed, 5-fold, binary, ensemble):
   python -m v3.scripts.main.run_cv \
       --run-name my_run \
       --eval-mode binary \
-      --tower-mode ensemble \
+      --hypertower-mode ensemble \
       --epochs 40 \
       --augment \
       --tune-binary-threshold \
@@ -22,7 +22,7 @@ Usage (10x5 rep-CV, seeds 100..1000):
       --rep-seed-start 100 \
       --rep-seed-step 100 \
       --eval-mode binary \
-      --tower-mode ensemble \
+      --hypertower-mode ensemble \
       --epochs 40 \
       --augment \
       --tune-binary-threshold \

v3/scripts/main/run_ntower_cv.py

@@ -0,0 +1,533 @@
+#!/usr/bin/env python
+"""
+NTowerHT + HyperBridge ensemble cross-validation runner.
+
+Reproduces phase5/embedding_mlp_head using the new module architecture:
+
+  Stage 1 — per-eye NTowerHT  (eye-level samples, BCD training)
+    img_tower + cd_tower → Bridge → z_fused  [B, fusion_dim]
+
+  Stage 2 — HyperBridge(embedding_mlp)  (patient-level bilateral samples)
+    cat([z_od, z_os]) → Linear(2*fusion_dim → hidden_dim) → ReLU → Dropout → Linear → logits
+
+Training mirrors v3_hypertower ensemble+fused_head:
+  - Warmup phases for NTowerHT (tower_warmup → fused_warmup → main)
+  - NTowerHT frozen; HyperBridge trained on bilateral samples
+
+Usage (phase5/embedding_mlp_head equivalent):
+  python -m v3.scripts.main.run_ntower_cv \\
+      --run-name ntower/ensemble_fused \\
+      --eval-mode binary \\
+      --epochs 30 \\
+      --fusion-epochs 10 \\
+      --in-memory-cache \\
+      --augment \\
+      --tune-binary-threshold \\
+      --backbone refugelike \\
+      --iop-corr-method ratio \\
+      --iop-drop-raw \\
+      --exclude-cols Axial_Length
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
+
+from v3.classes.hypertower_models import NTowerHT, train_ntower_epoch, collect_probs_ntower
+from v3.classes.towerbase import _to_label_tensor
+from v3.classes.bridges import HyperBridge
+from v3.classes.image_towers import ImageEncoder
+from v3.classes.clinical_towers import ClinicalEncoder
+from v3.classes.papila_builders import build_papila_data
+from v3.classes.profiles import build_papila_profile
+from v3.classes.split_manager import PatientFirstSplitManager
+from types import SimpleNamespace
+from v3.classes.loader_factory import (
+    filter_eye_samples,
+    filter_bilateral_samples,
+    make_loader,
+    build_balanced_sampler,
+)
+from v3.classes.metrics import _score_arrays, compute_extended_metrics, tune_binary_threshold
+from v3.classes.transforms import build_eval_transform
+from v3.classes.utils import seed_everything, choose_device
+from v3.classes.croppers import build_image_preprocessor_from_args
+from v3.classes.image_loader import CachedImageLoader
+
+REPO_ROOT    = Path(__file__).resolve().parents[3]
+IMAGE_DIR    = REPO_ROOT / "Papila" / "FundusImages"
+CLINICAL_DIR = REPO_ROOT / "Papila" / "ClinicalData"
+
+# Batch key mapping for per-eye NTowerHT training
+EYE_KEY_MAP = {"img": "image_1", "cd": "matrix_1"}
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def build_parser() -> argparse.ArgumentParser:
+    ap = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    ap.add_argument("--run-name",              default="ntower/ensemble_fused")
+    ap.add_argument("--output-root",           default="v3/results")
+    ap.add_argument("--eval-mode",             default="binary", choices=["binary", "multiclass"])
+    ap.add_argument("--epochs",                type=int,   default=30,
+                    help="NTowerHT main-phase epochs")
+    ap.add_argument("--fusion-epochs",         type=int,   default=10,
+                    help="HyperBridge training epochs (after NTowerHT is frozen)")
+    ap.add_argument("--warmup-cd-epochs",      type=int,   default=40,
+                    help="Pre-train cd tower + aux head only (no image tower)")
+    ap.add_argument("--folds",                 type=int,   default=5)
+    ap.add_argument("--fold-seed",             type=int,   default=100,
+                    help="Seed for patient splits (rep00=100, rep01=200, ...)")
+    ap.add_argument("--seed",                  type=int,   default=1234,
+                    help="Seed for model init / per-fold RNG (matches V3HyperTower default)")
+    ap.add_argument("--batch-size",            type=int,   default=16)
+    ap.add_argument("--lr",                    type=float, default=1e-4)
+    ap.add_argument("--backbone",              default="refugelike")
+    ap.add_argument("--freeze-ratio",          type=float, default=0.0)
+    ap.add_argument("--augment",               action="store_true")
+    ap.add_argument("--fusion-dim",            type=int,   default=256)
+    ap.add_argument("--hyper-hidden-dim",      type=int,   default=256,
+                    help="HyperBridge hidden dim (default matches EmbeddingMLPEnsembleHT)")
+    ap.add_argument("--cd-hidden-dim",         type=int,   default=128)
+    ap.add_argument("--bcd-prob",              type=float, default=0.5)
+    ap.add_argument("--warmup-tower-epochs",   type=int,   default=3)
+    ap.add_argument("--warmup-fused-epochs",   type=int,   default=3)
+    ap.add_argument("--label-col",             default="Diagnosis")
+    ap.add_argument("--iop-corr-method",       default="ratio")
+    ap.add_argument("--iop-drop-raw",          action="store_true")
+    ap.add_argument("--exclude-cols",          nargs="*",  default=[])
+    ap.add_argument("--num-workers",           type=int,   default=0)
+    ap.add_argument("--in-memory-cache",       action="store_true")
+    ap.add_argument("--tune-binary-threshold", action="store_true")
+    ap.add_argument("--device",                default=None)
+    return ap
+
+
+# ---------------------------------------------------------------------------
+# Data
+# ---------------------------------------------------------------------------
+
+def load_data(args):
+    return build_papila_data(
+        image_dir=str(IMAGE_DIR),
+        clinical_dir=str(CLINICAL_DIR),
+        label_col=args.label_col,
+        cat_cols=["Gender", "Phakic/Pseudophakic"],
+        iop_corr_method=args.iop_corr_method,
+        iop_drop_raw=args.iop_drop_raw,
+        exclude_cols=args.exclude_cols or [],
+    )
+
+
+# ---------------------------------------------------------------------------
+# Model factories
+# ---------------------------------------------------------------------------
+
+def build_nt(data, num_classes: int, args) -> NTowerHT:
+    """Per-eye NTowerHT: image + clinical → bridge."""
+    img_enc = ImageEncoder(backbone=args.backbone, freeze_ratio=args.freeze_ratio, augment=args.augment)
+    cd_enc  = ClinicalEncoder(clinical_data=data, hidden_dim=args.cd_hidden_dim)
+    return NTowerHT(
+        towers={"img": img_enc, "cd": cd_enc},
+        num_classes=num_classes,
+        fusion_dim=args.fusion_dim,
+    )
+
+
+def build_hb(num_classes: int, args) -> HyperBridge:
+    """HyperBridge(embedding_mlp): cat([z_od, z_os]) → MLP → logits."""
+    return HyperBridge(
+        input_dims={"od": args.fusion_dim, "os": args.fusion_dim},
+        num_classes=num_classes,
+        hidden_dim=args.hyper_hidden_dim,
+        mode="embedding_mlp",
+    )
+
+
+
+
+# ---------------------------------------------------------------------------
+# Per-tower pre-warmup + HyperBridge training/inference
+# ---------------------------------------------------------------------------
+
+def train_tower_pre_warmup(
+    nt: NTowerHT,
+    tower_idx: int,
+    loader,
+    batch_key: str,
+    opt,
+    device: torch.device,
+) -> tuple[float, float]:
+    """Pre-train a single tower (by index) + its bridge aux head only.
+
+    Everything else is frozen.  Caller is responsible for passing a loader
+    that omits unnecessary slots (e.g. cd_only_loader strips image_1).
+    """
+    tower_name = list(nt.towers.keys())[tower_idx]
+
+    for p in nt.parameters():
+        p.requires_grad_(False)
+    for p in nt.towers[tower_name].parameters():
+        p.requires_grad_(True)
+    for p in nt.bridge.aux_heads[tower_idx].parameters():
+        p.requires_grad_(True)
+
+    nt.train()
+    total_loss = total_correct = total_n = 0
+    for batch in loader:
+        x = batch.get(batch_key)
+        y = batch.get("label_1")
+        if not torch.is_tensor(x):
+            continue
+        y_t    = _to_label_tensor(y, device)
+        z      = nt.towers[tower_name](x.to(device))
+        logits = nt.bridge.aux_heads[tower_idx](z)
+        loss   = F.cross_entropy(logits, y_t)
+        opt.zero_grad(); loss.backward(); opt.step()
+        total_loss    += loss.item() * len(y_t)
+        total_correct += int((logits.argmax(1) == y_t).sum())
+        total_n       += len(y_t)
+
+    for p in nt.parameters():
+        p.requires_grad_(True)
+
+    return (
+        total_loss / total_n if total_n else float("nan"),
+        total_correct / total_n if total_n else float("nan"),
+    )
+
+
+def _encode_eye(nt: NTowerHT, batch: dict, batch_key_map: dict[str, str], device) -> torch.Tensor:
+    """Run all towers from a single-eye batch dict and return z_fused."""
+    embeddings = {name: nt.towers[name](batch[key].to(device))
+                  for name, key in batch_key_map.items()}
+    return nt.encode(embeddings)
+
+
+def train_hb_epoch(
+    nt: NTowerHT,
+    hb: HyperBridge,
+    loader,
+    opt,
+    device: torch.device,
+) -> tuple[float, float]:
+    """Train HyperBridge with NTowerHT frozen.
+
+    Each bilateral batch provides both eyes; we encode each through NTowerHT
+    to get z_od, z_os, then train HyperBridge to fuse them.
+    """
+    nt.eval()
+    hb.train()
+    total_loss = total_correct = total_n = 0
+    for batch in loader:
+        x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
+        x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
+        y  = batch.get("label_1")
+        if not all(torch.is_tensor(t) for t in (x1, m1, x2, m2)):
+            continue
+        y_t = _to_label_tensor(y, device)
+        # Build per-eye batch dicts (keyed by batch key, as _encode_eye expects)
+        od_batch = {"image_1": x1, "matrix_1": m1}
+        os_batch = {"image_1": x2, "matrix_1": m2}
+        with torch.no_grad():
+            z_od = _encode_eye(nt, od_batch, EYE_KEY_MAP, device)
+            z_os = _encode_eye(nt, os_batch, EYE_KEY_MAP, device)
+        logits, _ = hb({"od": z_od, "os": z_os})
+        loss = F.cross_entropy(logits, y_t)
+        opt.zero_grad(); loss.backward(); opt.step()
+        total_loss    += loss.item() * len(y_t)
+        total_correct += int((logits.argmax(1) == y_t).sum())
+        total_n       += len(y_t)
+    return (
+        total_loss / total_n if total_n else float("nan"),
+        total_correct / total_n if total_n else float("nan"),
+    )
+
+
+def collect_probs_hb(
+    nt: NTowerHT,
+    hb: HyperBridge,
+    loader,
+    device: torch.device,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Collect HyperBridge predictions (patient-level)."""
+    nt.eval(); hb.eval()
+    y_all, p_all = [], []
+    with torch.no_grad():
+        for batch in loader:
+            x1 = batch.get("image_1"); m1 = batch.get("matrix_1")
+            x2 = batch.get("image_2"); m2 = batch.get("matrix_2")
+            y  = batch.get("label_1")
+            if not all(torch.is_tensor(t) for t in (x1, m1, x2, m2)):
+                continue
+            y_t  = _to_label_tensor(y, device)
+            z_od = _encode_eye(nt, {"image_1": x1, "matrix_1": m1}, EYE_KEY_MAP, device)
+            z_os = _encode_eye(nt, {"image_1": x2, "matrix_1": m2}, EYE_KEY_MAP, device)
+            logits, _ = hb({"od": z_od, "os": z_os})
+            y_all.append(y_t.cpu().numpy())
+            p_all.append(F.softmax(logits, dim=1).cpu().numpy())
+    if not y_all:
+        return np.zeros(0, dtype=np.int64), np.zeros((0, 0), dtype=np.float32)
+    return np.concatenate(y_all), np.concatenate(p_all, axis=0)
+
+
+# ---------------------------------------------------------------------------
+# Fold runner
+# ---------------------------------------------------------------------------
+
+def phase_for_epoch(epoch: int, warmup_tower: int, warmup_fused: int) -> str:
+    if epoch < warmup_tower:
+        return "tower_warmup"
+    if epoch < warmup_tower + warmup_fused:
+        return "fused_warmup"
+    return "main"
+
+
+def run_fold(fold: int, plans, data, num_classes: int, device, args,
+             profile_eye, profile_patient, image_preprocessor, image_cache) -> dict:
+    nan = float("nan")
+    seed_everything(args.seed + fold * 100)
+
+    split = plans[fold]  # PatientSplit with .train/.val/.test DataFrames
+
+    # Eye-level splits (for NTowerHT training)
+    eye_train = filter_eye_samples(profile_eye.build_samples(df=split.train, clinical=data))
+    eye_val   = filter_eye_samples(profile_eye.build_samples(df=split.val,   clinical=data))
+
+    # Patient-level splits (for HyperBridge training/eval)
+    bilat_train = filter_bilateral_samples(profile_patient.build_samples(df=split.train, clinical=data))
+    bilat_val   = filter_bilateral_samples(profile_patient.build_samples(df=split.val,   clinical=data))
+    bilat_test  = filter_bilateral_samples(profile_patient.build_samples(
+        df=split.test, clinical=data)) if split.test is not None else []
+
+    if not bilat_val:
+        print(f"  fold{fold+1}: no bilateral val samples, skipping.", flush=True)
+        return {"fold": fold, "val_auc": nan, "val_acc": nan, "val_n": 0,
+                "val_kappa": nan, "val_mcc": nan, "val_f1": nan,
+                "val_threshold": 0.5, "test_auc": nan, "test_acc": nan, "test_n": nan}
+
+    # Build Stage 1 model only — HyperBridge is built after Stage 1 completes,
+    # matching phase5 where EmbeddingMLPEnsembleHT is constructed at Phase 2 start.
+    # Building hb here would consume random state and shift all subsequent dropout ops.
+    nt = build_nt(data, num_classes, args).to(device)
+    opt_nt = torch.optim.Adam(nt.parameters(), lr=args.lr)
+
+    slots_eye     = profile_eye.slot_descriptors()
+    slots_patient = profile_patient.slot_descriptors()
+
+    loader_kw = dict(batch_size=args.batch_size, num_workers=args.num_workers,
+                     image_cache=image_cache, persistent_workers=args.num_workers > 0)
+    eval_transform = build_eval_transform(args.backbone)
+
+    # Unified eye-level loader (all slots, shuffle=True — matches phase5 exactly)
+    train_eye_loader = make_loader(
+        eye_train, slots_eye, image_transform=nt.transform,
+        image_preprocessor=image_preprocessor, shuffle=True, **loader_kw,
+    )
+    # cd-only loader for warmup: strips image_1 so image decoding is skipped entirely
+    slots_cd_only = {k: v for k, v in slots_eye.items() if k != "image_1"}
+    cd_warmup_loader = make_loader(
+        eye_train, slots_cd_only, image_transform=None,
+        image_preprocessor=None, shuffle=True,
+        sampler=build_balanced_sampler(eye_train), **loader_kw,
+    )
+
+    val_eye_loader = make_loader(
+        eye_val, slots_eye, image_transform=eval_transform,
+        image_preprocessor=image_preprocessor, shuffle=False, **loader_kw,
+    )
+    train_bilat_loader = make_loader(
+        bilat_train, slots_patient, image_transform=nt.transform,
+        image_preprocessor=image_preprocessor, shuffle=True, **loader_kw,
+    )
+    val_bilat_loader = make_loader(
+        bilat_val, slots_patient, image_transform=eval_transform,
+        image_preprocessor=image_preprocessor, shuffle=False, **loader_kw,
+    )
+    test_bilat_loader = make_loader(
+        bilat_test, slots_patient, image_transform=eval_transform,
+        image_preprocessor=image_preprocessor, shuffle=False, **loader_kw,
+    ) if bilat_test else None
+
+    # ── Stage 1: train NTowerHT ───────────────────────────────────────────
+    tower_names = list(nt.towers.keys())
+    tower_keys  = list(EYE_KEY_MAP.values())
+
+    # Per-tower pre-warmup using the cd-only loader (no image loading overhead)
+    pre_warmup_epochs = [args.warmup_cd_epochs if name == "cd" else 0
+                         for name in tower_names]
+    warmup_loaders    = {"cd": cd_warmup_loader}
+    for idx, (name, n_epochs) in enumerate(zip(tower_names, pre_warmup_epochs)):
+        if n_epochs == 0:
+            continue
+        for epoch in range(n_epochs):
+            tr_loss, tr_acc = train_tower_pre_warmup(
+                nt, idx, warmup_loaders[name], tower_keys[idx], opt_nt, device,
+            )
+            print(
+                f"  fold{fold+1} [NT] ep{epoch+1:03d}/{n_epochs} [{name}_warmup     ]"
+                f"  loss={tr_loss:.4f}  tr_acc={tr_acc:.3f}",
+                flush=True,
+            )
+
+    total_nt_epochs = args.warmup_tower_epochs + args.warmup_fused_epochs + args.epochs
+
+    for epoch in range(total_nt_epochs):
+        phase = phase_for_epoch(epoch, args.warmup_tower_epochs, args.warmup_fused_epochs)
+        tr_loss, tr_acc = train_ntower_epoch(
+            nt, train_eye_loader, opt_nt, device,
+            batch_key_map=EYE_KEY_MAP, phase=phase, bcd_prob=args.bcd_prob,
+        )
+        y_v, p_v = collect_probs_ntower(nt, val_eye_loader, device, batch_key_map=EYE_KEY_MAP)
+        _, val_auc, _ = _score_arrays(y_v, p_v, num_classes)
+        print(
+            f"  fold{fold+1} [NT] ep{epoch+1:03d}/{total_nt_epochs} [{phase:14s}]"
+            f"  loss={tr_loss:.4f}  tr_acc={tr_acc:.3f}  val_auc={val_auc:.4f}",
+            flush=True,
+        )
+
+    # Freeze NTowerHT (final-epoch weights, matching phase5 — no best-checkpoint restore)
+    for p in nt.parameters():
+        p.requires_grad_(False)
+
+    # ── Stage 2: train HyperBridge ────────────────────────────────────────
+    # Build here (not at fold start) to match phase5 random-state sequence
+    hb = build_hb(num_classes, args).to(device)
+    opt_hb = torch.optim.Adam(hb.parameters(), lr=args.lr)
+
+    for epoch in range(args.fusion_epochs):
+        tr_loss, tr_acc = train_hb_epoch(nt, hb, train_bilat_loader, opt_hb, device)
+        y_v, p_v = collect_probs_hb(nt, hb, val_bilat_loader, device)
+        _, val_auc, _ = _score_arrays(y_v, p_v, num_classes)
+        print(
+            f"  fold{fold+1} [HB] ep{epoch+1:02d}/{args.fusion_epochs} [fusion       ]"
+            f"  loss={tr_loss:.4f}  tr_acc={tr_acc:.3f}  val_auc={val_auc:.4f}",
+            flush=True,
+        )
+    # Final-epoch weights used (no best-checkpoint restore, matching phase5)
+
+    # ── Final eval ────────────────────────────────────────────────────────
+    y_val, p_val = collect_probs_hb(nt, hb, val_bilat_loader, device)
+    val_acc, val_auc, val_n = _score_arrays(y_val, p_val, num_classes)
+    ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {}
+
+    val_threshold = 0.5
+    if args.tune_binary_threshold and num_classes == 2 and y_val.size >= 2:
+        val_threshold = tune_binary_threshold(y_val, p_val[:, 1])
+
+    test_auc = test_acc = test_n = nan
+    if test_bilat_loader is not None:
+        y_te, p_te = collect_probs_hb(nt, hb, test_bilat_loader, device)
+        test_acc, test_auc, test_n = _score_arrays(y_te, p_te, num_classes)
+
+    return {
+        "fold":          fold,
+        "val_auc":       val_auc,
+        "val_acc":       val_acc,
+        "val_n":         val_n,
+        "val_kappa":     ext.get("kappa", nan),
+        "val_mcc":       ext.get("mcc", nan),
+        "val_f1":        ext.get("macro_f1", nan),
+        "val_threshold": val_threshold,
+        "test_auc":      test_auc,
+        "test_acc":      test_acc,
+        "test_n":        test_n,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    ap   = build_parser()
+    args = ap.parse_args()
+
+    num_classes = 2 if args.eval_mode == "binary" else 3
+    device      = choose_device(args.device)
+    print(f"Device: {device}", flush=True)
+
+    print("Loading data ...", flush=True)
+    data = load_data(args)
+    print(f"  feature_dim={data.feature_dim}", flush=True)
+
+    image_cache        = CachedImageLoader() if args.in_memory_cache else None
+    image_preprocessor = build_image_preprocessor_from_args(args)
+
+    profile_eye     = build_papila_profile(patient_col="Patient ID", label_col=args.label_col, sample_mode="eye")
+    profile_patient = build_papila_profile(patient_col="Patient ID", label_col=args.label_col, sample_mode="patient")
+
+    # Filter to binary labels before splitting (mirrors V3HyperTower)
+    df_mode = data.df.copy()
+    if args.eval_mode == "binary":
+        df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
+
+    split_mgr   = PatientFirstSplitManager(patient_col="Patient ID", label_col=args.label_col)
+    split_args  = SimpleNamespace(eval_mode=args.eval_mode, n_splits=args.folds, fold_seed=args.fold_seed)
+    clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col)
+    plans       = split_mgr.build_plans(clinical=clinical_ns, args=split_args, profile=None)
+
+    out_dir = REPO_ROOT / args.output_root / args.run_name / "binary" / "ntower"
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    fold_results = []
+    t0 = time.time()
+
+    for fold in range(args.folds):
+        split     = plans[fold]
+        bilat_val = filter_bilateral_samples(profile_patient.build_samples(df=split.val, clinical=data))
+        print(
+            f"\n── fold {fold+1}/{args.folds}"
+            f"  train_patients={split.train['Patient ID'].nunique()}"
+            f"  val={len(bilat_val)} ──",
+            flush=True,
+        )
+        result = run_fold(fold, plans, data, num_classes, device, args,
+                          profile_eye, profile_patient, image_preprocessor, image_cache)
+        fold_results.append(result)
+        print(
+            f"  fold{fold+1} DONE  val_auc={result['val_auc']:.4f}"
+            f"  test_auc={result['test_auc']:.4f}",
+            flush=True,
+        )
+
+    if fold_results:
+        val_aucs  = [r["val_auc"]  for r in fold_results if not np.isnan(r["val_auc"])]
+        test_aucs = [r["test_auc"] for r in fold_results if not np.isnan(r["test_auc"])]
+        summary = {
+            "run_name":       args.run_name,
+            "backbone":       args.backbone,
+            "nt_epochs":      args.epochs,
+            "fusion_epochs":  args.fusion_epochs,
+            "folds":          args.folds,
+            "mean_val_auc":   float(np.mean(val_aucs))  if val_aucs  else float("nan"),
+            "std_val_auc":    float(np.std(val_aucs))   if val_aucs  else float("nan"),
+            "mean_test_auc":  float(np.mean(test_aucs)) if test_aucs else float("nan"),
+            "std_test_auc":   float(np.std(test_aucs))  if test_aucs else float("nan"),
+            "elapsed_s":      round(time.time() - t0, 1),
+            "fold_results":   fold_results,
+        }
+        summary_path = out_dir / "summary.json"
+        summary_path.write_text(json.dumps(summary, indent=2))
+        print(f"\n{'='*60}", flush=True)
+        print(f"Val  AUC: {summary['mean_val_auc']:.4f} ± {summary['std_val_auc']:.4f}", flush=True)
+        print(f"Test AUC: {summary['mean_test_auc']:.4f} ± {summary['std_test_auc']:.4f}", flush=True)
+        print(f"Saved: {summary_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

v3/scripts/output_analysis/explainability/gradcam_phase5.py

@@ -239,7 +239,7 @@ def make_disc_attention_detail(
 def build_model(ckpt_path: Path, device: torch.device):
     """Reconstruct SingleEyeHT from checkpoint and load weights."""
     from types import SimpleNamespace
-    from v3.classes.models import SingleEyeHT
+    from v3.classes.hypertower_models import SingleEyeHT
     sd = torch.load(ckpt_path, map_location="cpu")
     # ClinicalTower only reads clinical_data.feature_dim at init time
     cd_in = sd["cd_tower.block0.0.weight"].shape[1]

v3/scripts/output_analysis/explainability/permutation_importance_phase5.py

@@ -44,7 +44,7 @@ SEED = 0
 # ── Model ─────────────────────────────────────────────────────────────────────
 
 def build_model(ckpt_path: Path, device: torch.device):
-    from v3.classes.models import SingleEyeHT
+    from v3.classes.hypertower_models import SingleEyeHT
     sd = torch.load(ckpt_path, map_location="cpu")
     cd_in = sd["cd_tower.block0.0.weight"].shape[1]
     model = SingleEyeHT(
@@ -169,7 +169,7 @@ def run_fold(rep_idx: int, fold_idx: int, model, data, device: torch.device,
     # Baseline AUC
     with torch.no_grad():
         md_feats = model.cd_tower(meta_all.to(device))
-        out_f, _, _ = model.bridge(img_feats, md_feats)
+        out_f, _ = model.bridge.fuse([img_feats, md_feats])
         probs_base = F.softmax(out_f, dim=1)[:, 1].cpu().numpy()
     baseline_auc = roc_auc_score(y_true, probs_base)
     print(f"  fold{fold_idx}: baseline AUC={baseline_auc:.4f}  N={len(y_true)}")
@@ -186,7 +186,7 @@ def run_fold(rep_idx: int, fold_idx: int, model, data, device: torch.device,
             meta_perm[:, dims] = meta_perm[perm_idx][:, dims]
             with torch.no_grad():
                 md_p = model.cd_tower(meta_perm.to(device))
-                out_p, _, _ = model.bridge(img_feats, md_p)
+                out_p, _ = model.bridge.fuse([img_feats, md_p])
                 probs_p = F.softmax(out_p, dim=1)[:, 1].cpu().numpy()
             try:
                 drops.append(baseline_auc - roc_auc_score(y_true, probs_p))

v4/classes/accessory/backbones.py

@@ -0,0 +1,217 @@
+"""backbones — backbone registry and builder for v4 image towers."""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Dict, List
+
+import torch
+from torch import nn
+from torchvision import models, transforms
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+
+REFUGELIKE_BACKBONE_PATH    = _REPO_ROOT / "models/v2/refuge/refugelike_backbone.pt"
+REFUGE_DENSENET_PATH        = _REPO_ROOT / "models/refuge/classifier/refuge_densenet_backbone.pt"
+REFUGE_EFFICIENT_B0_PATH    = _REPO_ROOT / "models/refuge/classifier/refuge_efficient_b0_backbone.pt"
+REFUGE_EFFICIENT_B7_PATH    = _REPO_ROOT / "models/refuge/classifier/refuge_efficient_b7_backbone.pt"
+
+
+@dataclass(frozen=True)
+class BackboneSpec:
+    ctor:            Callable
+    weights_default: object
+    strip:           Callable[[nn.Module], tuple]
+    blocks:          Callable[[nn.Module], List[nn.Module]]
+
+
+# ---------------------------------------------------------------------------
+# Strip helpers
+# ---------------------------------------------------------------------------
+
+def _strip_efficientnet(m: models.EfficientNet):
+    out_dim = m.classifier[1].in_features
+    m.classifier = nn.Identity()
+    return out_dim, m
+
+def _strip_resnet(m: models.ResNet):
+    out_dim = m.fc.in_features
+    m.fc = nn.Identity()
+    return out_dim, m
+
+def _strip_densenet(m: models.DenseNet):
+    out_dim = m.classifier.in_features
+    m.classifier = nn.Identity()
+    return out_dim, m
+
+def _strip_vgg(m: models.VGG):
+    out_dim = m.classifier[0].in_features
+    m.classifier = nn.Identity()
+    return out_dim, m
+
+def _strip_mobilenet_v2(m: models.MobileNetV2):
+    out_dim = m.classifier[1].in_features
+    m.classifier = nn.Identity()
+    return out_dim, m
+
+def _strip_inception_v3(m: models.Inception3):
+    out_dim = m.fc.in_features
+    m.fc = nn.Identity()
+    m.aux_logits = False
+    m.AuxLogits = None
+    return out_dim, m
+
+
+# ---------------------------------------------------------------------------
+# Block splitters for ratio-based freezing
+# ---------------------------------------------------------------------------
+
+def _blocks_efficientnet(m: models.EfficientNet):
+    return list(m.features)
+
+def _blocks_resnet(m: models.ResNet):
+    stem = nn.Sequential(m.conv1, m.bn1, m.relu, m.maxpool)
+    return [stem, m.layer1, m.layer2, m.layer3, m.layer4]
+
+def _blocks_densenet(m: models.DenseNet):
+    f = m.features
+    stem = nn.Sequential(f.conv0, f.norm0, f.relu0, f.pool0)
+    return [stem, f.denseblock1, f.transition1, f.denseblock2, f.transition2,
+            f.denseblock3, f.transition3, f.denseblock4, f.norm5]
+
+def _blocks_vgg(m: models.VGG):
+    stages, cur = [], []
+    for mod in m.features:
+        cur.append(mod)
+        if isinstance(mod, nn.MaxPool2d):
+            stages.append(nn.Sequential(*cur)); cur = []
+    if cur:
+        stages.append(nn.Sequential(*cur))
+    return stages
+
+def _blocks_mobilenet_v2(m: models.MobileNetV2):
+    return list(m.features)
+
+def _blocks_inception_v3(m: models.Inception3):
+    return [child for name, child in m.named_children()
+            if name not in ("fc", "AuxLogits")]
+
+
+# ---------------------------------------------------------------------------
+# Registry
+# ---------------------------------------------------------------------------
+
+BACKBONES: Dict[str, BackboneSpec] = {
+    "efficientnet_b0": BackboneSpec(
+        ctor=models.efficientnet_b0,
+        weights_default=models.EfficientNet_B0_Weights.DEFAULT,
+        strip=_strip_efficientnet,
+        blocks=_blocks_efficientnet,
+    ),
+    "resnet50": BackboneSpec(
+        ctor=models.resnet50,
+        weights_default=models.ResNet50_Weights.DEFAULT,
+        strip=_strip_resnet,
+        blocks=_blocks_resnet,
+    ),
+    "densenet121": BackboneSpec(
+        ctor=models.densenet121,
+        weights_default=models.DenseNet121_Weights.DEFAULT,
+        strip=_strip_densenet,
+        blocks=_blocks_densenet,
+    ),
+    "vgg16": BackboneSpec(
+        ctor=models.vgg16,
+        weights_default=models.VGG16_Weights.DEFAULT,
+        strip=_strip_vgg,
+        blocks=_blocks_vgg,
+    ),
+    "mobilenet_v2": BackboneSpec(
+        ctor=models.mobilenet_v2,
+        weights_default=models.MobileNet_V2_Weights.DEFAULT,
+        strip=_strip_mobilenet_v2,
+        blocks=_blocks_mobilenet_v2,
+    ),
+    "inception_v3": BackboneSpec(
+        ctor=models.inception_v3,
+        weights_default=models.Inception_V3_Weights.DEFAULT,
+        strip=_strip_inception_v3,
+        blocks=_blocks_inception_v3,
+    ),
+    "refugelike": BackboneSpec(
+        ctor=models.resnet50,
+        weights_default=None,
+        strip=_strip_resnet,
+        blocks=_blocks_resnet,
+    ),
+    "refuge_densenet": BackboneSpec(
+        ctor=models.densenet121,
+        weights_default=None,
+        strip=_strip_densenet,
+        blocks=_blocks_densenet,
+    ),
+    "refuge_efficient_b0": BackboneSpec(
+        ctor=models.efficientnet_b0,
+        weights_default=None,
+        strip=_strip_efficientnet,
+        blocks=_blocks_efficientnet,
+    ),
+    "refuge_efficient_b7": BackboneSpec(
+        ctor=models.efficientnet_b7,
+        weights_default=None,
+        strip=_strip_efficientnet,
+        blocks=_blocks_efficientnet,
+    ),
+}
+
+
+def list_names() -> List[str]:
+    return list(BACKBONES.keys())
+
+
+def load_backbone_weights(key: str, model: nn.Module) -> None:
+    paths = {
+        "refugelike":         REFUGELIKE_BACKBONE_PATH,
+        "refuge_densenet":    REFUGE_DENSENET_PATH,
+        "refuge_efficient_b0": REFUGE_EFFICIENT_B0_PATH,
+        "refuge_efficient_b7": REFUGE_EFFICIENT_B7_PATH,
+    }
+    path = paths.get(key)
+    if path is None:
+        return
+    if not path.exists():
+        raise FileNotFoundError(
+            f"Custom backbone weights not found at {path}. "
+            "Export them via refuge_build.py --export-backbone first."
+        )
+    state = torch.load(path, map_location="cpu")
+    model.load_state_dict(state, strict=False)
+
+
+def build_backbone(name: str, freeze_ratio: float = 0.0) -> tuple[nn.Module, int, list]:
+    """Instantiate a backbone, strip its classifier head, apply freeze ratio.
+
+    Returns (model, out_dim, blocks) where blocks is the ordered list of
+    freezable units — callers use it to dynamically adjust freeze_ratio later.
+    """
+    key = (name or "").lower()
+    if key not in BACKBONES:
+        raise ValueError(f"Unknown backbone '{name}'. Available: {list_names()}")
+
+    spec = BACKBONES[key]
+    if spec.weights_default is not None:
+        m = spec.ctor(weights=spec.weights_default)
+    else:
+        m = spec.ctor(weights=None)
+    out_dim, m = spec.strip(m)
+    load_backbone_weights(key, m)
+
+    blocks  = spec.blocks(m)
+    fr      = max(0.0, min(1.0, float(freeze_ratio)))
+    n_freeze = int(math.floor(len(blocks) * fr))
+    for b in blocks[:n_freeze]:
+        for p in b.parameters():
+            p.requires_grad = False
+
+    return m, out_dim, blocks

v4/classes/accessory/se_block.py

@@ -0,0 +1,126 @@
+"""se_block — Squeeze-and-Excitation channel gating."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class SEBlock(nn.Module):
+    """SE-style channel gating for vectors, feature maps, and sequences.
+
+    Input shapes:
+      [N, C]       (vector)     — squeeze is identity
+      [N, C, H, W] (image map)  — squeeze over H, W
+      [N, T, C]    (sequence)   — squeeze over T
+
+    Gate modes:
+      residual (default): gate = 1 + tanh(MLP(s))  in (0, 2), identity at init
+      plain:              gate = sigmoid(MLP(s))    in (0, 1)
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        reduction: int = 16,
+        residual: bool = True,
+        identity_init: bool = True,
+    ):
+        super().__init__()
+        hid = max(1, dim // max(1, reduction))
+        self.fc1      = nn.Linear(dim, hid, bias=True)
+        self.act      = nn.ReLU(inplace=True)
+        self.fc2      = nn.Linear(hid, dim, bias=True)
+        self.residual = residual
+
+        if residual and identity_init:
+            nn.init.zeros_(self.fc2.weight)
+            nn.init.zeros_(self.fc2.bias)
+
+    def _squeeze(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() == 2:
+            return x
+        if x.dim() == 4:
+            return x.mean(dim=(2, 3))
+        if x.dim() == 3:
+            return x.mean(dim=1)
+        return x.view(x.size(0), -1)
+
+    def _broadcast(self, gate: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
+        if like.dim() == 2:
+            return gate
+        if like.dim() == 3:
+            return gate.unsqueeze(1)
+        if like.dim() == 4:
+            return gate.unsqueeze(-1).unsqueeze(-1)
+        return gate.view_as(like)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        s = self._squeeze(x)
+        u = self.fc2(self.act(self.fc1(s)))
+        gate = (1.0 + torch.tanh(u)) if self.residual else torch.sigmoid(u)
+        return x * self._broadcast(gate, x), gate
+
+
+class SEGateLogger:
+    """Running stats over SE gate activations across batches."""
+
+    def __init__(
+        self,
+        enabled: bool = True,
+        track_channels: bool = False,
+        dim: int | None = None,
+    ):
+        self.enabled        = enabled
+        self.track_channels = track_channels
+        self.dim            = dim
+        self.reset()
+
+    def reset(self) -> None:
+        self._n      = 0
+        self._sum    = 0.0
+        self._sum2   = 0.0
+        self._lt02   = 0
+        self._gt08   = 0
+        self._ch_sum   = None
+        self._ch_count = 0
+        if self.track_channels and self.dim is not None:
+            self._ch_sum = torch.zeros(self.dim, dtype=torch.float32)
+
+    @torch.no_grad()
+    def accumulate(self, gates: torch.Tensor) -> None:
+        if not self.enabled:
+            return
+        if gates.dim() == 4:
+            g = gates.mean(dim=(2, 3))
+        elif gates.dim() == 3:
+            g = gates.mean(dim=1)
+        elif gates.dim() == 2:
+            g = gates
+        else:
+            g = gates.view(gates.size(0), -1)
+        g = g.detach()
+        self._n    += g.numel()
+        self._sum  += g.sum().item()
+        self._sum2 += (g * g).sum().item()
+        self._lt02 += (g < 0.2).sum().item()
+        self._gt08 += (g > 0.8).sum().item()
+        if self._ch_sum is not None:
+            self._ch_sum  += g.sum(dim=0).cpu()
+            self._ch_count += g.size(0)
+
+    def get(self, reset: bool = True) -> dict | None:
+        if self._n == 0:
+            return None
+        mean = self._sum / self._n
+        var  = max(0.0, self._sum2 / self._n - mean * mean)
+        out  = {
+            "mean":        mean,
+            "std":         var ** 0.5,
+            "pct_lt_0.2":  self._lt02 / self._n,
+            "pct_gt_0.8":  self._gt08 / self._n,
+        }
+        if self._ch_sum is not None and self._ch_count > 0:
+            out["channel_mean"] = (self._ch_sum / float(self._ch_count)).tolist()
+        if reset:
+            self.reset()
+        return out

v4/classes/accessory/transforms.py

@@ -0,0 +1,66 @@
+"""transforms — image transform utilities for v4 towers."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+
+from torchvision import transforms
+
+from v4.classes.accessory.backbones import BACKBONES
+
+IMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406)
+IMAGENET_STD:  Tuple[float, float, float] = (0.229, 0.224, 0.225)
+
+
+@dataclass
+class ImageTransformConfig:
+    crop_size:    int                              = 224
+    resize_size:  int                              = 256
+    mean:         Tuple[float, float, float]       = IMAGENET_MEAN
+    std:          Tuple[float, float, float]       = IMAGENET_STD
+    augment:      bool                             = True
+    rotation_deg: int                              = 15
+    color_jitter: Tuple[float, float, float, float] = (0.1, 0.1, 0.1, 0.05)
+    hflip:        bool                             = True
+    vflip:        bool                             = True
+
+    def build(self) -> transforms.Compose:
+        ops = [
+            transforms.Resize(self.resize_size),
+            transforms.CenterCrop(self.crop_size),
+        ]
+        if self.augment:
+            if self.hflip:
+                ops.append(transforms.RandomHorizontalFlip())
+            if self.vflip:
+                ops.append(transforms.RandomVerticalFlip())
+            if self.rotation_deg:
+                ops.append(transforms.RandomRotation(self.rotation_deg))
+            if self.color_jitter:
+                ops.append(transforms.ColorJitter(*self.color_jitter))
+        ops += [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=self.mean, std=self.std),
+        ]
+        return transforms.Compose(ops)
+
+
+def backbone_transform_config(backbone_name: str, augment: bool = True) -> ImageTransformConfig:
+    """Build an ImageTransformConfig using the backbone's default normalisation stats."""
+    key = (backbone_name or "").lower()
+    if key not in BACKBONES:
+        raise ValueError(f"Unknown backbone '{backbone_name}'.")
+    spec = BACKBONES[key]
+    mean = getattr(spec.weights_default, "meta", {}).get("mean", IMAGENET_MEAN)
+    std  = getattr(spec.weights_default, "meta", {}).get("std",  IMAGENET_STD)
+    crop = 299 if key == "inception_v3" else 224
+    return ImageTransformConfig(crop_size=crop, mean=mean, std=std, augment=augment)
+
+
+def build_backbone_transform(backbone_name: str, augment: bool = True) -> transforms.Compose:
+    return backbone_transform_config(backbone_name, augment=augment).build()
+
+
+def build_eval_transform(backbone_name: str) -> transforms.Compose:
+    """Deterministic eval transform — no augmentation, backbone-matched normalisation."""
+    return build_backbone_transform(backbone_name, augment=False)

v4/classes/bridges/fusion_bridge.py

@@ -0,0 +1,57 @@
+"""fusion_bridge — FusionBridge: N-input Hadamard-product fusion, embedding output only."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from v4.classes.accessory.se_block import SEBlock
+
+
+class FusionBridge(nn.Module):
+    """Project N input embeddings to a shared dim, fuse via element-wise product.
+
+    Pure embedding producer — no classification head.  Attach a head stage in
+    the pipeline config to produce logits.
+
+    Parameters
+    ----------
+    input_dims   : ordered list of input embedding dims
+    fusion_dim   : projection / output dimension
+    use_se       : SE gate on the fused vector
+    se_reduction : SE reduction factor
+    se_pre_norm  : LayerNorm before each projection; else Identity
+    """
+
+    def __init__(
+        self,
+        input_dims:   list[int],
+        fusion_dim:   int  = 256,
+        use_se:       bool = True,
+        se_reduction: int  = 16,
+        se_pre_norm:  bool = True,
+    ):
+        super().__init__()
+        self.out_dim = fusion_dim
+        self.W  = nn.ModuleList([nn.Linear(d, fusion_dim) for d in input_dims])
+        self.ln = nn.ModuleList(
+            [nn.LayerNorm(fusion_dim) if se_pre_norm else nn.Identity()
+             for _ in input_dims]
+        )
+        self.se = SEBlock(fusion_dim, reduction=se_reduction, residual=True) if use_se else None
+
+    def forward(self, embeddings: list[torch.Tensor]) -> torch.Tensor:
+        assert len(embeddings) == len(self.W), (
+            f"FusionBridge expects {len(self.W)} inputs, got {len(embeddings)}"
+        )
+        h = self.ln[0](self.W[0](embeddings[0]))
+        for i in range(1, len(embeddings)):
+            h = h * self.ln[i](self.W[i](embeddings[i]))
+        if self.se is not None:
+            h, _ = self.se(h)
+        return h
+
+    def set_phase(self, phase: str) -> None:
+        """Freeze bridge during tower_warmup; trainable otherwise."""
+        enabled = phase not in ("tower_warmup", "cd_warmup")
+        for p in self.parameters():
+            p.requires_grad_(enabled)

v4/classes/bridges/hyperbridge.py

@@ -0,0 +1,56 @@
+"""hyperbridge — HyperBridge: bilateral fusion over paired embeddings, embedding output only."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class HyperBridge(nn.Module):
+    """Fuse side embeddings (e.g. two z_fused vectors) into a single embedding.
+
+    Pure embedding producer — no classification head.  Attach a head stage in
+    the pipeline config to produce logits.
+
+    Modes
+    -----
+    embedding_mlp (default)
+        Linear projection of concatenated inputs → hidden_dim embedding.
+
+    classic_bridge
+        Per-side projection → Hadamard product → hidden_dim embedding.
+
+    Parameters
+    ----------
+    input_dims  : {side_key: dim} — e.g. {"a": 256, "b": 256}
+    hidden_dim  : output embedding dimension
+    mode        : "embedding_mlp" | "classic_bridge"
+    """
+
+    def __init__(
+        self,
+        input_dims: dict[str, int],
+        hidden_dim: int = 256,
+        mode:       str = "embedding_mlp",
+    ):
+        super().__init__()
+        self.input_names = list(input_dims.keys())
+        self.mode        = mode
+        self.out_dim     = hidden_dim
+        dims             = list(input_dims.values())
+
+        if mode == "embedding_mlp":
+            self.proj = nn.Linear(sum(dims), hidden_dim)
+        elif mode == "classic_bridge":
+            self.W  = nn.ModuleList([nn.Linear(d, hidden_dim) for d in dims])
+            self.ln = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in dims])
+        else:
+            raise ValueError(f"Unknown HyperBridge mode: {mode!r}")
+
+    def forward(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+        ordered = [inputs[name] for name in self.input_names]
+        if self.mode == "embedding_mlp":
+            return self.proj(torch.cat(ordered, dim=1))
+        h = self.ln[0](self.W[0](ordered[0]))
+        for i in range(1, len(ordered)):
+            h = h * self.ln[i](self.W[i](ordered[i]))
+        return h

v4/classes/dataset.py

@@ -0,0 +1,248 @@
+"""dataset — data packaging for v4: shells, DataBundle, HTDataset.
+
+ShellEntry / LoaderShell
+    Minimal, data-free structures representing *who* to sample and in what
+    order.  entity_id is opaque to the orchestrator; towers interpret it.
+
+DataBundle
+    Accumulates per-eye DataFrames and derives scalar/categorical stats used
+    by ClinicalDataView.  No kfold, no vectorization — those live in the
+    profile and orchestrator respectively.
+
+HTDataset / ht_collate
+    PyTorch Dataset that delegates sample retrieval to towers via get_sample.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+
+# ---------------------------------------------------------------------------
+# Shell
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ShellEntry:
+    """One sample slot in a LoaderShell.
+
+    entity_id : opaque — defined by the profile, interpreted by towers.
+    label     : integer class label.
+    meta      : per-entry context a profile wants to pass through.
+    """
+    entity_id: Any
+    label:     int
+    meta:      dict = field(default_factory=dict)
+
+
+@dataclass
+class LoaderShell:
+    """An ordered sequence of ShellEntry objects for one split/fold."""
+    entries: list[ShellEntry]
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def __iter__(self):
+        return iter(self.entries)
+
+
+# ---------------------------------------------------------------------------
+# DataBundle
+# ---------------------------------------------------------------------------
+
+class DataBundle:
+    """Accumulates per-eye DataFrames and derives feature metadata.
+
+    Responsibilities: column type inference, scalar stats (min/max/median),
+    categorical index maps, and feature dim.  Everything else (splits,
+    vectorization, image paths) lives in the profile that uses this bundle.
+    """
+
+    def __init__(
+        self,
+        *,
+        image_dir: str,
+        clinical_dir: Optional[str] = None,
+        label_col: str,
+        patient_col: str = "Patient ID",
+        cat_cols: Optional[Iterable[str]] = None,
+        max_unique_for_cat: int = 4,
+        n_splits: int = 5,
+        random_seed: int = 42,
+        filename_template: str = "RET{pid:03d}{eye}.jpg",
+    ) -> None:
+        self.image_dir         = Path(image_dir)
+        self.label_col         = label_col
+        self.patient_col       = patient_col
+        self.max_unique_for_cat = max_unique_for_cat
+        self.filename_template = filename_template
+        self.clinical_dir      = Path(clinical_dir) if clinical_dir else None
+
+        self.frames:       List[pd.DataFrame]           = []
+        self.df:           pd.DataFrame                 = pd.DataFrame()
+        self.scalar_cols:  List[str]                    = []
+        self.cat_cols:     List[str]                    = list(cat_cols) if cat_cols else []
+        self.scalar_stats: Dict[str, Dict[str, float]]  = {}
+        self.cat_maps:     Dict[str, Dict[object, int]] = {}
+        self.feature_dim:  int                          = 0
+
+    def add_df(
+        self,
+        df: pd.DataFrame,
+        *,
+        id_column: Optional[str] = None,
+        exclude_cols: Optional[Iterable[str]] = None,
+    ) -> None:
+        df = df.copy()
+        self._ensure_patient_id(df, id_column)
+        if self.label_col not in df.columns:
+            raise ValueError(f"label_col '{self.label_col}' not found in added dataframe")
+        self.frames.append(df)
+        self._refresh_master_df(exclude_cols=exclude_cols)
+        self._infer_or_validate_feature_types(exclude_cols=exclude_cols)
+        self._compute_numeric_stats()
+        self._build_cat_maps()
+        self._compute_feature_dim()
+
+    # ── Internal ─────────────────────────────────────────────────────────────
+
+    def _ensure_patient_id(self, df: pd.DataFrame, id_column: Optional[str]) -> None:
+        if self.patient_col in df.columns:
+            return
+        if id_column and id_column in df.columns:
+            df.rename(columns={id_column: self.patient_col}, inplace=True)
+            return
+        candidates = [
+            c for c in df.columns
+            if c.lower().replace(" ", "") in {"patientid", "patient", "pid"}
+        ]
+        if len(candidates) == 1:
+            df.rename(columns={candidates[0]: self.patient_col}, inplace=True)
+            return
+        raise ValueError(
+            f"A '{self.patient_col}' column is required; "
+            f"provide id_column=... if it has a different name."
+        )
+
+    def _refresh_master_df(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
+        self.df = pd.concat(self.frames, axis=0, ignore_index=True)
+        if exclude_cols:
+            self.df = self.df.drop(columns=[c for c in exclude_cols if c in self.df.columns])
+
+    def _infer_or_validate_feature_types(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
+        excluded = set(exclude_cols or []) | {self.label_col, self.patient_col}
+        candidates = [c for c in self.df.columns if c not in excluded]
+        cats = set(self.cat_cols)
+        scalars: set[str] = set()
+        for c in candidates:
+            if c in cats:
+                continue
+            s       = self.df[c]
+            as_num  = pd.to_numeric(s, errors="coerce")
+            n_uniq  = s.dropna().nunique()
+            if as_num.notna().any() and as_num.isna().mean() < 1.0 and n_uniq > self.max_unique_for_cat:
+                scalars.add(c)
+            else:
+                cats.add(c)
+        self.cat_cols    = sorted(cats)
+        self.scalar_cols = sorted(scalars)
+
+    def _compute_numeric_stats(self) -> None:
+        self.scalar_stats.clear()
+        for col in self.scalar_cols:
+            vals = pd.to_numeric(self.df[col], errors="coerce").dropna().astype(float).values
+            if vals.size == 0:
+                lo, hi, med = 0.0, 1.0, 0.0
+            else:
+                lo, hi = float(np.min(vals)), float(np.max(vals))
+                med    = float(np.median(vals))
+                if hi <= lo:
+                    hi = lo + 1.0
+            self.scalar_stats[col] = {"min": lo, "max": hi, "median": med}
+
+    def _build_cat_maps(self) -> None:
+        self.cat_maps.clear()
+        for col in self.cat_cols:
+            cats = [v for v in self.df[col].dropna().unique().tolist()]
+            try:
+                cats = sorted(cats)
+            except Exception:
+                pass
+            mapping: Dict[object, int] = {"<UNK>": 0}
+            for i, v in enumerate(cats, start=1):
+                mapping[v] = i
+            self.cat_maps[col] = mapping
+
+    def _compute_feature_dim(self) -> None:
+        self.feature_dim = (
+            len(self.scalar_cols)
+            + sum(len(m) for m in self.cat_maps.values())
+            + len(self.scalar_cols)
+        )
+
+
+# ---------------------------------------------------------------------------
+# HTDataset / ht_collate
+# ---------------------------------------------------------------------------
+
+class HTDataset(Dataset):
+    """PyTorch Dataset backed by a LoaderShell.
+
+    Delegates sample retrieval to each tower's get_sample(entry).
+    Batch keys are tower names plus "label".
+    """
+
+    def __init__(self, shell: LoaderShell, towers: dict) -> None:
+        self.entries = shell.entries
+        self.towers  = towers
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        entry  = self.entries[idx]
+        sample = {
+            "label":     torch.tensor(entry.label, dtype=torch.long),
+            "entity_id": entry.entity_id,
+        }
+        for name, tower in self.towers.items():
+            sample[name] = tower.get_sample(entry)
+        return sample
+
+
+def to_label_tensor(labels, device: torch.device) -> torch.Tensor:
+    """Normalise a batch of labels (tensor or list) to a long tensor on device."""
+    if torch.is_tensor(labels):
+        return labels.to(device=device, dtype=torch.long)
+    return torch.as_tensor(labels, dtype=torch.long, device=device)
+
+
+def ht_collate(batch: list[dict[str, Any]]) -> dict[str, Any]:
+    """Collate HTDataset samples.
+
+    Tensor values are stacked; dict values (side dicts from patient-level
+    shells) are stacked per inner key; everything else becomes a list.
+    """
+    if not batch:
+        return {}
+    result: dict[str, Any] = {}
+    for key in batch[0]:
+        vals  = [s[key] for s in batch]
+        first = vals[0]
+        if isinstance(first, torch.Tensor):
+            result[key] = torch.stack(vals, dim=0)
+        elif isinstance(first, dict):
+            result[key] = {
+                side: torch.stack([v[side] for v in vals], dim=0)
+                for side in first
+            }
+        else:
+            result[key] = vals
+    return result

v4/classes/heads/classifier.py

@@ -0,0 +1,24 @@
+"""classifier — ClassificationHead output head."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class ClassificationHead(nn.Module):
+    """Minimal classification head: ReLU → Dropout → Linear(in_dim → num_classes).
+
+    Used as the output stage of bridges and any module that needs a reusable,
+    swappable task head producing class logits.
+    """
+
+    def __init__(self, in_dim: int, num_classes: int, dropout: float = 0.5):
+        super().__init__()
+        self.head = nn.Sequential(
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(in_dim, num_classes),
+        )
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        return self.head(z)

v4/classes/loaders/image_loader.py

@@ -0,0 +1,132 @@
+"""image_loader — CachedImageLoader: disk I/O with optional in-memory cache.
+
+A single instance is shared across all towers/datasets for a run so that
+images are decoded from disk at most once.  The optional preprocessor runs
+at cache-fill time (resize, deterministic crop, etc.) so that per-batch
+transforms in the image tower only need to apply stochastic augmentations.
+
+Usage
+-----
+    loader = CachedImageLoader(enabled=True, workers=4)
+    loader.warm(paths, preprocessor=resize_fn)   # optional parallel pre-fill
+    img = loader.load(path, preprocessor=resize_fn)  # returns PIL Image
+"""
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Callable, Iterable, Optional
+
+import numpy as np
+from PIL import Image
+
+
+class CachedImageLoader:
+    """Loads PIL Images from disk with an optional shared in-memory cache.
+
+    The cache stores decoded, pre-preprocessed images as uint8 numpy arrays
+    (RGB, HWC).  Storing after the preprocessor runs means the deterministic
+    resize/crop step executes only once per image across all folds and epochs.
+
+    Parameters
+    ----------
+    enabled : bool
+        When False the cache is bypassed and every call hits disk.
+    workers : int
+        Thread count for ``warm()``.  0 or 1 → single-threaded.
+    """
+
+    def __init__(self, *, enabled: bool = True, workers: int = 4) -> None:
+        self._cache: dict[str, np.ndarray] | None = {} if enabled else None
+        self._workers = workers
+
+    @property
+    def enabled(self) -> bool:
+        return self._cache is not None
+
+    def __len__(self) -> int:
+        return len(self._cache) if self._cache is not None else 0
+
+    def load(
+        self,
+        path: str | Path,
+        preprocessor: Optional[Callable[..., Image.Image]] = None,
+    ) -> Image.Image:
+        """Return a PIL Image for *path*, using the cache when enabled."""
+        key = str(path)
+        if self._cache is not None:
+            cached = self._cache.get(key)
+            if cached is not None:
+                return Image.fromarray(cached, mode="RGB")
+
+        img = Image.open(path).convert("RGB")
+        if preprocessor is not None:
+            img = call_preprocessor(preprocessor, img, path)
+
+        if self._cache is not None:
+            self._cache[key] = np.asarray(img, dtype=np.uint8)
+
+        return img
+
+    def warm(
+        self,
+        paths: Iterable[str | Path],
+        preprocessor: Optional[Callable[..., Image.Image]] = None,
+    ) -> None:
+        """Pre-populate the cache for all *paths* (no-op when disabled).
+
+        Already-cached paths are skipped, so calling warm() multiple times
+        (e.g. once per fold) is safe.
+        """
+        if self._cache is None:
+            return
+        paths = list(paths)
+        to_warm = [str(p) for p in paths if str(p) not in self._cache]
+        if not to_warm:
+            return
+        already = len(paths) - len(to_warm)
+        print(
+            f"[image_cache] warming {len(to_warm)} images"
+            + (f" ({already} already cached)" if already else ""),
+            flush=True,
+        )
+
+        def _warm_one(path_str: str) -> None:
+            if path_str in self._cache:
+                return
+            img = Image.open(path_str).convert("RGB")
+            if preprocessor is not None:
+                img = call_preprocessor(preprocessor, img, Path(path_str))
+            self._cache[path_str] = np.asarray(img, dtype=np.uint8)
+
+        try:
+            from tqdm import tqdm
+        except ImportError:
+            tqdm = None
+
+        if self._workers <= 1:
+            it = tqdm(to_warm, desc="Warm image cache", unit="img") if tqdm else to_warm
+            for p in it:
+                _warm_one(p)
+        else:
+            with ThreadPoolExecutor(max_workers=self._workers) as ex:
+                futures = {ex.submit(_warm_one, p): p for p in to_warm}
+                it = (
+                    tqdm(as_completed(futures), total=len(futures),
+                         desc="Warm image cache", unit="img")
+                    if tqdm else as_completed(futures)
+                )
+                for fut in it:
+                    fut.result()
+
+
+def call_preprocessor(
+    fn: Callable[..., Image.Image],
+    img: Image.Image,
+    path: Path | str,
+) -> Image.Image:
+    """Call preprocessor with (img, path) or just (img) depending on arity."""
+    try:
+        return fn(img, path)
+    except TypeError:
+        return fn(img)

v4/classes/logging/prediction_store.py

@@ -0,0 +1,394 @@
+"""prediction_store — per-epoch logit and embedding recording across folds and phases.
+
+PredictionStore  — records logits for every head, fold, phase, and epoch.
+FeatureStore     — records embeddings (opt-in); same structure but per-head
+                   tensors since embedding dims vary across heads.
+
+HDF5 layout — PredictionStore
+------------------------------
+/{phase}/logits          float32  (n_folds, n_epochs, n_samples, n_heads, n_classes)
+/{phase}/head_names      str      (n_heads,)
+/{phase}/y_true          int64    (n_samples,)
+/{phase}/entity_id_{k}  int64|str (n_samples,)  — one dataset per id component
+/{phase}/split           str      (n_folds, n_samples)
+/{phase}/loss            float32  (n_folds, n_epochs)
+
+HDF5 layout — FeatureStore
+---------------------------
+/{phase}/{head_name}     float32  (n_folds, n_epochs, n_samples, embedding_dim)
+/{phase}/y_true          int64    (n_samples,)
+/{phase}/entity_id_{k}  int64|str (n_samples,)
+/{phase}/split           str      (n_folds, n_samples)
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Sequence
+
+import numpy as np
+
+try:
+    import h5py
+except ImportError as e:
+    raise ImportError("PredictionStore requires h5py: pip install h5py") from e
+
+_STR_DT = h5py.string_dtype()
+
+
+# ---------------------------------------------------------------------------
+# Internal phase buffer
+# ---------------------------------------------------------------------------
+
+class _PhaseBuffer:
+    def __init__(
+        self,
+        entity_ids: list[tuple],
+        y_true:     np.ndarray,
+        head_names: list[str],
+        n_epochs:   int,
+        n_folds:    int,
+        n_classes:  int,
+    ):
+        n_s = len(entity_ids)
+        n_h = len(head_names)
+        self.entity_ids = list(entity_ids)
+        self.y_true     = np.asarray(y_true, dtype=np.int64)
+        self.head_names = list(head_names)
+        self.n_epochs   = n_epochs
+        self.logits     = np.full((n_folds, n_epochs, n_s, n_h, n_classes), np.nan, dtype=np.float32)
+        self.split      = np.full((n_folds, n_s), "", dtype=object)
+        self.loss       = np.full((n_folds, n_epochs), np.nan, dtype=np.float32)
+        self._sid       = {str(eid): i for i, eid in enumerate(entity_ids)}
+        self._hid       = {h: i for i, h in enumerate(head_names)}
+
+
+# ---------------------------------------------------------------------------
+# Internal feature buffer (per-head, variable embedding_dim)
+# ---------------------------------------------------------------------------
+
+class _FeaturePhaseBuffer:
+    def __init__(
+        self,
+        entity_ids: list[tuple],
+        y_true:     np.ndarray,
+        n_folds:    int,
+    ):
+        self.entity_ids = list(entity_ids)
+        self.y_true     = np.asarray(y_true, dtype=np.int64)
+        self.split      = np.full((n_folds, len(entity_ids)), "", dtype=object)
+        self._sid       = {str(eid): i for i, eid in enumerate(entity_ids)}
+        # head_name → (buffer array, n_epochs)
+        self._heads: dict[str, tuple[np.ndarray, int]] = {}
+
+    def register_head(self, head: str, n_epochs: int, embedding_dim: int, n_folds: int) -> None:
+        n_s = len(self.entity_ids)
+        self._heads[head] = (
+            np.full((n_folds, n_epochs, n_s, embedding_dim), np.nan, dtype=np.float32),
+            n_epochs,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _write_entity_ids(grp: h5py.Group, entity_ids: list[tuple]) -> None:
+    if not entity_ids:
+        return
+    n_components = max(len(eid) for eid in entity_ids)
+    for k in range(n_components):
+        vals = [eid[k] if k < len(eid) else "" for eid in entity_ids]
+        if all(isinstance(v, (int, np.integer)) for v in vals):
+            grp.create_dataset(f"entity_id_{k}", data=np.array(vals, dtype=np.int64))
+        else:
+            grp.create_dataset(f"entity_id_{k}", data=np.array(vals, dtype=object), dtype=_STR_DT)
+
+
+def _read_entity_ids(grp: h5py.Group, n_samples: int) -> list[tuple]:
+    k, components = 0, []
+    while f"entity_id_{k}" in grp:
+        arr = grp[f"entity_id_{k}"][:]
+        if arr.dtype.kind in ("S", "O", "U"):
+            arr = np.array([v.decode() if isinstance(v, bytes) else str(v) for v in arr])
+        components.append(arr)
+        k += 1
+    if not components:
+        return [() for _ in range(n_samples)]
+    return [tuple(c[i] for c in components) for i in range(n_samples)]
+
+
+def _decode_str_array(arr: np.ndarray) -> list[str]:
+    return [v.decode() if isinstance(v, bytes) else str(v) for v in arr.flat]
+
+
+# ---------------------------------------------------------------------------
+# PredictionStore
+# ---------------------------------------------------------------------------
+
+class PredictionStore:
+    """Records per-epoch logits across all folds and phases, saves to HDF5.
+
+    The store is generic — it knows nothing about what heads or phases exist.
+    The orchestrator registers phases and records whatever heads it builds.
+    """
+
+    def __init__(self, n_folds: int, n_classes: int) -> None:
+        self.n_folds   = n_folds
+        self.n_classes = n_classes
+        self._phases: dict[str, _PhaseBuffer] = {}
+
+    def register_phase(
+        self,
+        phase:      str,
+        entity_ids: list[tuple],
+        y_true:     Sequence[int],
+        head_names: list[str],
+        n_epochs:   int,
+    ) -> None:
+        """Register a training phase before recording begins."""
+        self._phases[phase] = _PhaseBuffer(
+            entity_ids=list(entity_ids),
+            y_true=np.asarray(y_true, dtype=np.int64),
+            head_names=list(head_names),
+            n_epochs=n_epochs,
+            n_folds=self.n_folds,
+            n_classes=self.n_classes,
+        )
+
+    def record(
+        self,
+        phase:      str,
+        fold:       int,
+        epoch:      int,
+        entity_ids: Sequence[tuple],
+        head:       str,
+        logits:     np.ndarray,
+    ) -> None:
+        """Record a batch of logits for one head at one epoch."""
+        buf  = self._phases[phase]
+        hidx = buf._hid.get(head)
+        if hidx is None:
+            return
+        for i, eid in enumerate(entity_ids):
+            sidx = buf._sid.get(str(eid))
+            if sidx is not None:
+                buf.logits[fold, epoch, sidx, hidx, :] = logits[i]
+
+    def record_loss(self, phase: str, fold: int, epoch: int, loss: float) -> None:
+        self._phases[phase].loss[fold, epoch] = float(loss)
+
+    def set_split(
+        self,
+        phase:      str,
+        fold:       int,
+        entity_ids: Sequence[tuple],
+        label:      str,
+    ) -> None:
+        """Mark samples as 'train', 'val', or 'test' for a fold."""
+        buf = self._phases[phase]
+        for eid in entity_ids:
+            sidx = buf._sid.get(str(eid))
+            if sidx is not None:
+                buf.split[fold, sidx] = label
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+
+    def save(self, path: str | Path) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with h5py.File(path, "w") as f:
+            for phase, buf in self._phases.items():
+                grp = f.create_group(phase)
+                grp.create_dataset("logits",     data=buf.logits, compression="gzip", compression_opts=4)
+                grp.create_dataset("y_true",     data=buf.y_true)
+                grp.create_dataset("loss",       data=buf.loss)
+                grp.create_dataset("head_names", data=np.array(buf.head_names, dtype=object), dtype=_STR_DT)
+                grp.create_dataset("split",      data=buf.split.astype(str),                  dtype=_STR_DT)
+                _write_entity_ids(grp, buf.entity_ids)
+
+    @classmethod
+    def load(cls, path: str | Path) -> "PredictionStore":
+        """Load all phases into memory."""
+        with h5py.File(path, "r") as f:
+            first  = next(iter(f.values()))
+            n_folds, _, _, _, n_classes = first["logits"].shape
+            store  = cls(n_folds=n_folds, n_classes=n_classes)
+            for phase in f:
+                grp        = f[phase]
+                logits     = grp["logits"][:]
+                n_folds_, n_epochs, n_samples, n_heads, _ = logits.shape
+                head_names = _decode_str_array(grp["head_names"][:])
+                entity_ids = _read_entity_ids(grp, n_samples)
+                buf            = _PhaseBuffer(
+                    entity_ids=entity_ids,
+                    y_true=grp["y_true"][:],
+                    head_names=head_names,
+                    n_epochs=n_epochs,
+                    n_folds=n_folds_,
+                    n_classes=n_classes,
+                )
+                buf.logits     = logits
+                buf.loss       = grp["loss"][:]
+                split_raw      = grp["split"][:]
+                buf.split      = np.array(
+                    [[v.decode() if isinstance(v, bytes) else str(v) for v in row]
+                     for row in split_raw],
+                    dtype=object,
+                )
+                store._phases[phase] = buf
+        return store
+
+    # ------------------------------------------------------------------
+    # Query helpers
+    # ------------------------------------------------------------------
+
+    @property
+    def phases(self) -> list[str]:
+        return list(self._phases.keys())
+
+    def head_names(self, phase: str) -> list[str]:
+        return self._phases[phase].head_names
+
+    def entity_ids(self, phase: str) -> list[tuple]:
+        return self._phases[phase].entity_ids
+
+    def get_logits(
+        self,
+        phase: str,
+        head:  str,
+        fold:  int | None = None,
+        epoch: int | None = None,
+    ) -> np.ndarray:
+        """Slice logits for one head. Unspecified dims return the full axis.
+
+        Returns shape (folds, epochs, samples, classes) by default,
+        with leading dims dropped for each specified index.
+        """
+        buf  = self._phases[phase]
+        hidx = buf._hid[head]
+        data = buf.logits[:, :, :, hidx, :]       # (folds, epochs, samples, classes)
+        if fold  is not None: data = data[fold]   # (epochs, samples, classes)
+        if epoch is not None: data = data[..., epoch, :, :]  if fold is None else data[epoch]
+        return data
+
+    def get_split(self, phase: str, fold: int) -> dict[str, list[tuple]]:
+        """Return {'train': [...], 'val': [...], 'test': [...]} entity_id lists."""
+        buf    = self._phases[phase]
+        labels = buf.split[fold]
+        out: dict[str, list[tuple]] = {}
+        for eid, lbl in zip(buf.entity_ids, labels):
+            out.setdefault(lbl, []).append(eid)
+        return out
+
+
+# ---------------------------------------------------------------------------
+# FeatureStore
+# ---------------------------------------------------------------------------
+
+class FeatureStore:
+    """Records per-epoch embeddings (variable dim per head), saves to HDF5.
+
+    Opt-in companion to PredictionStore. Typically written only on checkpoint
+    runs where you want to do dimensionality reduction or cluster analysis.
+    """
+
+    def __init__(self, n_folds: int) -> None:
+        self.n_folds = n_folds
+        self._phases: dict[str, _FeaturePhaseBuffer] = {}
+
+    def register_phase(
+        self,
+        phase:      str,
+        entity_ids: list[tuple],
+        y_true:     Sequence[int],
+    ) -> None:
+        self._phases[phase] = _FeaturePhaseBuffer(
+            entity_ids=list(entity_ids),
+            y_true=np.asarray(y_true, dtype=np.int64),
+            n_folds=self.n_folds,
+        )
+
+    def register_head(
+        self,
+        phase:         str,
+        head:          str,
+        n_epochs:      int,
+        embedding_dim: int,
+    ) -> None:
+        self._phases[phase].register_head(head, n_epochs, embedding_dim, self.n_folds)
+
+    def record(
+        self,
+        phase:      str,
+        fold:       int,
+        epoch:      int,
+        entity_ids: Sequence[tuple],
+        head:       str,
+        embeddings: np.ndarray,
+    ) -> None:
+        buf   = self._phases[phase]
+        arr, _= buf._heads[head]
+        for i, eid in enumerate(entity_ids):
+            sidx = buf._sid.get(str(eid))
+            if sidx is not None:
+                arr[fold, epoch, sidx, :] = embeddings[i]
+
+    def set_split(
+        self,
+        phase:      str,
+        fold:       int,
+        entity_ids: Sequence[tuple],
+        label:      str,
+    ) -> None:
+        buf = self._phases[phase]
+        for eid in entity_ids:
+            sidx = buf._sid.get(str(eid))
+            if sidx is not None:
+                buf.split[fold, sidx] = label
+
+    def save(self, path: str | Path) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with h5py.File(path, "w") as f:
+            for phase, buf in self._phases.items():
+                grp = f.create_group(phase)
+                grp.create_dataset("y_true", data=buf.y_true)
+                grp.create_dataset("split",  data=buf.split.astype(str), dtype=_STR_DT)
+                _write_entity_ids(grp, buf.entity_ids)
+                for head, (arr, _) in buf._heads.items():
+                    grp.create_dataset(head, data=arr, compression="gzip", compression_opts=4)
+
+    @classmethod
+    def load(cls, path: str | Path) -> "FeatureStore":
+        with h5py.File(path, "r") as f:
+            n_folds = next(
+                arr.shape[0]
+                for grp in f.values()
+                for k, arr in grp.items()
+                if k not in ("y_true", "split") and not k.startswith("entity_id_")
+            )
+            store = cls(n_folds=n_folds)
+            _meta = {"y_true", "split"}
+            for phase in f:
+                grp        = f[phase]
+                n_samples  = grp["y_true"].shape[0]
+                entity_ids = _read_entity_ids(grp, n_samples)
+                buf        = _FeaturePhaseBuffer(
+                    entity_ids=entity_ids,
+                    y_true=grp["y_true"][:],
+                    n_folds=n_folds,
+                )
+                buf.split = np.array(
+                    [[v.decode() if isinstance(v, bytes) else str(v) for v in row]
+                     for row in grp["split"][:]],
+                    dtype=object,
+                )
+                for key in grp:
+                    if key in _meta or key.startswith("entity_id_"):
+                        continue
+                    arr = grp[key][:]
+                    buf._heads[key] = (arr, arr.shape[1])
+                store._phases[phase] = buf
+        return store

v4/classes/metrics.py

@@ -0,0 +1,175 @@
+"""metrics — loss, scoring, calibration, and threshold/bias tuning."""
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from sklearn.metrics import (
+    cohen_kappa_score,
+    f1_score,
+    matthews_corrcoef,
+    recall_score,
+    roc_auc_score,
+    roc_curve,
+)
+
+
+# ---------------------------------------------------------------------------
+# Loss
+# ---------------------------------------------------------------------------
+
+def focal_loss(
+    logits:    torch.Tensor,
+    targets:   torch.Tensor,
+    gamma:     float = 0.0,
+    weight:    Optional[torch.Tensor] = None,
+    reduction: str = "mean",
+) -> torch.Tensor:
+    """Focal loss; reduces to cross-entropy when gamma=0."""
+    if gamma <= 0:
+        return F.cross_entropy(logits, targets, weight=weight, reduction=reduction)
+    log_probs = F.log_softmax(logits, dim=1)
+    probs     = log_probs.exp()
+    targets   = targets.long().view(-1, 1)
+    logpt     = log_probs.gather(1, targets)
+    pt        = probs.gather(1, targets)
+    loss      = -(((1.0 - pt).clamp_min(0.0) ** gamma) * logpt)
+    if weight is not None:
+        loss = loss * weight.gather(0, targets.view(-1)).view(-1, 1)
+    loss = loss.view(-1)
+    if reduction == "sum":  return loss.sum()
+    if reduction == "mean": return loss.mean()
+    return loss
+
+
+# ---------------------------------------------------------------------------
+# Basic array scoring
+# ---------------------------------------------------------------------------
+
+def score_arrays(y_true: np.ndarray, probs: np.ndarray, num_classes: int):
+    """Return (acc, auc, n)."""
+    if y_true.size == 0:
+        return float("nan"), float("nan"), 0
+    acc = float((probs.argmax(1) == y_true).mean())
+    try:
+        auc = (
+            float(roc_auc_score(y_true, probs[:, 1]))
+            if num_classes == 2
+            else float(roc_auc_score(y_true, probs, multi_class="ovr", average="macro"))
+        )
+    except Exception:
+        auc = float("nan")
+    return acc, auc, int(len(y_true))
+
+
+# ---------------------------------------------------------------------------
+# Calibration
+# ---------------------------------------------------------------------------
+
+def compute_ece(y_true: np.ndarray, probs: np.ndarray, n_bins: int = 10) -> float:
+    """Expected Calibration Error: weighted mean |confidence − accuracy| per bin."""
+    if y_true.size == 0:
+        return float("nan")
+    confidences = probs.max(axis=1)
+    predictions = probs.argmax(axis=1)
+    bin_edges   = np.linspace(0.0, 1.0, n_bins + 1)
+    ece = 0.0
+    n   = len(y_true)
+    for i, (lo, hi) in enumerate(zip(bin_edges[:-1], bin_edges[1:])):
+        mask = (confidences >= lo) & (
+            confidences <= hi if i == n_bins - 1 else confidences < hi
+        )
+        if not mask.any():
+            continue
+        ece += float(mask.sum()) / n * abs(
+            float(confidences[mask].mean()) - float((predictions[mask] == y_true[mask]).mean())
+        )
+    return float(ece)
+
+
+def compute_extended_metrics(
+    y_true:          np.ndarray,
+    probs:           np.ndarray,
+    num_classes:     int,
+    n_bins:          int = 10,
+    preds_override:  Optional[np.ndarray] = None,
+) -> dict:
+    nan = float("nan")
+    if y_true.size == 0:
+        return dict(
+            kappa=nan, mcc=nan, macro_f1=nan,
+            per_class_recall=np.full(num_classes, nan), ece=nan,
+        )
+    preds = preds_override if preds_override is not None else probs.argmax(axis=1)
+    try:    kappa    = float(cohen_kappa_score(y_true, preds))
+    except: kappa    = nan
+    try:    mcc      = float(matthews_corrcoef(y_true, preds))
+    except: mcc      = nan
+    try:    macro_f1 = float(f1_score(y_true, preds, average="macro", zero_division=0))
+    except: macro_f1 = nan
+    try:
+        pcr = recall_score(
+            y_true, preds, average=None,
+            labels=list(range(num_classes)), zero_division=0,
+        ).astype(float)
+    except:
+        pcr = np.full(num_classes, nan)
+    return dict(
+        kappa=kappa, mcc=mcc, macro_f1=macro_f1,
+        per_class_recall=pcr, ece=compute_ece(y_true, probs, n_bins=n_bins),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Threshold / bias tuning
+# ---------------------------------------------------------------------------
+
+def tune_binary_threshold(y_true: np.ndarray, p1: np.ndarray) -> float:
+    """Pick threshold via Youden's J (sensitivity + specificity − 1).
+
+    Class-distribution independent; falls back to 0.5 if fewer than two
+    classes are present in y_true.
+    """
+    if y_true.size == 0 or len(np.unique(y_true)) < 2:
+        return 0.5
+    fpr, tpr, thresholds = roc_curve(y_true, p1)
+    return float(thresholds[np.argmax(tpr + (1.0 - fpr) - 1.0)])
+
+
+def multiclass_acc_with_bias(
+    y_true: np.ndarray, probs: np.ndarray, bias: np.ndarray
+) -> float:
+    """Balanced accuracy (mean per-class recall) after applying log-space bias."""
+    if y_true.size == 0:
+        return float("nan")
+    logits  = np.log(np.clip(probs, 1e-8, 1.0)) + bias.reshape(1, -1)
+    preds   = np.argmax(logits, axis=1)
+    classes = np.unique(y_true)
+    return float(np.mean([(preds[y_true == c] == c).mean() for c in classes]))
+
+
+def tune_multiclass_bias(
+    y_true: np.ndarray, probs: np.ndarray, *, iters: int = 2
+) -> np.ndarray:
+    """Grid-search per-class log-space bias to maximise balanced accuracy."""
+    if y_true.size == 0 or probs.size == 0:
+        return np.zeros((0,), dtype=float)
+    c    = probs.shape[1]
+    bias = np.zeros((c,), dtype=float)
+    grid = np.linspace(-1.0, 1.0, 41)
+    for _ in range(iters):
+        for k in range(c):
+            best_v   = bias[k]
+            best_acc = multiclass_acc_with_bias(y_true, probs, bias)
+            old      = bias[k]
+            for v in grid:
+                bias[k] = float(v)
+                acc = multiclass_acc_with_bias(y_true, probs, bias)
+                if acc > best_acc or (acc == best_acc and abs(v) < abs(best_v)):
+                    best_acc, best_v = acc, float(v)
+            bias[k] = best_v
+            if np.isnan(best_acc):
+                bias[k] = old
+    return bias

v4/classes/profiles/v4papila.py

@@ -0,0 +1,663 @@
+"""v4papila — self-contained PAPILA data module.
+
+Public contract (v4 orchestrator interface)
+-------------------------------------------
+    bundle = build_data(args: dict) -> PapilaBundle
+
+PapilaBundle exposes:
+    .df                     full preprocessed DataFrame (for split building)
+    .label_col, .patient_col, .feature_dim
+    .id_names               tuple of semantic names for each entity_id slot
+                            e.g. ("patient_id", "eye") — used by orchestrator for logging
+    .matrix                 ClinicalDataView (all eyes)
+    .matrix.od / .matrix.os scoped views (OD or OS only)
+    .image                  ImageDataView (all eyes)
+    .image.od  / .image.os  scoped views
+    .build_shells(df, *, level) -> LoaderShell
+
+DataView interface (consumed by towers' get_sample)
+----------------------------------------------------
+Both views accept positional id slots (*ids) matching entity_id tuple positions.
+Semantic names for each position are in view.id_names.
+
+ClinicalDataView:
+    .feature_dim
+    .id_names               e.g. ("patient_id", "eye")
+    .vectorize_entity(*ids) -> np.ndarray
+    .side_map               -> dict mapping generic keys {"a", "b"} to id_1 values
+    .od, .os                -> scoped ClinicalDataView
+
+ImageDataView:
+    .id_names               e.g. ("patient_id", "eye")
+    .get_image_path(*ids)   -> Path
+    .load_image(*ids)       -> PIL.Image
+    .side_map               -> dict mapping generic keys {"a", "b"} to id_1 values
+    .od, .os                -> scoped ImageDataView
+"""
+from __future__ import annotations
+
+import sys
+from functools import cached_property
+from pathlib import Path
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from v4.classes.loaders.image_loader import CachedImageLoader, call_preprocessor
+from v4.classes.dataset import DataBundle, LoaderShell, ShellEntry
+
+
+# ---------------------------------------------------------------------------
+# Pachymetry → IOP correction (PAPILA Table 3)
+# ---------------------------------------------------------------------------
+
+_PACHY_TABLE: Dict[int, int] = {
+    475: +5, 485: +4, 495: +4, 505: +3, 515: +2,
+    525: +1, 535: +1, 545:  0, 555: -1, 565: -1,
+    575: -2, 585: -3, 595: -4, 605: -4, 615: -5,
+}
+_PACHY_KEYS = np.array(sorted(_PACHY_TABLE.keys()))
+
+
+def _nearest_pachy_key(x: float) -> int:
+    return int(_PACHY_KEYS[int(np.argmin(np.abs(_PACHY_KEYS - float(x))))])
+
+
+def _fit_perkins_converter(
+    frames: List[pd.DataFrame], method: str
+) -> Callable[[float, Optional[float]], float]:
+    combined  = pd.concat(frames, ignore_index=True)
+    paired    = combined.dropna(subset=["Pneumatic", "Perkins"])
+    if len(paired) == 0:
+        raise ValueError("No paired Pneumatic+Perkins rows; cannot fit converter.")
+    pneumatic = paired["Pneumatic"].values.astype(float)
+    perkins   = paired["Perkins"].values.astype(float)
+
+    if method == "ratio":
+        ratio = float((pneumatic / perkins).mean())
+        def _conv(p: float, pachy: Optional[float] = None) -> float:
+            return p * ratio
+        return _conv
+
+    elif method == "ols":
+        from scipy import stats as _stats
+        slope, intercept, *_ = _stats.linregress(perkins, pneumatic)
+        slope, intercept = float(slope), float(intercept)
+        def _conv(p: float, pachy: Optional[float] = None) -> float:
+            return p * slope + intercept
+        return _conv
+
+    elif method == "lad":
+        from scipy import stats as _stats
+        from scipy.optimize import minimize as _minimize
+        slope0, intercept0, *_ = _stats.linregress(perkins, pneumatic)
+        def _lad_loss(params):
+            a, b = params
+            return np.abs(pneumatic - (a * perkins + b)).mean()
+        res = _minimize(_lad_loss, x0=[slope0, intercept0], method="Nelder-Mead")
+        slope, intercept = float(res.x[0]), float(res.x[1])
+        def _conv(p: float, pachy: Optional[float] = None) -> float:
+            return p * slope + intercept
+        return _conv
+
+    elif method == "multi":
+        from numpy.linalg import lstsq as _lstsq
+        pm = combined.dropna(subset=["Pneumatic", "Perkins", "Pachymetry"])
+        if len(pm) == 0:
+            raise ValueError("No Pneumatic+Perkins+Pachymetry rows; cannot fit multi.")
+        pneu  = pm["Pneumatic"].values.astype(float)
+        perk  = pm["Perkins"].values.astype(float)
+        pv    = pm["Pachymetry"].values.astype(float)
+        X     = np.column_stack([perk, pv, np.ones(len(perk))])
+        coeffs, *_ = _lstsq(X, pneu, rcond=None)
+        slope, pachy_coef, intercept = float(coeffs[0]), float(coeffs[1]), float(coeffs[2])
+        fallback = float(pv.mean())
+        def _conv(p: float, pachy: Optional[float] = None) -> float:
+            pval = pachy if (pachy is not None and not np.isnan(pachy)) else fallback
+            return p * slope + pachy_coef * pval + intercept
+        return _conv
+
+    else:
+        raise ValueError(f"Unknown iop_corr_method: {method!r}. Choose ratio/ols/lad/multi.")
+
+
+def _pick_iop(row: pd.Series, converter: Callable) -> float:
+    pneumatic = row.get("Pneumatic", np.nan)
+    if not pd.isna(pneumatic):
+        return float(pneumatic)
+    perkins = row.get("Perkins", np.nan)
+    if pd.isna(perkins):
+        return np.nan
+    pachy = row.get("Pachymetry", np.nan)
+    return converter(float(perkins), None if pd.isna(pachy) else float(pachy))
+
+
+def _correct_iop(raw_iop: float, pachy: float) -> float:
+    if pd.isna(raw_iop):
+        return np.nan
+    if pd.isna(pachy):
+        return float(raw_iop)
+    key = _nearest_pachy_key(float(pachy))
+    return float(raw_iop) + float(_PACHY_TABLE[key])
+
+
+def _apply_iop_and_drop_md(
+    df: pd.DataFrame, converter: Callable, drop_raw: bool = False
+) -> pd.DataFrame:
+    df["IOP_raw"] = df.apply(lambda row: _pick_iop(row, converter), axis=1)
+    pachy = df.get("Pachymetry", pd.Series(np.nan, index=df.index))
+    df["IOP_corr"] = [
+        _correct_iop(r, p) for r, p in zip(df["IOP_raw"].values, pachy.values)
+    ]
+    drop = [c for c in ("Pneumatic", "Perkins", "VF_MD") if c in df.columns]
+    if drop_raw:
+        drop.append("IOP_raw")
+    if drop:
+        df.drop(columns=drop, inplace=True)
+    return df
+
+
+def _canonicalize_eye_column(df: pd.DataFrame) -> None:
+    if "eyeID" in df.columns:
+        src = "eyeID"
+    else:
+        src = next((c for c in df.columns if "eye" in c.lower()), None)
+        if src is None:
+            df["eyeID"] = "OS"
+            return
+
+    def norm(v):
+        if pd.isna(v):
+            return None
+        x = str(v).strip().upper()
+        if x in {"OS", "L", "LEFT", "0"}:  return "OS"
+        if x in {"OD", "R", "RIGHT", "1"}: return "OD"
+        try:
+            num = int(float(x))
+            return "OD" if num % 2 == 1 else "OS"
+        except Exception:
+            return None
+
+    mapped = df[src].map(norm)
+    uniq   = {u for u in mapped.dropna().unique().tolist()}
+    if not uniq.issubset({"OS", "OD"}):
+        raise ValueError(f"eyeID must be binary; found {sorted(uniq)}")
+    df["eyeID"] = mapped.fillna("OS")
+
+
+# ---------------------------------------------------------------------------
+# DataView classes
+# ---------------------------------------------------------------------------
+
+class ClinicalDataView:
+    """Tabular feature view over a (possibly eye-filtered) clinical DataFrame.
+
+    Exposes feature_dim and vectorize_entity so towers can retrieve
+    feature vectors by entity identity without knowing about the DataFrame.
+
+    Scoped views (OD or OS only) are accessed via .od and .os properties.
+
+    id_names gives semantic labels for each positional slot in an entity_id tuple,
+    e.g. ("patient_id", "eye").  The orchestrator uses this for logging without
+    needing to know PAPILA-specific field names itself.
+    """
+
+    # PAPILA canonical side keys used in ShellEntry entity_ids
+    SIDE_A = "OD"
+    SIDE_B = "OS"
+
+    # Semantic name for each entity_id position (id_0, id_1, ...)
+    id_names: tuple[str, ...] = ("patient_id", "eye")
+
+    def __init__(
+        self,
+        df:           pd.DataFrame,
+        patient_col:  str,
+        scalar_cols:  list[str],
+        cat_cols:     list[str],
+        scalar_stats: dict,
+        cat_maps:     dict,
+        *,
+        eye_filter: str | None = None,  # "OD", "OS", or None (all eyes)
+    ):
+        self._df          = df
+        self.patient_col  = patient_col
+        self.scalar_cols  = scalar_cols
+        self.cat_cols     = cat_cols
+        self.scalar_stats = scalar_stats
+        self.cat_maps     = cat_maps
+        self._eye_filter  = eye_filter
+
+        # Build a (patient_id, eyeID) → row index for fast lookup
+        if "eyeID" in df.columns:
+            self._idx = df.set_index([patient_col, "eyeID"])
+        else:
+            self._idx = df.set_index(patient_col)
+
+    # ── Public API ──────────────────────────────────────────────────────────
+
+    @property
+    def feature_dim(self) -> int:
+        n_scalar = len(self.scalar_cols)
+        n_cat    = sum(len(m) for m in self.cat_maps.values())
+        return n_scalar + n_cat + n_scalar  # scalars + one-hots + missing flags
+
+    def vectorize_entity(self, *ids) -> np.ndarray:
+        """Return the feature vector for an entity identified by positional ids.
+
+        Positional slots match entity_id tuple positions (see id_names).
+        For PAPILA: ids = (id_0, id_1) = (patient_id, eye).
+        """
+        try:
+            row = self._idx.loc[ids if len(ids) > 1 else ids[0]].copy()
+        except KeyError:
+            names = self.id_names[:len(ids)]
+            raise KeyError(
+                f"ClinicalDataView: no row found for {dict(zip(names, ids))}"
+            )
+        # set_index removes index-level columns from the row; restore any that
+        # _vectorize_row needs (e.g. eyeID is a cat feature AND an index level)
+        idx_names = (self._idx.index.names
+                     if hasattr(self._idx.index, 'names')
+                     else [self._idx.index.name])
+        for name, val in zip(idx_names, ids if len(ids) > 1 else [ids[0]]):
+            if name not in row.index:
+                row[name] = val
+        return self._vectorize_row(row)
+
+    # ── Scoped views ─────────────────────────────────────────────────────────
+
+    @property
+    def side_map(self) -> dict[str, str]:
+        """Generic side-key → dataset side string. Towers use this for patient-level shells."""
+        return {"a": self.SIDE_A, "b": self.SIDE_B}
+
+    @cached_property
+    def od(self) -> "ClinicalDataView":
+        return self._scoped(self.SIDE_A)
+
+    @cached_property
+    def os(self) -> "ClinicalDataView":
+        return self._scoped(self.SIDE_B)
+
+    def _scoped(self, eye: str) -> "ClinicalDataView":
+        sub = self._df[self._df["eyeID"] == eye].reset_index(drop=True)
+        return ClinicalDataView(
+            df=sub,
+            patient_col=self.patient_col,
+            scalar_cols=self.scalar_cols,
+            cat_cols=self.cat_cols,
+            scalar_stats=self.scalar_stats,
+            cat_maps=self.cat_maps,
+            eye_filter=eye,
+        )
+
+    # ── Internal ─────────────────────────────────────────────────────────────
+
+    def _vectorize_row(self, row: pd.Series) -> np.ndarray:
+        feats: list[float] = []
+        miss:  list[float] = []
+        for col in self.scalar_cols:
+            v = pd.to_numeric(row.get(col), errors="coerce")
+            if pd.isna(v):
+                miss.append(1.0)
+                v = self.scalar_stats[col]["median"]
+            else:
+                miss.append(0.0)
+            lo = self.scalar_stats[col]["min"]
+            hi = self.scalar_stats[col]["max"]
+            feats.append((float(v) - lo) / (hi - lo) if hi > lo else 0.0)
+        for col in self.cat_cols:
+            mapping = self.cat_maps[col]
+            one = [0.0] * len(mapping)
+            key = row.get(col)
+            one[mapping.get(key, 0)] = 1.0
+            feats.extend(one)
+        feats.extend(miss)
+        return np.asarray(feats, dtype=np.float32)
+
+
+class ImageDataView:
+    """Image path and loading view over a (possibly eye-filtered) DataFrame.
+
+    Provides get_image_path and load_image keyed by positional id slots.
+    Scoped views (.od, .os) are available for single-side towers.
+    The optional image_cache is a shared CachedImageLoader for the run.
+
+    id_names gives semantic labels for each positional slot in an entity_id tuple,
+    e.g. ("patient_id", "eye").
+    """
+
+    SIDE_A = "OD"
+    SIDE_B = "OS"
+
+    id_names: tuple[str, ...] = ("patient_id", "eye")
+
+    def __init__(
+        self,
+        df:                pd.DataFrame,
+        patient_col:       str,
+        image_dir:         str,
+        filename_template: str,
+        preprocessor:      Callable | None = None,
+        image_cache:       CachedImageLoader | None = None,
+        *,
+        eye_filter: str | None = None,
+    ):
+        self._df               = df
+        self.patient_col       = patient_col
+        self.image_dir         = Path(image_dir)
+        self.filename_template = filename_template
+        self.preprocessor      = preprocessor
+        self.image_cache       = image_cache
+        self._eye_filter       = eye_filter
+
+        if "eyeID" in df.columns:
+            self._idx = df.set_index([patient_col, "eyeID"])
+        else:
+            self._idx = df.set_index(patient_col)
+
+    # ── Public API ──────────────────────────────────────────────────────────
+
+    def get_image_path(self, *ids) -> Path:
+        """Return the image path for an entity identified by positional ids.
+
+        For PAPILA: ids = (id_0, id_1) = (patient_id, eye).
+        """
+        id_0, id_1 = ids  # PAPILA always uses two slots
+        try:
+            row = self._idx.loc[(id_0, id_1)]
+        except KeyError:
+            names = self.id_names[:len(ids)]
+            raise KeyError(
+                f"ImageDataView: no row for {dict(zip(names, ids))}"
+            )
+        pid = int(row[self.patient_col]) if self.patient_col in row.index else int(id_0)
+        return self.image_dir / self.filename_template.format(pid=pid, eye=id_1)
+
+    def load_image(self, *ids) -> Image.Image:
+        """Load image for an entity identified by positional ids."""
+        path = self.get_image_path(*ids)
+        if self.image_cache is not None:
+            return self.image_cache.load(path, preprocessor=self.preprocessor)
+        img = Image.open(path).convert("RGB")
+        if self.preprocessor is not None:
+            img = call_preprocessor(self.preprocessor, img, path)
+        return img
+
+    # ── Scoped views ─────────────────────────────────────────────────────────
+
+    @cached_property
+    def od(self) -> "ImageDataView":
+        return self._scoped(self.SIDE_A)
+
+    @cached_property
+    def os(self) -> "ImageDataView":
+        return self._scoped(self.SIDE_B)
+
+    @property
+    def side_map(self) -> dict[str, str]:
+        """Generic side-key → dataset side string. Towers use this for patient-level shells."""
+        return {"a": self.SIDE_A, "b": self.SIDE_B}
+
+    def _scoped(self, eye: str) -> "ImageDataView":
+        sub = self._df[self._df["eyeID"] == eye].reset_index(drop=True)
+        return ImageDataView(
+            df=sub,
+            patient_col=self.patient_col,
+            image_dir=str(self.image_dir),
+            filename_template=self.filename_template,
+            preprocessor=self.preprocessor,
+            image_cache=self.image_cache,
+            eye_filter=eye,
+        )
+
+
+# ---------------------------------------------------------------------------
+# PapilaBundle — the v4 DataBundle returned by build_data
+# ---------------------------------------------------------------------------
+
+class PapilaBundle:
+    """V4 DataBundle for PAPILA.
+
+    Wraps the v3 DataBundle for backward compatibility (df, feature_dim,
+    vectorize_row, get_image_path, patient_col, label_col) while adding
+    the v4 DataView interface and build_shells().
+    """
+
+    def __init__(
+        self,
+        bundle:       DataBundle,
+        image_dir:    str,
+        preprocessor: Callable | None = None,
+        image_cache:  CachedImageLoader | None = None,
+    ):
+        self._bundle    = bundle
+        self._image_dir = image_dir
+
+        # ── ClinicalDataView (all eyes) ──────────────────────────────────────
+        self.matrix = ClinicalDataView(
+            df=bundle.df,
+            patient_col=bundle.patient_col,
+            scalar_cols=bundle.scalar_cols,
+            cat_cols=bundle.cat_cols,
+            scalar_stats=bundle.scalar_stats,
+            cat_maps=bundle.cat_maps,
+        )
+
+        # ── ImageDataView (all eyes) ─────────────────────────────────────────
+        self.image = ImageDataView(
+            df=bundle.df,
+            patient_col=bundle.patient_col,
+            image_dir=image_dir,
+            filename_template=bundle.filename_template,
+            preprocessor=preprocessor,
+            image_cache=image_cache,
+        )
+
+    # ── Entity-id metadata (for orchestrator logging) ────────────────────────
+
+    @property
+    def id_names(self) -> tuple[str, ...]:
+        """Semantic names for each entity_id position, e.g. ('patient_id', 'eye').
+
+        Orchestrators use this to decode entity_ids for logging without
+        hardcoding dataset-specific field names.
+        """
+        return self.matrix.id_names  # both views share the same structure
+
+    # ── Identity column registry (for orchestrator split_identity_level) ────────
+
+    @property
+    def identity_cols(self) -> list[str]:
+        """Ordered list of grouping columns, one per identity level.
+
+        identity_level=1 → identity_cols[0] → patient column (group by patient)
+        identity_level=2 → identity_cols[1] → eye column (group by patient+eye)
+        """
+        return [self._bundle.patient_col, "eyeID"]
+
+    # ── Backward-compat delegates ────────────────────────────────────────────
+
+    @property
+    def df(self) -> pd.DataFrame:
+        return self._bundle.df
+
+    @property
+    def label_col(self) -> str:
+        return self._bundle.label_col
+
+    @property
+    def patient_col(self) -> str:
+        return self._bundle.patient_col
+
+    @property
+    def feature_dim(self) -> int:
+        return self._bundle.feature_dim
+
+    def vectorize_row(self, row: pd.Series) -> np.ndarray:
+        return self._bundle.vectorize_row(row)
+
+    def get_image_path(self, row: pd.Series):
+        return self._bundle.get_image_path(row)
+
+    # ── Shell building ───────────────────────────────────────────────────────
+
+    def build_shells(
+        self,
+        df: pd.DataFrame,
+        *,
+        level: str = "eye",
+        label_filter: list[int] | None = None,
+    ) -> LoaderShell:
+        """Build a LoaderShell from a split DataFrame.
+
+        level="eye"     — one ShellEntry per eye row.
+                          entity_id = (patient_id, side_key)
+                          e.g. (42, "OD") or (42, "OS")
+
+        level="patient" — one ShellEntry per patient.
+                          entity_id = (patient_id,)  — 1-tuple
+                          Towers that need both sides use data_view.side_map
+                          to assemble them in get_sample.
+                          Patients missing either eye are excluded.
+        """
+        pc = self.patient_col
+        lc = self.label_col
+
+        if label_filter is not None:
+            df = df[df[lc].isin(label_filter)]
+
+        entries: list[ShellEntry] = []
+
+        if level == "eye":
+            for _, row in df.iterrows():
+                pid   = int(row[pc])
+                label = int(row[lc])
+                side  = str(row.get("eyeID", "OD"))
+                entries.append(ShellEntry(entity_id=(pid, side), label=label))
+
+        elif level == "patient":
+            for pid, grp in df.groupby(pc):
+                if "eyeID" in grp.columns:
+                    eyes = set(grp["eyeID"].unique())
+                    if "OD" not in eyes or "OS" not in eyes:
+                        continue
+                label_mode = grp[lc].mode()
+                label = int(label_mode.iloc[0]) if not label_mode.empty else int(grp[lc].iloc[0])
+                entries.append(ShellEntry(entity_id=(int(pid),), label=label))
+
+        else:
+            raise ValueError(f"Unknown shell level: {level!r}. Choose 'eye' or 'patient'.")
+
+        return LoaderShell(entries=entries)
+
+
+# ---------------------------------------------------------------------------
+# Resolve helper — used by the orchestrator to inject DataView into towers
+# ---------------------------------------------------------------------------
+
+def resolve_data_source(bundle: PapilaBundle, path: str):
+    """Resolve a dot-path data source string against a PapilaBundle.
+
+    Examples
+    --------
+    "matrix"     → bundle.matrix
+    "matrix.od"  → bundle.matrix.od
+    "image"      → bundle.image
+    "image.os"   → bundle.image.os
+    """
+    obj = bundle
+    for part in path.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+# ---------------------------------------------------------------------------
+# Public contract: build_data(args: dict) -> PapilaBundle
+# ---------------------------------------------------------------------------
+
+_DEFAULT_CAT_COLS = ["Gender", "Phakic/Pseudophakic"]
+
+
+def build_data(args: dict) -> PapilaBundle:
+    """Build and return a PapilaBundle.
+
+    args keys
+    ---------
+    image_dir        (required)
+    clinical_dir     (required)
+    label_col        (default: "Diagnosis")
+    iop_corr_method  (default: "ratio")
+    iop_drop_raw     (default: False)
+    exclude_cols     (default: [])
+    cat_cols         (default: ["Gender", "Phakic/Pseudophakic"])
+    n_splits         (default: 5)
+    random_seed      (default: 42)
+    in_memory_cache  (default: False) — enable shared image cache for the run
+    """
+    image_dir      = args["image_dir"]
+    clinical_dir   = args["clinical_dir"]
+    label_col      = args.get("label_col",       "Diagnosis")
+    iop_method     = args.get("iop_corr_method", "ratio")
+    iop_drop_raw   = bool(args.get("iop_drop_raw", False))
+    exclude_cols   = list(args.get("exclude_cols", []))
+    cat_cols       = list(args.get("cat_cols",    _DEFAULT_CAT_COLS))
+    n_splits       = int(args.get("n_splits",     5))
+    random_seed    = int(args.get("random_seed",  42))
+    use_cache      = bool(args.get("in_memory_cache", False))
+
+    effective_cat = [c for c in cat_cols if c not in exclude_cols]
+
+    bundle = DataBundle(
+        image_dir=image_dir,
+        clinical_dir=clinical_dir,
+        label_col=label_col,
+        patient_col="Patient ID",
+        cat_cols=effective_cat,
+        n_splits=n_splits,
+        random_seed=random_seed,
+        filename_template="RET{pid:03d}{eye}.jpg",
+    )
+
+    od  = pd.read_excel(f"{clinical_dir}/patient_data_od.xlsx", header=1)
+    od["eyeID"] = "OD"
+    os_ = pd.read_excel(f"{clinical_dir}/patient_data_os.xlsx", header=1)
+    os_["eyeID"] = "OS"
+
+    for frame in (od, os_):
+        if "Patient ID" not in frame.columns and "ID" in frame.columns:
+            frame.rename(columns={"ID": "Patient ID"}, inplace=True)
+        frame["Patient ID"] = (
+            frame["Patient ID"].astype(str).str.extract(r"(\d+)")[0].astype(int)
+        )
+        _canonicalize_eye_column(frame)
+
+    bundle.add_df(od,  id_column="ID", exclude_cols=exclude_cols or None)
+    bundle.add_df(os_, id_column="ID", exclude_cols=exclude_cols or None)
+
+    converter = _fit_perkins_converter(bundle.frames, method=iop_method)
+    for i in range(len(bundle.frames)):
+        bundle.frames[i] = _apply_iop_and_drop_md(
+            bundle.frames[i], converter=converter, drop_raw=iop_drop_raw
+        )
+
+    bundle._refresh_master_df(exclude_cols=exclude_cols or None)
+    bundle._infer_or_validate_feature_types(exclude_cols=exclude_cols or None)
+    bundle._compute_numeric_stats()
+    bundle._build_cat_maps()
+    bundle._compute_feature_dim()
+
+    image_cache = CachedImageLoader() if use_cache else None
+
+    return PapilaBundle(
+        bundle=bundle,
+        image_dir=image_dir,
+        image_cache=image_cache,
+    )

v4/classes/split_manager.py

@@ -0,0 +1,177 @@
+"""split_manager — generic stratified k-fold splitter.
+
+SplitManager splits a DataFrame into train/val/test folds using an
+outer/inner k-fold scheme.  The grouping identity is controlled by
+``group_col``:
+
+    group_col=None       — row-level splits (each row is its own identity)
+    group_col="Patient ID" — group-level splits (all rows sharing a group
+                             key land in the same fold)
+
+The translation from a conceptual "identity level" to a concrete column
+name belongs in the caller (typically the orchestrator), which has access
+to the data bundle's column schema.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Iterable, Optional
+
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import KFold, StratifiedKFold
+
+
+# ---------------------------------------------------------------------------
+# Data structures
+# ---------------------------------------------------------------------------
+
+@dataclass(frozen=True)
+class SplitPlan:
+    train_ids: set[Any]
+    val_ids:   set[Any]
+    test_ids:  set[Any]
+
+
+@dataclass
+class Split:
+    """One fold's train/val/test DataFrames."""
+    train: pd.DataFrame
+    val:   pd.DataFrame
+    test:  Optional[pd.DataFrame] = None
+
+
+# ---------------------------------------------------------------------------
+# Core splitter
+# ---------------------------------------------------------------------------
+
+def _can_stratify(labels: np.ndarray, n_splits: int) -> bool:
+    if labels.size == 0:
+        return False
+    unique, counts = np.unique(labels, return_counts=True)
+    return len(unique) >= 2 and bool(np.all(counts >= n_splits))
+
+
+def _build_split_plans(
+    ids:      np.ndarray,
+    labels:   np.ndarray,
+    n_splits: int,
+    seed:     int,
+) -> list[SplitPlan]:
+    """Outer/inner k-fold: test=fold k, val=fold (k+1)%n, train=remaining."""
+    if ids.size == 0:
+        raise ValueError("No samples available for splitting")
+    if len(set(ids.tolist())) != ids.size:
+        raise ValueError("ids must be unique")
+    if n_splits < 3:
+        raise ValueError("n_splits must be >= 3 for outer/inner k-fold")
+
+    if _can_stratify(labels, n_splits):
+        splitter = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
+        folds = list(splitter.split(ids, labels))
+    else:
+        splitter = KFold(n_splits=n_splits, shuffle=True, random_state=seed)
+        folds = list(splitter.split(ids))
+
+    fold_sets = [set(ids[test_idx].tolist()) for _, test_idx in folds]
+
+    plans = []
+    for k in range(n_splits):
+        test_ids  = fold_sets[k]
+        val_ids   = fold_sets[(k + 1) % n_splits]
+        train_ids = set().union(*(fold_sets[j] for j in range(n_splits)
+                                  if j != k and j != (k + 1) % n_splits))
+        plans.append(SplitPlan(train_ids=train_ids, val_ids=val_ids, test_ids=test_ids))
+    return plans
+
+
+# ---------------------------------------------------------------------------
+# SplitManager
+# ---------------------------------------------------------------------------
+
+class SplitManager:
+    """Generic stratified k-fold split manager.
+
+    Parameters
+    ----------
+    group_col : str | None
+        Column whose values define the grouping identity for fold assignment.
+        ``None`` splits on individual rows (no grouping).
+    label_col : str | None
+        Column used for stratification.  Resolved from the DataFrame at
+        ``build_plans`` time if not provided here.
+    """
+
+    def __init__(
+        self,
+        group_col:  Optional[str] = None,
+        label_col:  Optional[str] = None,
+    ) -> None:
+        self.group_col = group_col
+        self.label_col = label_col
+
+    def build_plans(
+        self,
+        df:       pd.DataFrame,
+        *,
+        n_splits: int  = 5,
+        seed:     int  = 42,
+        label_col: Optional[str] = None,
+    ) -> list[Split]:
+        """Build n_splits fold plans from df.
+
+        Parameters
+        ----------
+        df        : full DataFrame (pre-filtered to the desired eval mode)
+        n_splits  : number of folds (must be >= 3)
+        seed      : random seed for reproducibility
+        label_col : override for stratification column (falls back to
+                    ``self.label_col``, then raises)
+        """
+        lc = label_col or self.label_col
+        if lc is None:
+            raise ValueError("label_col must be provided to build_plans or SplitManager")
+        if lc not in df.columns:
+            raise ValueError(f"label_col {lc!r} not found in DataFrame")
+
+        df = df.copy().reset_index(drop=True)
+
+        if self.group_col is None:
+            # Row-level: each row is its own identity
+            ids    = df.index.to_numpy()
+            labels = df[lc].to_numpy()
+            plans  = _build_split_plans(ids, labels, n_splits, seed)
+            return [
+                Split(
+                    train=df[df.index.isin(p.train_ids)].reset_index(drop=True),
+                    val  =df[df.index.isin(p.val_ids)  ].reset_index(drop=True),
+                    test =df[df.index.isin(p.test_ids)  ].reset_index(drop=True),
+                )
+                for p in plans
+            ]
+        else:
+            gc = self.group_col
+            if gc not in df.columns:
+                raise ValueError(f"group_col {gc!r} not found in DataFrame")
+            # Group-level: collapse to one row per group, then split
+            group_table = (
+                df.groupby(gc, as_index=False)[lc]
+                .agg(lambda s: s.mode().iloc[0] if not s.mode().empty else s.iloc[0])
+                .rename(columns={lc: "_label"})
+                .sort_values(gc)
+                .reset_index(drop=True)
+            )
+            plans = _build_split_plans(
+                group_table[gc].to_numpy(),
+                group_table["_label"].to_numpy(),
+                n_splits,
+                seed,
+            )
+            return [
+                Split(
+                    train=df[df[gc].isin(p.train_ids)].reset_index(drop=True),
+                    val  =df[df[gc].isin(p.val_ids)  ].reset_index(drop=True),
+                    test =df[df[gc].isin(p.test_ids)  ].reset_index(drop=True),
+                )
+                for p in plans
+            ]

v4/classes/stages/fusion.py

@@ -0,0 +1,273 @@
+"""stages/fusion — fusion stage runner: trains a bridge + associated head stages."""
+from __future__ import annotations
+
+import importlib
+from random import choice, random as _random
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from v4.classes.dataset import LoaderShell, to_label_tensor
+from v4.classes.metrics import score_arrays, compute_extended_metrics, tune_binary_threshold
+from v4.classes.stages.helpers import (
+    encode_embedding, get_out_dim, resolve_input_dims, phase_for_epoch,
+)
+
+
+def collect_probs(
+    bridge,
+    primary_head,
+    stage_cfg:    dict,
+    towers:       dict,
+    stage_models: dict,
+    cfg_stages:   list[dict],
+    loader,
+    device,
+    num_classes:  int,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Eval pass for one fusion stage; returns (y_true, softmax_probs)."""
+    from v4.classes.dataset import to_label_tensor
+    bridge.eval(); primary_head.eval()
+    for t in towers.values():
+        t.eval()
+    inputs       = stage_cfg["inputs"]
+    is_bilateral = isinstance(inputs, dict)
+    y_all, p_all = [], []
+
+    with torch.no_grad():
+        for batch in loader:
+            y = batch.get("label")
+            if not torch.is_tensor(y):
+                continue
+
+            if is_bilateral:
+                side_embs = {
+                    side: encode_embedding(src, batch, side, towers, stage_models, cfg_stages, device)
+                    for side, src in inputs.items()
+                }
+                z = bridge(side_embs)
+            else:
+                embs = [
+                    encode_embedding(n, batch, None, towers, stage_models, cfg_stages, device)
+                    for n in inputs
+                ]
+                z = bridge(embs)
+
+            logits = primary_head(z)
+            y_all.append(to_label_tensor(y, device).cpu().numpy())
+            p_all.append(F.softmax(logits, dim=1).cpu().numpy())
+
+    if not y_all:
+        return np.zeros(0, dtype=np.int64), np.zeros((0, num_classes), dtype=np.float32)
+    return np.concatenate(y_all), np.concatenate(p_all, axis=0)
+
+
+def run(
+    stage_cfg:    dict,
+    cfg:          dict,
+    towers:       dict,
+    stage_models: dict,
+    data,
+    split,
+    label_filter,
+    num_classes:  int,
+    device,
+    fold:         int,
+    cfg_stages:   list[dict],
+    _make_loader,
+) -> tuple[dict, dict]:
+    """Train one fusion stage + its associated head stages.
+
+    Returns (updated_stage_models, metrics_dict).
+    """
+    nan          = float("nan")
+    name         = stage_cfg["name"]
+    level        = stage_cfg["level"]
+    epochs       = stage_cfg["epochs"]
+    inputs       = stage_cfg["inputs"]
+    is_bilateral = isinstance(inputs, dict)
+
+    s_train = data.build_shells(split.train, level=level, label_filter=label_filter)
+    s_val   = data.build_shells(split.val,   level=level, label_filter=label_filter)
+    s_test  = (data.build_shells(split.test, level=level, label_filter=label_filter)
+               if split.test is not None else LoaderShell(entries=[]))
+
+    if not s_val.entries:
+        print(f"  fold{fold+1}: no val samples for stage {name!r}, skipping.", flush=True)
+        return stage_models, {}
+
+    bs           = cfg["training"]["batch_size"]
+    train_loader = _make_loader(s_train, towers, batch_size=bs, shuffle=True)
+    val_loader   = _make_loader(s_val,   towers, batch_size=bs, shuffle=False)
+    test_loader  = (_make_loader(s_test, towers, batch_size=bs, shuffle=False)
+                    if s_test.entries else None)
+
+    # ── Bridge ───────────────────────────────────────────────────────────────
+    input_dims = resolve_input_dims(inputs, towers, stage_models)
+    bmod   = importlib.import_module(stage_cfg["module"])
+    bridge = getattr(bmod, stage_cfg["class"])(input_dims, **stage_cfg.get("args", {})).to(device)
+
+    # ── Head stages ──────────────────────────────────────────────────────────
+    head_stage_cfgs = [s for s in cfg_stages if s["type"] == "head"
+                       and s.get("train_with") == name]
+    head_models: dict[str, torch.nn.Module] = {}
+    for hs in head_stage_cfgs:
+        h_dim = get_out_dim(hs["input"], towers, {**stage_models, name: bridge})
+        h_mod = importlib.import_module(hs.get("module", "v4.classes.heads.classifier"))
+        h_cls = getattr(h_mod, hs.get("class", "ClassificationHead"))
+        head_models[hs["name"]] = h_cls(h_dim, num_classes).to(device)
+
+    primary_hs_cfg = next((hs for hs in head_stage_cfgs if not hs.get("bcd", False)), None)
+    bcd_head_cfgs  = [hs for hs in head_stage_cfgs if hs.get("bcd", False)]
+
+    if primary_hs_cfg is None:
+        print(f"  WARNING: no primary head for stage {name!r}; skipping.", flush=True)
+        return stage_models, {}
+
+    primary_head = head_models[primary_hs_cfg["name"]]
+
+    # ── Freeze prior stages ───────────────────────────────────────────────────
+    for m in stage_models.values():
+        for p in m.parameters():
+            p.requires_grad_(False)
+        m.eval()
+
+    # ── Optimizer ────────────────────────────────────────────────────────────
+    train_towers = stage_cfg.get("train_towers", False)
+    opt_params = (
+        ([p for t in towers.values() for p in t.parameters()] if train_towers else []) +
+        list(bridge.parameters()) +
+        [p for h in head_models.values() for p in h.parameters()]
+    )
+    opt = torch.optim.Adam(opt_params, lr=cfg["training"]["lr"])
+
+    warmup_cfg = stage_cfg.get("warmup", {})
+    wt         = 0 if is_bilateral else warmup_cfg.get("tower_epochs", 0)
+    wf         = 0 if is_bilateral else warmup_cfg.get("fused_epochs",  0)
+    bcd_prob   = cfg["training"].get("bcd_prob", 0.5)
+
+    # ── Epoch loop ────────────────────────────────────────────────────────────
+    for epoch in range(epochs):
+        bridge.train()
+        for h in head_models.values():
+            h.train()
+        for t in towers.values():
+            if train_towers:
+                t.train()
+            else:
+                t.eval()
+
+        if not is_bilateral:
+            phase = phase_for_epoch(epoch, wt, wf)
+            if hasattr(bridge, "set_phase"):
+                bridge.set_phase(phase)
+            for t in towers.values():
+                if hasattr(t, "set_phase"):
+                    t.set_phase(phase)
+        else:
+            phase = "fusion"
+
+        total_loss = total_correct = total_n = 0
+
+        for batch in train_loader:
+            y = batch.get("label")
+            if not torch.is_tensor(y):
+                continue
+            y_t = to_label_tensor(y, device)
+            if y_t.numel() == 0:
+                continue
+
+            if is_bilateral:
+                side_embs  = {
+                    side: encode_embedding(src, batch, side, towers, stage_models, cfg_stages, device)
+                    for side, src in inputs.items()
+                }
+                local_embs = {name: bridge(side_embs)}
+            else:
+                local_embs = {n: towers[n](batch[n].to(device)) for n in inputs
+                              if n in batch and torch.is_tensor(batch[n])}
+                if len(local_embs) != len(inputs):
+                    continue
+                local_embs[name] = bridge(list(local_embs[n] for n in inputs))
+
+            head_logits = {
+                hs["name"]: head_models[hs["name"]](local_embs[hs["input"]])
+                for hs in head_stage_cfgs
+                if hs["input"] in local_embs
+            }
+
+            if is_bilateral or phase == "fused_warmup":
+                logits = head_logits.get(primary_hs_cfg["name"])
+            elif phase == "tower_warmup" and bcd_head_cfgs:
+                losses = [F.cross_entropy(head_logits[hs["name"]], y_t)
+                          for hs in bcd_head_cfgs if hs["name"] in head_logits]
+                if not losses:
+                    continue
+                loss = sum(losses) / len(losses)
+                opt.zero_grad(); loss.backward(); opt.step()
+                total_loss += loss.item() * len(y_t)
+                total_n    += len(y_t)
+                continue
+            else:
+                if bcd_head_cfgs and _random() < bcd_prob:
+                    logits = head_logits.get(choice(bcd_head_cfgs)["name"])
+                else:
+                    logits = head_logits.get(primary_hs_cfg["name"])
+
+            if logits is None:
+                continue
+            loss = F.cross_entropy(logits, y_t)
+            opt.zero_grad(); loss.backward(); opt.step()
+            total_correct += int((logits.argmax(1) == y_t).sum())
+            total_loss    += loss.item() * len(y_t)
+            total_n       += len(y_t)
+
+        tr_loss = total_loss / total_n if total_n else nan
+        tr_acc  = total_correct / total_n if total_n else nan
+
+        y_v, p_v = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                  stage_models, cfg_stages, val_loader, device, num_classes)
+        _, val_auc, _ = score_arrays(y_v, p_v, num_classes) if y_v.size else (nan, nan, nan)
+        print(
+            f"  fold{fold+1} [{name}] ep{epoch+1:03d}/{epochs} [{phase:14s}]"
+            f"  loss={tr_loss:.4f}  acc={tr_acc:.3f}  val_auc={val_auc:.4f}",
+            flush=True,
+        )
+
+    # ── Final eval ────────────────────────────────────────────────────────────
+    y_val, p_val = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                  stage_models, cfg_stages, val_loader, device, num_classes)
+    val_acc, val_auc, val_n = (score_arrays(y_val, p_val, num_classes)
+                               if y_val.size else (nan, nan, nan))
+    ext = compute_extended_metrics(y_val, p_val, num_classes) if y_val.size else {}
+
+    val_threshold = 0.5
+    if (cfg["training"].get("tune_binary_threshold")
+            and num_classes == 2 and y_val.size >= 2):
+        val_threshold = tune_binary_threshold(y_val, p_val[:, 1])
+
+    test_auc = test_acc = test_n = nan
+    if test_loader is not None:
+        y_te, p_te = collect_probs(bridge, primary_head, stage_cfg, towers,
+                                    stage_models, cfg_stages, test_loader, device, num_classes)
+        test_acc, test_auc, test_n = (score_arrays(y_te, p_te, num_classes)
+                                      if y_te.size else (nan, nan, nan))
+
+    updated = dict(stage_models)
+    updated[name] = bridge
+    updated.update(head_models)
+
+    metrics = {
+        f"{name}_val_auc":       val_auc,
+        f"{name}_val_acc":       val_acc,
+        f"{name}_val_n":         val_n,
+        f"{name}_val_kappa":     ext.get("kappa", nan),
+        f"{name}_val_mcc":       ext.get("mcc", nan),
+        f"{name}_val_f1":        ext.get("macro_f1", nan),
+        f"{name}_val_threshold": val_threshold,
+        f"{name}_test_auc":      test_auc,
+        f"{name}_test_acc":      test_acc,
+        f"{name}_test_n":        test_n,
+    }
+    return updated, metrics

v4/classes/stages/helpers.py

@@ -0,0 +1,66 @@
+"""stages/helpers — shared utilities for stage runners."""
+from __future__ import annotations
+
+import torch
+
+
+def get_out_dim(name: str, towers: dict, stage_models: dict) -> int:
+    if name in towers:
+        return towers[name].out_dim
+    if name in stage_models:
+        return stage_models[name].out_dim
+    raise KeyError(f"No out_dim for input {name!r}")
+
+
+def resolve_input_dims(inputs, towers: dict, stage_models: dict):
+    """Return list[int] for list inputs, dict[str, int] for dict inputs."""
+    if isinstance(inputs, list):
+        return [get_out_dim(n, towers, stage_models) for n in inputs]
+    return {k: get_out_dim(v, towers, stage_models) for k, v in inputs.items()}
+
+
+def encode_embedding(
+    name:         str,
+    batch:        dict,
+    side:         str | None,
+    towers:       dict,
+    stage_models: dict,
+    cfg_stages:   list[dict],
+    device,
+) -> torch.Tensor:
+    """Return the embedding for a named tower or prior frozen fusion stage.
+
+    For bilateral batches, *side* selects which side dict entry to use.
+    Fusion stages are re-encoded recursively from their own inputs.
+    Runs under torch.no_grad() — only used for frozen passes.
+    """
+    if name in towers:
+        t = batch[name]
+        if side is not None and isinstance(t, dict):
+            t = t[side]
+        with torch.no_grad():
+            return towers[name](t.to(device))
+
+    s_cfg  = next(s for s in cfg_stages if s["name"] == name)
+    model  = stage_models[name]
+    inputs = s_cfg["inputs"]
+
+    if isinstance(inputs, list):
+        embeddings = [
+            encode_embedding(n, batch, side, towers, stage_models, cfg_stages, device)
+            for n in inputs
+        ]
+        with torch.no_grad():
+            return model(embeddings)
+
+    raise NotImplementedError(
+        f"Stage {name!r} uses dict inputs and cannot be used as a fusion input."
+    )
+
+
+def phase_for_epoch(epoch: int, warmup_tower: int, warmup_fused: int) -> str:
+    if epoch < warmup_tower:
+        return "tower_warmup"
+    if epoch < warmup_tower + warmup_fused:
+        return "fused_warmup"
+    return "main"

v4/classes/stages/warm.py

@@ -0,0 +1,73 @@
+"""stages/warm — warm stage runner: pre-trains a single tower with a temporary probe."""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from v4.classes.dataset import to_label_tensor
+from v4.classes.stages.helpers import phase_for_epoch
+
+
+def run(
+    stage_cfg:   dict,
+    towers:      dict,
+    data,
+    split,
+    label_filter,
+    cfg:         dict,
+    num_classes: int,
+    device,
+    fold:        int,
+    _make_loader,
+    _balanced_sampler,
+) -> None:
+    """Pre-train one tower using a temporary linear probe (probe discarded after)."""
+    tower_name = stage_cfg["tower"]
+    n_epochs   = stage_cfg.get("epochs", 0)
+    level      = stage_cfg["level"]
+
+    if n_epochs == 0:
+        return
+
+    s_train = data.build_shells(split.train, level=level, label_filter=label_filter)
+    bs      = cfg["training"]["batch_size"]
+    loader  = _make_loader(
+        s_train, {tower_name: towers[tower_name]},
+        batch_size=bs, shuffle=False,
+        sampler=_balanced_sampler(s_train),
+    )
+
+    for n, t in towers.items():
+        for p in t.parameters():
+            p.requires_grad_(n == tower_name)
+
+    probe = torch.nn.Linear(towers[tower_name].out_dim, num_classes).to(device)
+    opt   = torch.optim.Adam(
+        list(towers[tower_name].parameters()) + list(probe.parameters()),
+        lr=cfg["training"]["lr"],
+    )
+
+    towers[tower_name].train()
+    for epoch in range(n_epochs):
+        total_loss = total_correct = total_n = 0
+        for batch in loader:
+            y = batch.get("label")
+            x = batch.get(tower_name)
+            if not torch.is_tensor(y) or not torch.is_tensor(x):
+                continue
+            y_t    = to_label_tensor(y, device)
+            logits = probe(towers[tower_name](x.to(device)))
+            loss   = F.cross_entropy(logits, y_t)
+            opt.zero_grad(); loss.backward(); opt.step()
+            total_loss    += loss.item() * len(y_t)
+            total_correct += int((logits.argmax(1) == y_t).sum())
+            total_n       += len(y_t)
+        print(
+            f"  fold{fold+1} [warm/{tower_name}] ep{epoch+1:03d}/{n_epochs}"
+            f"  loss={total_loss/total_n:.4f}  acc={total_correct/total_n:.3f}",
+            flush=True,
+        )
+
+    for t in towers.values():
+        for p in t.parameters():
+            p.requires_grad_(True)

v4/classes/towerbase.py

@@ -0,0 +1,57 @@
+"""towerbase — v4 TowerBase ABC.
+
+All v4 towers inherit from TowerBase.  The only required interface is:
+
+    out_dim    : int property — embedding dimensionality
+    _side_map  : dict property — generic side keys {"a","b"} → dataset-specific ids
+    _get(*ids) : retrieve and transform one sample by entity_id slots
+
+get_sample handles the eye-level / patient-level dispatch automatically:
+  len(entity_id) > 1  — full key; calls _get(*entity_id) directly
+  len(entity_id) == 1 — patient key; builds {"a": _get(...), "b": _get(...)}
+                        using _side_map to expand the missing slot
+
+Towers may optionally implement early_pass(context) for cross-tower
+communication before loaders are built.
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import torch
+from torch import nn
+
+from v4.classes.dataset import ShellEntry
+
+
+class TowerBase(nn.Module, ABC):
+
+    @property
+    @abstractmethod
+    def out_dim(self) -> int: ...
+
+    @property
+    @abstractmethod
+    def _side_map(self) -> dict[str, str]: ...
+
+    @abstractmethod
+    def _get(self, *ids) -> torch.Tensor: ...
+
+    def get_sample(self, entry: ShellEntry) -> torch.Tensor | dict[str, torch.Tensor]:
+        eid = entry.entity_id
+        if len(eid) > 1:
+            return self._get(*eid)
+        return {key: self._get(eid[0], id_1) for key, id_1 in self._side_map.items()}
+
+    def set_phase(self, phase: str) -> None:
+        """Called by the orchestrator at the start of each training epoch.
+
+        Default: freeze all parameters during fused_warmup, train otherwise.
+        Override to implement tower-specific phase behaviour.
+        """
+        trainable = phase != "fused_warmup"
+        for p in self.parameters():
+            p.requires_grad_(trainable)
+
+    def early_pass(self, context) -> None:
+        pass

v4/classes/towers/clinical_tower.py

@@ -0,0 +1,94 @@
+"""clinical_tower — ClinicalEncoder for v4.
+
+Self-contained: no v3 dependencies.
+Inherits get_sample dispatch from TowerBase.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch import nn
+
+from v4.classes.towerbase import TowerBase
+from v4.classes.accessory.se_block import SEBlock
+
+
+class ClinicalEncoder(TowerBase):
+    """MLP over tabular clinical features.
+
+    clinical_data : ClinicalDataView — provides feature_dim, vectorize_entity, side_map
+    hidden_dim    : output embedding dimensionality
+    dropout       : applied after the first linear block
+    use_se        : wrap output with SEBlock channel gating
+    se_reduction  : SEBlock bottleneck factor
+    se_pre_norm   : apply LayerNorm before SEBlock
+    """
+
+    def __init__(
+        self,
+        clinical_data,
+        hidden_dim: int = 128,
+        dropout: float = 0.1,
+        use_se: bool = False,
+        se_reduction: int = 16,
+        se_pre_norm: bool = True,
+    ):
+        super().__init__()
+        self.clinical_data = clinical_data
+        self._out_dim      = hidden_dim
+        feature_dim        = clinical_data.feature_dim
+
+        self.block0 = nn.Sequential(
+            nn.Linear(feature_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+        )
+        self.block1 = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True),
+        )
+        self.net = nn.Sequential(self.block0, self.block1)
+
+        self.tower_ln = nn.LayerNorm(hidden_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = SEBlock(hidden_dim, reduction=se_reduction, residual=True) if use_se else None
+
+    # ── TowerBase interface ──────────────────────────────────────────────────
+
+    @property
+    def out_dim(self) -> int:
+        return self._out_dim
+
+    @property
+    def _side_map(self) -> dict[str, str]:
+        return self.clinical_data.side_map
+
+    def _get(self, *ids) -> torch.Tensor:
+        arr = self.clinical_data.vectorize_entity(*ids)
+        return torch.from_numpy(arr.astype(np.float32, copy=False))
+
+    # ── nn.Module forward ────────────────────────────────────────────────────
+
+    def forward(self, x) -> torch.Tensor:
+        if not isinstance(x, torch.Tensor):
+            x = torch.as_tensor(x, dtype=torch.float32)
+        h = self.net(x)
+        if self.tower_se is not None:
+            h, _ = self.tower_se(self.tower_ln(h))
+        return h
+
+    # ── Utilities ────────────────────────────────────────────────────────────
+
+    def set_freeze_ratio(self, ratio: float) -> None:
+        """Freeze the earliest MLP block proportionally."""
+        r = max(0.0, min(1.0, float(ratio)))
+        for p in self.block0.parameters():
+            p.requires_grad = True
+        for p in self.block1.parameters():
+            p.requires_grad = True
+        if r >= 0.5:
+            for p in self.block0.parameters():
+                p.requires_grad = False
+        if r >= 1.0:
+            for p in self.block1.parameters():
+                p.requires_grad = False

v4/classes/towers/image_tower.py

@@ -0,0 +1,83 @@
+"""image_tower — ImageEncoder for v4.
+
+Self-contained: no v3 dependencies.
+Inherits get_sample dispatch from TowerBase.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+from v4.classes.towerbase import TowerBase
+from v4.classes.accessory.backbones import build_backbone
+from v4.classes.accessory.se_block import SEBlock
+from v4.classes.accessory.transforms import build_backbone_transform, build_eval_transform
+
+
+class ImageEncoder(TowerBase):
+    """Vision backbone → pooled feature vector.
+
+    image_data   : ImageDataView — provides load_image(*ids) and side_map
+    backbone     : backbone key (see accessory/backbones.py)
+    freeze_ratio : fraction of early blocks to freeze in [0, 1]
+    use_se       : apply SE attention over the pooled feature vector
+    augment      : include random flip/rotation/jitter in the train transform
+    """
+
+    def __init__(
+        self,
+        image_data,
+        backbone:     str   = "efficientnet_b0",
+        freeze_ratio: float = 0.0,
+        use_se:       bool  = False,
+        se_reduction: int   = 16,
+        se_pre_norm:  bool  = True,
+        augment:      bool  = True,
+    ):
+        super().__init__()
+        self.image_data     = image_data
+        self._name          = backbone
+        self.backbone, self._base_dim, self._blocks = build_backbone(backbone, freeze_ratio)
+        self.transform      = build_backbone_transform(backbone, augment=augment)
+        self.eval_transform = build_eval_transform(backbone)
+
+        self.tower_ln = nn.LayerNorm(self._base_dim) if se_pre_norm else nn.Identity()
+        self.tower_se = SEBlock(self._base_dim, reduction=se_reduction, residual=True) if use_se else None
+
+    # ── TowerBase interface ──────────────────────────────────────────────────
+
+    @property
+    def out_dim(self) -> int:
+        return self._base_dim
+
+    @property
+    def _side_map(self) -> dict[str, str]:
+        return self.image_data.side_map
+
+    def _get(self, *ids) -> torch.Tensor:
+        img = self.image_data.load_image(*ids)
+        t   = self.transform if self.training else self.eval_transform
+        return t(img)
+
+    # ── nn.Module forward ────────────────────────────────────────────────────
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.backbone(x)
+        if self.tower_se is not None:
+            y, _ = self.tower_se(self.tower_ln(y))
+        return y
+
+    # ── Utilities ────────────────────────────────────────────────────────────
+
+    def set_freeze_ratio(self, ratio: float) -> None:
+        """Dynamically freeze the earliest floor(N * ratio) backbone blocks."""
+        r        = max(0.0, min(1.0, float(ratio)))
+        n_freeze = int(math.floor(len(self._blocks) * r))
+        for b in self._blocks:
+            for p in b.parameters():
+                p.requires_grad = True
+        for b in self._blocks[:n_freeze]:
+            for p in b.parameters():
+                p.requires_grad = False

v4/classes/utils.py

@@ -0,0 +1,59 @@
+"""utils — general-purpose pipeline utilities."""
+from __future__ import annotations
+
+import random as pyrandom
+
+import numpy as np
+import pandas as pd
+import torch
+
+
+def seed_everything(seed: int) -> None:
+    pyrandom.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def choose_device(device_arg: str | None) -> torch.device:
+    if device_arg and device_arg != "auto":
+        return torch.device(device_arg)
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def drop_mixed_label_patients(
+    df: pd.DataFrame, *, patient_col: str, label_col: str
+):
+    """Remove patients whose rows carry conflicting labels.
+
+    Returns (clean_df, mixed_pids).
+    """
+    per_patient = (
+        df.groupby(patient_col)[label_col]
+        .agg(lambda s: set(pd.to_numeric(s, errors="coerce").dropna().astype(int).tolist()))
+    )
+    mixed = [pid for pid, labels in per_patient.items() if len(labels) > 1]
+    if not mixed:
+        return df, []
+    return df[~df[patient_col].isin(mixed)].reset_index(drop=True), mixed
+
+
+def relabel_mixed_patients_to_max(
+    df: pd.DataFrame, *, patient_col: str, label_col: str
+):
+    """Set all rows for each patient to that patient's max observed label.
+
+    Returns (df, changed_rows, still_mixed_pids).
+    """
+    out         = df.copy()
+    labels      = pd.to_numeric(out[label_col], errors="coerce")
+    patient_max = labels.groupby(out[patient_col]).transform("max")
+    changed_rows = int((labels != patient_max).fillna(False).sum())
+    out[label_col] = patient_max.astype(int)
+    still_mixed = (
+        out.groupby(patient_col)[label_col]
+        .nunique(dropna=True)
+        .pipe(lambda s: s[s > 1].index.tolist())
+    )
+    return out.reset_index(drop=True), changed_rows, still_mixed

v4/classes/v4_hypertower.py

@@ -0,0 +1,285 @@
+#!/usr/bin/env python
+"""
+V4 HyperTower orchestrator — config-driven stage pipeline.
+
+Stage logic lives in v4/classes/stages/:
+  warm.py    — pre-trains a single tower with a temporary linear probe
+  fusion.py  — trains a bridge + associated head stages, freezes for downstream use
+  helpers.py — encode_embedding, resolve_input_dims, phase_for_epoch, etc.
+
+Usage:
+  python -m v4.classes.v4_hypertower --config v4/configs/ensemble_fused.json
+"""
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from torch.utils.data import DataLoader
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(REPO_ROOT))
+
+from v4.classes.dataset import LoaderShell, HTDataset, ht_collate
+from v4.classes.utils import seed_everything, choose_device
+from v4.classes.split_manager import SplitManager
+from v4.classes.stages import warm, fusion
+
+
+# ---------------------------------------------------------------------------
+# Early-pass protocol
+# ---------------------------------------------------------------------------
+
+class EarlyPassContext:
+    def __init__(self) -> None:
+        self._store: dict[str, Any] = {}
+
+    def put(self, key: str, value: Any) -> None:
+        self._store[key] = value
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self._store.get(key, default)
+
+    def require(self, key: str) -> Any:
+        if key not in self._store:
+            raise KeyError(
+                f"EarlyPassContext: required key '{key}' not present. "
+                f"Available: {sorted(self._store.keys())}"
+            )
+        return self._store[key]
+
+    def keys(self) -> set[str]:
+        return set(self._store.keys())
+
+
+def _validate_epc_requests(towers_cfg: list[dict], provided_keys: set[str]) -> None:
+    available = set(provided_keys)
+    for t in towers_cfg:
+        for req in t.get("epc_requests", []):
+            if req not in available:
+                raise ValueError(
+                    f"Tower '{t['name']}' requests EPC key '{req}' "
+                    f"but no supplier provides it. Available: {sorted(available)}"
+                )
+        available.update(t.get("epc_supplies", []))
+
+
+# ---------------------------------------------------------------------------
+# Data + tower helpers
+# ---------------------------------------------------------------------------
+
+cfg_ref: dict = {}
+
+
+def load_data(cfg: dict):
+    data_cfg = cfg["data"]
+    args = dict(data_cfg.get("args", {}))
+    for key in ("image_dir", "clinical_dir"):
+        if key in args:
+            p = Path(args[key])
+            if not p.is_absolute():
+                args[key] = str(REPO_ROOT / p)
+    mod = importlib.import_module(data_cfg["module"])
+    return mod.build_data(args)
+
+
+def build_towers(towers_cfg: list[dict], data) -> dict:
+    def _resolve(path: str):
+        mod_name = cfg_ref.get("data", {}).get("module", "")
+        if mod_name:
+            try:
+                m = importlib.import_module(mod_name)
+                if hasattr(m, "resolve_data_source"):
+                    return m.resolve_data_source(data, path)
+            except Exception:
+                pass
+        obj = data
+        for part in path.split("."):
+            obj = getattr(obj, part)
+        return obj
+
+    towers = {}
+    for t in towers_cfg:
+        mod = importlib.import_module(t["module"])
+        cls = getattr(mod, t["class"])
+        kwargs = dict(t.get("args", {}))
+        if "data_source" in t:
+            towers[t["name"]] = cls(_resolve(t["data_source"]), **kwargs)
+        elif "data_arg" in t:
+            kwargs[t["data_arg"]] = data
+            towers[t["name"]] = cls(**kwargs)
+        else:
+            towers[t["name"]] = cls(**kwargs)
+    return towers
+
+
+def _make_loader(shell: LoaderShell, towers: dict, *, batch_size: int,
+                 shuffle: bool, sampler=None) -> DataLoader:
+    dataset = HTDataset(shell, towers)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=(shuffle and sampler is None),
+        sampler=sampler,
+        collate_fn=ht_collate,
+        num_workers=0,
+        persistent_workers=False,
+    )
+
+
+def _balanced_sampler(shell: LoaderShell):
+    from torch.utils.data import WeightedRandomSampler
+    labels  = [e.label for e in shell.entries]
+    counts  = {}
+    for l in labels:
+        counts[l] = counts.get(l, 0) + 1
+    weights = [1.0 / counts[l] for l in labels]
+    return WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
+
+
+# ---------------------------------------------------------------------------
+# Fold runner
+# ---------------------------------------------------------------------------
+
+def run_fold(fold: int, splits, cfg: dict, data, num_classes: int, device) -> dict:
+    seed_everything(cfg["seed"] + fold * 100)
+    split        = splits[fold]
+    label_filter = cfg.get("label_filter", None)
+    cfg_stages   = cfg["stages"]
+
+    towers = build_towers(cfg["towers"], data)
+    for t in towers.values():
+        t.to(device)
+
+    context = EarlyPassContext()
+    context.put("device",       device)
+    context.put("data",         data)
+    context.put("split",        split)
+    context.put("label_filter", label_filter)
+    _validate_epc_requests(cfg["towers"], context.keys())
+    for tower in towers.values():
+        if hasattr(tower, "early_pass"):
+            tower.early_pass(context)
+
+    stage_models: dict = {}
+    fold_result         = {"fold": fold}
+
+    for stage_cfg in cfg_stages:
+        stype = stage_cfg["type"]
+
+        if stype == "warm":
+            warm.run(stage_cfg, towers, data, split, label_filter,
+                     cfg, num_classes, device, fold,
+                     _make_loader, _balanced_sampler)
+
+        elif stype == "fusion":
+            stage_models, metrics = fusion.run(
+                stage_cfg, cfg, towers, stage_models, data, split,
+                label_filter, num_classes, device, fold, cfg_stages,
+                _make_loader,
+            )
+            fold_result.update(metrics)
+
+        # head stages are handled inside fusion.run
+
+    return fold_result
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    global cfg_ref
+
+    ap = argparse.ArgumentParser(description=__doc__,
+                                 formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--device", default=None)
+    args = ap.parse_args()
+
+    with open(args.config) as f:
+        cfg = json.load(f)
+    cfg_ref = cfg
+
+    device      = choose_device(args.device or cfg.get("device"))
+    num_classes = cfg.get("num_classes", 2)
+    label_filter = cfg.get("label_filter", None)
+    print(f"Device: {device}", flush=True)
+
+    print("Loading data ...", flush=True)
+    data = load_data(cfg)
+    print(f"  feature_dim={data.feature_dim}", flush=True)
+
+    label_col = cfg["data"]["args"].get("label_col", "Diagnosis")
+    df_mode   = data.df.copy()
+    if label_filter is not None:
+        df_mode = df_mode[df_mode[label_col].isin(label_filter)].reset_index(drop=True)
+
+    identity_level = cfg.get("split_identity_level", 1)
+    identity_cols  = getattr(data, "identity_cols", [])
+    group_col = identity_cols[identity_level - 1] if identity_level and identity_cols else None
+
+    splits = SplitManager(group_col=group_col).build_plans(
+        df_mode,
+        label_col=label_col,
+        n_splits=cfg.get("folds", 5),
+        seed=cfg.get("fold_seed", 100),
+    )
+
+    out_dir_tags = cfg.get("out_dir_tags", [])
+    out_dir = REPO_ROOT / cfg.get("output_root", "v4/results") / cfg["run_name"]
+    for tag in out_dir_tags:
+        out_dir = out_dir / tag
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    eval_stage   = cfg.get("eval_stage", "hb")
+    fold_results = []
+    t0 = time.time()
+
+    for fold in range(cfg.get("folds", 5)):
+        split   = splits[fold]
+        n_train = split.train[group_col].nunique() if group_col else len(split.train)
+        print(f"\n── fold {fold+1}/{cfg.get('folds', 5)}  train_groups={n_train} ──",
+              flush=True)
+        result = run_fold(fold, splits, cfg, data, num_classes, device)
+        fold_results.append(result)
+        print(
+            f"  fold{fold+1} DONE"
+            f"  val_auc={result.get(f'{eval_stage}_val_auc', float('nan')):.4f}"
+            f"  test_auc={result.get(f'{eval_stage}_test_auc', float('nan')):.4f}",
+            flush=True,
+        )
+
+    if fold_results:
+        val_aucs  = [r.get(f"{eval_stage}_val_auc",  float("nan")) for r in fold_results]
+        test_aucs = [r.get(f"{eval_stage}_test_auc", float("nan")) for r in fold_results]
+        val_aucs  = [v for v in val_aucs  if not np.isnan(v)]
+        test_aucs = [v for v in test_aucs if not np.isnan(v)]
+        summary = {
+            "run_name":      cfg["run_name"],
+            "eval_stage":    eval_stage,
+            "config":        cfg,
+            "mean_val_auc":  float(np.mean(val_aucs))  if val_aucs  else float("nan"),
+            "std_val_auc":   float(np.std(val_aucs))   if val_aucs  else float("nan"),
+            "mean_test_auc": float(np.mean(test_aucs)) if test_aucs else float("nan"),
+            "std_test_auc":  float(np.std(test_aucs))  if test_aucs else float("nan"),
+            "elapsed_s":     round(time.time() - t0, 1),
+            "fold_results":  fold_results,
+        }
+        summary_path = out_dir / "summary.json"
+        summary_path.write_text(json.dumps(summary, indent=2))
+        print(f"\n{'='*60}", flush=True)
+        print(f"Val  AUC: {summary['mean_val_auc']:.4f} ± {summary['std_val_auc']:.4f}", flush=True)
+        print(f"Test AUC: {summary['mean_test_auc']:.4f} ± {summary['std_test_auc']:.4f}", flush=True)
+        print(f"Saved: {summary_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

v4/classes/vehicles-deprecated/htbase.py

@@ -0,0 +1,40 @@
+"""htbase — HTBase: abstract base for all v4 vehicle classes."""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+
+
+class HTBase(nn.Module, ABC):
+    """Shared interface for all HyperTower vehicles.
+
+    Subclasses must implement ``encode`` and ``forward``.
+
+    ``transform`` walks ``self.towers`` (if present) and returns the transform
+    from the first tower that exposes one — used by data loaders.
+
+    ``forward`` contract: returns ``(logits, aux_dict)`` where
+    ``aux_dict`` maps a name or index to per-component logits.
+    HTMono returns an empty dict to keep the signature uniform.
+    """
+
+    @abstractmethod
+    def encode(self, inputs) -> torch.Tensor:
+        """Return the pre-classifier embedding."""
+
+    @abstractmethod
+    def forward(self, inputs) -> tuple[torch.Tensor, dict]:
+        """Return (logits, aux_dict)."""
+
+    @property
+    def transform(self):
+        towers = getattr(self, "towers", None) or {}
+        for t in (towers.values() if hasattr(towers, "values") else []):
+            if hasattr(t, "transform"):
+                return t.transform
+        encoder = getattr(self, "encoder", None)
+        if encoder is not None:
+            return getattr(encoder, "transform", None)
+        return None

v4/classes/vehicles-deprecated/htfusion.py

@@ -0,0 +1,61 @@
+"""htfusion — HTFusion: N named towers fused through a FusionBridge."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from v4.classes.bridges.fusion_bridge import FusionBridge
+from v4.classes.vehicles.htbase import HTBase
+
+
+class HTFusion(HTBase):
+    """General N-tower fusion vehicle.
+
+    Each named encoder is registered as a submodule; the FusionBridge
+    projects and Hadamard-fuses their embeddings.
+
+    Parameters
+    ----------
+    towers      : ordered dict ``{name: encoder}``.  Each encoder must
+                  expose ``.out_dim``.
+    num_classes : output classes
+    fusion_dim  : bridge projection dimensionality
+    dropout     : bridge dropout
+    use_se      : SE gate on the fused vector
+
+    Forward contract
+    ----------------
+    ``forward(embeddings)`` takes a ``dict[str, Tensor]`` of pre-computed
+    per-tower embeddings and returns ``(logits_fused, aux_dict)`` where
+    ``aux_dict`` maps each tower name to its auxiliary head logits.
+    """
+
+    def __init__(
+        self,
+        towers:      dict[str, nn.Module],
+        num_classes: int,
+        fusion_dim:  int   = 256,
+        dropout:     float = 0.5,
+        use_se:      bool  = False,
+    ):
+        super().__init__()
+        self.towers = nn.ModuleDict(towers)
+        self.bridge = FusionBridge(
+            tower_dims=[t.out_dim for t in self.towers.values()],
+            num_classes=num_classes,
+            fusion_dim=fusion_dim,
+            dropout=dropout,
+            use_se=use_se,
+        )
+
+    def encode(self, embeddings: dict[str, torch.Tensor]) -> torch.Tensor:
+        """Return z_fused (pre-classifier) from a dict of per-tower embeddings."""
+        return self.bridge.encode([embeddings[name] for name in self.towers])
+
+    def forward(
+        self,
+        embeddings: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        ordered = [embeddings[name] for name in self.towers]
+        logits, aux = self.bridge.fuse(ordered)
+        return logits, {name: aux[i] for i, name in enumerate(self.towers)}

v4/classes/vehicles-deprecated/htlateral.py

@@ -0,0 +1,67 @@
+"""htlateral — HTLateral: shared encoder over N same-type inputs."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from v4.classes.vehicles.htbase import HTBase
+
+
+class HTLateral(HTBase):
+    """N same-type inputs through a shared encoder, jointly compressed, then classified.
+
+    All inputs share the same encoder weights (one forward pass per input).
+    The joint MLP compresses the concatenated embeddings before classification.
+
+    Aux heads provide per-input logits before the joint MLP — useful for
+    BCD-style training.
+
+    Parameters
+    ----------
+    encoder     : shared encoder module with ``.out_dim``
+    input_names : ordered slot names (e.g. ``["od", "os"]``)
+    num_classes : output classes
+    fusion_dim  : joint MLP hidden dim
+    dropout     : dropout in MLP and classifier
+    """
+
+    def __init__(
+        self,
+        encoder:     nn.Module,
+        input_names: list[str],
+        num_classes: int,
+        fusion_dim:  int   = 256,
+        dropout:     float = 0.5,
+    ):
+        super().__init__()
+        self.encoder     = encoder
+        self.input_names = list(input_names)
+        n       = len(input_names)
+        in_dim: int = encoder.out_dim  # type: ignore[assignment]
+
+        self.joint = nn.Sequential(
+            nn.Linear(n * in_dim, fusion_dim), nn.LayerNorm(fusion_dim),
+            nn.ReLU(), nn.Dropout(dropout), nn.Linear(fusion_dim, in_dim),
+        )
+        self.aux_heads = nn.ModuleList([
+            nn.Linear(in_dim, num_classes) for _ in range(n)
+        ])
+        self.head = nn.Sequential(
+            nn.ReLU(), nn.Dropout(dropout), nn.Linear(in_dim, num_classes),
+        )
+
+    def encode(self, inputs: dict[str, torch.Tensor]) -> torch.Tensor:
+        """Return joint embedding (post-MLP, pre-classifier)."""
+        zs = [self.encoder(inputs[name]) for name in self.input_names]
+        return self.joint(torch.cat(zs, dim=1))
+
+    def forward(
+        self,
+        inputs: dict[str, torch.Tensor],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        zs      = [self.encoder(inputs[name]) for name in self.input_names]
+        z_joint = self.joint(torch.cat(zs, dim=1))
+        logits  = self.head(z_joint)
+        aux     = {name: head(z)
+                   for name, head, z in zip(self.input_names, self.aux_heads, zs)}
+        return logits, aux

v4/classes/vehicles-deprecated/htmono.py

@@ -0,0 +1,44 @@
+"""htmono — HTMono: single tower + ClassificationHead, no bridge."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from v4.classes.heads.classifier import ClassificationHead
+from v4.classes.vehicles.htbase import HTBase
+
+
+class HTMono(HTBase):
+    """Single-tower vehicle: tower embedding fed directly into a ClassificationHead.
+
+    No bridge or projection — the tower's output goes straight to
+    ReLU → Dropout → Linear.  Returns ``(logits, {})`` from ``forward``
+    to match the HTFusion / HTLateral interface.
+
+    Parameters
+    ----------
+    tower       : encoder module with ``.out_dim``
+    num_classes : output classes
+    dropout     : dropout before the output linear layer
+    """
+
+    def __init__(
+        self,
+        tower:       nn.Module,
+        num_classes: int,
+        dropout:     float = 0.5,
+    ):
+        super().__init__()
+        self.tower = tower
+        self.head  = ClassificationHead(tower.out_dim, num_classes, dropout)  # type: ignore[arg-type]
+
+    def encode(self, inputs) -> torch.Tensor:
+        """Return tower embedding (pre-classifier)."""
+        if isinstance(inputs, dict):
+            # single-entry dict from HTDataset eye-level pass
+            (z,) = inputs.values()
+            return self.tower(z) if torch.is_tensor(z) else self.tower(*z.values())
+        return self.tower(inputs)
+
+    def forward(self, inputs) -> tuple[torch.Tensor, dict]:
+        return self.head(self.encode(inputs)), {}

v4/configs/ensemble_fused.json

@@ -0,0 +1,121 @@
+{
+  "_notes": [
+    "V4 stage-pipeline: warm → fusion stages with parallel head stages.",
+    "Bridges are pure embedding producers; heads are separate swappable stages.",
+    "BCD-eligible heads are sampled during tower_warmup and main phases.",
+    "eval_stage names the fusion stage whose primary head is used for final metrics."
+  ],
+
+  "run_name":            "v4/ensemble_fused",
+  "num_classes":         2,
+  "label_filter":        [0, 1],
+  "split_identity_level": 1,
+  "eval_stage":          "hb",
+  "save_predictions":    false,
+  "seed":                1234,
+  "folds":               5,
+  "fold_seed":           100,
+  "output_root":         "v4/results",
+  "out_dir_tags":        ["binary", "ntower"],
+
+  "data": {
+    "module": "v4.classes.profiles.v4papila",
+    "args": {
+      "image_dir":        "Papila/FundusImages",
+      "clinical_dir":     "Papila/ClinicalData",
+      "label_col":        "Diagnosis",
+      "iop_corr_method":  "ratio",
+      "iop_drop_raw":     true,
+      "exclude_cols":     ["Axial_Length"],
+      "in_memory_cache":  true
+    }
+  },
+
+  "towers": [
+    {
+      "name":        "img",
+      "module":      "v4.classes.towers.image_tower",
+      "class":       "ImageEncoder",
+      "data_source": "image",
+      "args": {
+        "backbone":     "refugelike",
+        "freeze_ratio": 0.0,
+        "augment":      true
+      }
+    },
+    {
+      "name":        "cd",
+      "module":      "v4.classes.towers.clinical_tower",
+      "class":       "ClinicalEncoder",
+      "data_source": "matrix",
+      "args": {
+        "hidden_dim": 128
+      }
+    }
+  ],
+
+  "stages": [
+    {
+      "name":    "cd_warm",
+      "type":    "warm",
+      "tower":   "cd",
+      "level":   "eye",
+      "epochs":  40
+    },
+    {
+      "name":       "img_aux",
+      "type":       "head",
+      "input":      "img",
+      "train_with": "nt",
+      "bcd":        true
+    },
+    {
+      "name":       "cd_aux",
+      "type":       "head",
+      "input":      "cd",
+      "train_with": "nt",
+      "bcd":        true
+    },
+    {
+      "name":         "nt",
+      "type":         "fusion",
+      "module":       "v4.classes.bridges.fusion_bridge",
+      "class":        "FusionBridge",
+      "inputs":       ["img", "cd"],
+      "level":        "eye",
+      "epochs":       36,
+      "train_towers": true,
+      "warmup":       { "tower_epochs": 3, "fused_epochs": 3 },
+      "args":         { "fusion_dim": 256 }
+    },
+    {
+      "name":       "nt_head",
+      "type":       "head",
+      "input":      "nt",
+      "train_with": "nt"
+    },
+    {
+      "name":    "hb",
+      "type":    "fusion",
+      "module":  "v4.classes.bridges.hyperbridge",
+      "class":   "HyperBridge",
+      "inputs":  { "a": "nt", "b": "nt" },
+      "level":   "patient",
+      "epochs":  10,
+      "args":    { "hidden_dim": 256, "mode": "embedding_mlp" }
+    },
+    {
+      "name":       "hb_head",
+      "type":       "head",
+      "input":      "hb",
+      "train_with": "hb"
+    }
+  ],
+
+  "training": {
+    "lr":                    1e-4,
+    "batch_size":            16,
+    "bcd_prob":              0.5,
+    "tune_binary_threshold": true
+  }
+}

v4/results/v4/ensemble_fused/binary/ntower/gold_standard.json

@@ -0,0 +1,146 @@
+{
+  "run_name": "v4/ensemble_fused",
+  "config": {
+    "_notes": [
+      "V4 ensemble_fused: img + cd towers, HTFusion Stage 1, HyperBridge Stage 2.",
+      "Matches phase5/embedding_mlp_head setup for direct comparison.",
+      "data_source resolves against the PapilaBundle returned by build_data.",
+      "image_dir / clinical_dir are project-relative; orchestrator resolves against REPO_ROOT."
+    ],
+    "run_name": "v4/ensemble_fused",
+    "eval_mode": "binary",
+    "split_identity_level": 1,
+    "epochs": 30,
+    "fusion_epochs": 10,
+    "folds": 5,
+    "fold_seed": 100,
+    "seed": 1234,
+    "output_root": "v4/results",
+    "data": {
+      "module": "v4.classes.profiles.v4papila",
+      "args": {
+        "image_dir": "Papila/FundusImages",
+        "clinical_dir": "Papila/ClinicalData",
+        "label_col": "Diagnosis",
+        "iop_corr_method": "ratio",
+        "iop_drop_raw": true,
+        "exclude_cols": [
+          "Axial_Length"
+        ],
+        "in_memory_cache": true
+      }
+    },
+    "towers": [
+      {
+        "name": "img",
+        "module": "v4.classes.towers.image_tower",
+        "class": "ImageEncoder",
+        "data_source": "image",
+        "args": {
+          "backbone": "refugelike",
+          "freeze_ratio": 0.0,
+          "augment": true
+        },
+        "warmup_epochs": 0
+      },
+      {
+        "name": "cd",
+        "module": "v4.classes.towers.clinical_tower",
+        "class": "ClinicalEncoder",
+        "data_source": "matrix",
+        "args": {
+          "hidden_dim": 128
+        },
+        "warmup_epochs": 40,
+        "warmup_exclude_towers": [
+          "img"
+        ]
+      }
+    ],
+    "bridge": {
+      "mode": "embedding_mlp",
+      "fusion_dim": 256,
+      "hidden_dim": 256
+    },
+    "training": {
+      "lr": 0.0001,
+      "batch_size": 16,
+      "bcd_prob": 0.5,
+      "warmup_tower_epochs": 3,
+      "warmup_fused_epochs": 3,
+      "tune_binary_threshold": true
+    }
+  },
+  "mean_val_auc": 0.9036764705882353,
+  "std_val_auc": 0.02318218054475656,
+  "mean_test_auc": 0.9007352941176471,
+  "std_test_auc": 0.04178903599524356,
+  "elapsed_s": 1563.1,
+  "fold_results": [
+    {
+      "fold": 0,
+      "val_auc": 0.9044117647058824,
+      "val_acc": 0.8571428571428571,
+      "val_n": 42,
+      "val_kappa": 0.42727272727272725,
+      "val_mcc": 0.4622975667767223,
+      "val_f1": 0.7083333333333333,
+      "val_threshold": 0.1659889668226242,
+      "test_auc": 0.8566176470588236,
+      "test_acc": 0.9285714285714286,
+      "test_n": 42
+    },
+    {
+      "fold": 1,
+      "val_auc": 0.8970588235294118,
+      "val_acc": 0.8809523809523809,
+      "val_n": 42,
+      "val_kappa": 0.5945945945945946,
+      "val_mcc": 0.5965587590013045,
+      "val_f1": 0.7971014492753623,
+      "val_threshold": 0.017083797603845596,
+      "test_auc": 0.9044117647058824,
+      "test_acc": 0.8095238095238095,
+      "test_n": 42
+    },
+    {
+      "fold": 2,
+      "val_auc": 0.863970588235294,
+      "val_acc": 0.8333333333333334,
+      "val_n": 42,
+      "val_kappa": 0.43243243243243246,
+      "val_mcc": 0.4338609156373123,
+      "val_f1": 0.7159420289855072,
+      "val_threshold": 0.09843172132968903,
+      "test_auc": 0.9044117647058824,
+      "test_acc": 0.8571428571428571,
+      "test_n": 42
+    },
+    {
+      "fold": 3,
+      "val_auc": 0.9227941176470588,
+      "val_acc": 0.8809523809523809,
+      "val_n": 42,
+      "val_kappa": 0.631578947368421,
+      "val_mcc": 0.6333004963811236,
+      "val_f1": 0.8156277436347674,
+      "val_threshold": 0.01531070377677679,
+      "test_auc": 0.8639705882352942,
+      "test_acc": 0.8571428571428571,
+      "test_n": 42
+    },
+    {
+      "fold": 4,
+      "val_auc": 0.9301470588235294,
+      "val_acc": 0.9047619047619048,
+      "val_n": 42,
+      "val_kappa": 0.6181818181818182,
+      "val_mcc": 0.6688560540599386,
+      "val_f1": 0.8055555555555556,
+      "val_threshold": 0.19587548077106476,
+      "test_auc": 0.9742647058823529,
+      "test_acc": 0.9285714285714286,
+      "test_n": 42
+    }
+  ]
+}

v4/results/v4/ensemble_fused/binary/ntower/summary.json

@@ -0,0 +1,264 @@
+{
+  "run_name": "v4/ensemble_fused",
+  "eval_stage": "hb",
+  "config": {
+    "_notes": [
+      "V4 stage-pipeline: warm \u2192 fusion stages with parallel head stages.",
+      "Bridges are pure embedding producers; heads are separate swappable stages.",
+      "BCD-eligible heads are sampled during tower_warmup and main phases.",
+      "eval_stage names the fusion stage whose primary head is used for final metrics."
+    ],
+    "run_name": "v4/ensemble_fused",
+    "num_classes": 2,
+    "label_filter": [
+      0,
+      1
+    ],
+    "split_identity_level": 1,
+    "eval_stage": "hb",
+    "save_predictions": false,
+    "seed": 1234,
+    "folds": 5,
+    "fold_seed": 100,
+    "output_root": "v4/results",
+    "out_dir_tags": [
+      "binary",
+      "ntower"
+    ],
+    "data": {
+      "module": "v4.classes.profiles.v4papila",
+      "args": {
+        "image_dir": "Papila/FundusImages",
+        "clinical_dir": "Papila/ClinicalData",
+        "label_col": "Diagnosis",
+        "iop_corr_method": "ratio",
+        "iop_drop_raw": true,
+        "exclude_cols": [
+          "Axial_Length"
+        ],
+        "in_memory_cache": true
+      }
+    },
+    "towers": [
+      {
+        "name": "img",
+        "module": "v4.classes.towers.image_tower",
+        "class": "ImageEncoder",
+        "data_source": "image",
+        "args": {
+          "backbone": "refugelike",
+          "freeze_ratio": 0.0,
+          "augment": true
+        }
+      },
+      {
+        "name": "cd",
+        "module": "v4.classes.towers.clinical_tower",
+        "class": "ClinicalEncoder",
+        "data_source": "matrix",
+        "args": {
+          "hidden_dim": 128
+        }
+      }
+    ],
+    "stages": [
+      {
+        "name": "cd_warm",
+        "type": "warm",
+        "tower": "cd",
+        "level": "eye",
+        "epochs": 40
+      },
+      {
+        "name": "img_aux",
+        "type": "head",
+        "input": "img",
+        "train_with": "nt",
+        "bcd": true
+      },
+      {
+        "name": "cd_aux",
+        "type": "head",
+        "input": "cd",
+        "train_with": "nt",
+        "bcd": true
+      },
+      {
+        "name": "nt",
+        "type": "fusion",
+        "module": "v4.classes.bridges.fusion_bridge",
+        "class": "FusionBridge",
+        "inputs": [
+          "img",
+          "cd"
+        ],
+        "level": "eye",
+        "epochs": 36,
+        "train_towers": true,
+        "warmup": {
+          "tower_epochs": 3,
+          "fused_epochs": 3
+        },
+        "args": {
+          "fusion_dim": 256
+        }
+      },
+      {
+        "name": "nt_head",
+        "type": "head",
+        "input": "nt",
+        "train_with": "nt"
+      },
+      {
+        "name": "hb",
+        "type": "fusion",
+        "module": "v4.classes.bridges.hyperbridge",
+        "class": "HyperBridge",
+        "inputs": {
+          "a": "nt",
+          "b": "nt"
+        },
+        "level": "patient",
+        "epochs": 10,
+        "args": {
+          "hidden_dim": 256,
+          "mode": "embedding_mlp"
+        }
+      },
+      {
+        "name": "hb_head",
+        "type": "head",
+        "input": "hb",
+        "train_with": "hb"
+      }
+    ],
+    "training": {
+      "lr": 0.0001,
+      "batch_size": 16,
+      "bcd_prob": 0.5,
+      "tune_binary_threshold": true
+    }
+  },
+  "mean_val_auc": 0.8926470588235293,
+  "std_val_auc": 0.06105155516072669,
+  "mean_test_auc": 0.9022058823529413,
+  "std_test_auc": 0.035001853633564395,
+  "elapsed_s": 1525.9,
+  "fold_results": [
+    {
+      "fold": 0,
+      "nt_val_auc": 0.858647936786655,
+      "nt_val_acc": 0.8928571428571429,
+      "nt_val_n": 84,
+      "nt_val_kappa": 0.6272189349112426,
+      "nt_val_mcc": 0.6411186083279721,
+      "nt_val_f1": 0.8124534854874721,
+      "nt_val_threshold": 0.012987074442207813,
+      "nt_test_auc": 0.9157155399473222,
+      "nt_test_acc": 0.9047619047619048,
+      "nt_test_n": 84,
+      "hb_val_auc": 0.9080882352941178,
+      "hb_val_acc": 0.8333333333333334,
+      "hb_val_n": 42,
+      "hb_val_kappa": 0.2898550724637682,
+      "hb_val_mcc": 0.33633639699815626,
+      "hb_val_f1": 0.6338729763387297,
+      "hb_val_threshold": 0.011891158297657967,
+      "hb_test_auc": 0.9522058823529411,
+      "hb_test_acc": 0.9523809523809523,
+      "hb_test_n": 42
+    },
+    {
+      "fold": 1,
+      "nt_val_auc": 0.8282828282828283,
+      "nt_val_acc": 0.8928571428571429,
+      "nt_val_n": 84,
+      "nt_val_kappa": 0.6111111111111112,
+      "nt_val_mcc": 0.6633249580710799,
+      "nt_val_f1": 0.801418439716312,
+      "nt_val_threshold": 0.008082838729023933,
+      "nt_test_auc": 0.8726953467954346,
+      "nt_test_acc": 0.8690476190476191,
+      "nt_test_n": 84,
+      "hb_val_auc": 0.7904411764705882,
+      "hb_val_acc": 0.9047619047619048,
+      "hb_val_n": 42,
+      "hb_val_kappa": 0.6181818181818182,
+      "hb_val_mcc": 0.6688560540599386,
+      "hb_val_f1": 0.8055555555555556,
+      "hb_val_threshold": 0.013159404508769512,
+      "hb_test_auc": 0.8970588235294117,
+      "hb_test_acc": 0.8571428571428571,
+      "hb_test_n": 42
+    },
+    {
+      "fold": 2,
+      "nt_val_auc": 0.8906882591093117,
+      "nt_val_acc": 0.8452380952380952,
+      "nt_val_n": 84,
+      "nt_val_kappa": 0.5125,
+      "nt_val_mcc": 0.5217535056401378,
+      "nt_val_f1": 0.7548821548821549,
+      "nt_val_threshold": 0.052708517760038376,
+      "nt_test_auc": 0.856060606060606,
+      "nt_test_acc": 0.8333333333333334,
+      "nt_test_n": 84,
+      "hb_val_auc": 0.9117647058823529,
+      "hb_val_acc": 0.8333333333333334,
+      "hb_val_n": 42,
+      "hb_val_kappa": 0.36909871244635195,
+      "hb_val_mcc": 0.3833788364965519,
+      "hb_val_f1": 0.6814734561213435,
+      "hb_val_threshold": 0.0401872955262661,
+      "hb_test_auc": 0.8566176470588236,
+      "hb_test_acc": 0.8571428571428571,
+      "hb_test_n": 42
+    },
+    {
+      "fold": 3,
+      "nt_val_auc": 0.921875,
+      "nt_val_acc": 0.8928571428571429,
+      "nt_val_n": 84,
+      "nt_val_kappa": 0.6758147512864494,
+      "nt_val_mcc": 0.6797955088067001,
+      "nt_val_f1": 0.837593984962406,
+      "nt_val_threshold": 0.708580732345581,
+      "nt_test_auc": 0.8631578947368421,
+      "nt_test_acc": 0.8452380952380952,
+      "nt_test_n": 84,
+      "hb_val_auc": 0.9779411764705882,
+      "hb_val_acc": 0.9285714285714286,
+      "hb_val_n": 42,
+      "hb_val_kappa": 0.7567567567567568,
+      "hb_val_mcc": 0.7592566023652966,
+      "hb_val_f1": 0.8782608695652174,
+      "hb_val_threshold": 0.04653067886829376,
+      "hb_test_auc": 0.875,
+      "hb_test_acc": 0.8809523809523809,
+      "hb_test_n": 42
+    },
+    {
+      "fold": 4,
+      "nt_val_auc": 0.8279192273924496,
+      "nt_val_acc": 0.8571428571428571,
+      "nt_val_n": 84,
+      "nt_val_kappa": 0.46325878594249204,
+      "nt_val_mcc": 0.49610717544581684,
+      "nt_val_f1": 0.7269772481040087,
+      "nt_val_threshold": 0.06419441103935242,
+      "nt_test_auc": 0.8961397058823529,
+      "nt_test_acc": 0.8928571428571429,
+      "nt_test_n": 84,
+      "hb_val_auc": 0.875,
+      "hb_val_acc": 0.9047619047619048,
+      "hb_val_n": 42,
+      "hb_val_kappa": 0.6181818181818182,
+      "hb_val_mcc": 0.6688560540599386,
+      "hb_val_f1": 0.8055555555555556,
+      "hb_val_threshold": 0.23509082198143005,
+      "hb_test_auc": 0.9301470588235294,
+      "hb_test_acc": 0.8809523809523809,
+      "hb_test_n": 42
+    }
+  ]
+}

v4/scripts/debug_splits.py

@@ -0,0 +1,65 @@
+"""Compare fold assignments between v3 PatientFirstSplitManager and v4 SplitManager."""
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+
+from v3.classes.split_manager import PatientFirstSplitManager
+from v4.classes.split_manager import SplitManager
+from v4.classes.profiles.v4papila import build_data
+
+args = {
+    "image_dir":       "Papila/FundusImages",
+    "clinical_dir":    "Papila/ClinicalData",
+    "label_col":       "Diagnosis",
+    "iop_corr_method": "ratio",
+    "iop_drop_raw":    True,
+    "exclude_cols":    ["Axial_Length"],
+}
+# Resolve relative paths
+root = Path(__file__).resolve().parents[2]
+args["image_dir"]    = str(root / args["image_dir"])
+args["clinical_dir"] = str(root / args["clinical_dir"])
+
+data = build_data(args)
+
+label_col   = data.label_col
+patient_col = data.patient_col
+df_mode     = data.df[data.df[label_col].isin([0, 1])].reset_index(drop=True)
+
+# ── v3 splits ────────────────────────────────────────────────────────────────
+split_mgr_v3  = PatientFirstSplitManager(patient_col=patient_col, label_col=label_col)
+split_args_v3 = SimpleNamespace(eval_mode="binary", n_splits=5, fold_seed=100)
+clinical_ns   = SimpleNamespace(df=df_mode, label_col=label_col)
+splits_v3     = split_mgr_v3.build_plans(clinical=clinical_ns, args=split_args_v3, profile=None)
+
+# ── v4 splits ────────────────────────────────────────────────────────────────
+splits_v4 = SplitManager(group_col=patient_col).build_plans(
+    df_mode, label_col=label_col, n_splits=5, seed=100,
+)
+
+# ── Compare ──────────────────────────────────────────────────────────────────
+print(f"{'Fold':<6} {'Set':<6} {'v3 patients':<8} {'v4 patients':<8} {'Match'}")
+print("-" * 50)
+
+all_match = True
+for fold in range(5):
+    s3, s4 = splits_v3[fold], splits_v4[fold]
+    for label, df3, df4 in [
+        ("train", s3.train, s4.train),
+        ("val",   s3.val,   s4.val),
+        ("test",  s3.test,  s4.test),
+    ]:
+        ids3 = set(df3[patient_col].unique()) if df3 is not None else set()
+        ids4 = set(df4[patient_col].unique()) if df4 is not None else set()
+        match = ids3 == ids4
+        if not match:
+            all_match = False
+        print(f"{fold+1:<6} {label:<6} {len(ids3):<8} {len(ids4):<8} {'✓' if match else '✗ DIFF'}")
+        if not match:
+            print(f"       only in v3: {sorted(ids3 - ids4)[:10]}")
+            print(f"       only in v4: {sorted(ids4 - ids3)[:10]}")
+
+print()
+print("All folds match!" if all_match else "SPLITS DIFFER — fold assignments changed.")