Add BANIS affinity reproduction and chunked inference

Author: Donglai Wei

.claude/benchmark/SNEMI.md

@@ -158,14 +158,14 @@ No per-tutorial loss config needed — the pipeline profile handles it.
 
 | | DeepEM | PyTC (before fix) | PyTC (after fix) |
 |---|---|---|---|
-| **Border handling** | Per-channel `get_pair` crop + mask | Uniform `deepem_crop` (max-offset spatial crop) | Per-channel valid mask ✅ |
+| **Border handling** | Per-channel `get_pair` crop + mask | Old uniform max-offset spatial crop | `affinity_mode=deepem` per-channel valid mask ✅ |
 | **Augment padding on labels** | Mask propagated through augmentation | `RandAffined` reflection padding on labels → false affinities | Per-channel mask excludes border artifacts ✅ |
 
 **Problem:** Two interacting issues caused border artifacts in affinity targets:
 
 1. **Reflection padding on labels during augmentation**: `RandAffined` and `RandElasticd` use `padding_mode="reflection"` for all keys including labels. When spatial transforms rotate/scale/shear a patch, border pixels are filled with reflected label values. Computing affinity from these reflected labels creates false affinities — especially visible for long-range channels like ch11 (offset `0-27-0`) where the reflected region spans 27 voxels.
 
-2. **Uniform spatial crop vs per-channel masking**: The old `deepem_crop` computed the **union** of all offsets' invalid borders and uniformly cropped all channels to this smallest valid region. For the SNEMI 12-channel offsets, this meant cropping (4, 27, 27) from all channels — even short-range channels that only need 1 voxel cropped. This wasted ~35% of training data for short-range channels.
+2. **Uniform spatial crop vs per-channel masking**: The old crop path computed the **union** of all offsets' invalid borders and uniformly cropped all channels to this smallest valid region. For the SNEMI 12-channel offsets, this meant cropping (4, 27, 27) from all channels — even short-range channels that only need 1 voxel cropped. This wasted ~35% of training data for short-range channels.
 
 **DeepEM's approach**: In DeepEM, `get_pair(arr, edge)` extracts two aligned crops per channel, computing affinity only in the overlap region. A separate mask (propagated through augmentation) excludes padded regions from the loss. Each channel has its own valid region.
 

.claude/refactor/affinity.md

@@ -0,0 +1,68 @@
+# Affinity Modes
+
+Affinity targets must declare one explicit convention:
+
+```yaml
+data:
+  label_transform:
+    targets:
+      - name: affinity
+        kwargs:
+          offsets: ["0-0-1", "0-1-0", "1-0-0"]
+          affinity_mode: deepem  # or banis
+```
+
+There is no legacy crop flag. The mode controls target voxel placement,
+valid-border masking, visualization crop, and test-time prediction crop.
+
+## Modes
+
+| Mode | Edge storage | Valid side for positive offsets | Intended use |
+| --- | --- | --- | --- |
+| `deepem` | Destination voxel `v + offset` | Leading border invalid | DeepEM/SNEMI-style targets, zwatershed/ABISS destination-index affinity |
+| `banis` | Source voxel `v` | Trailing border invalid | BANIS-compatible targets and source-index connected-component decoding |
+
+For a positive x offset `0-0-1`, `deepem` produces valid affinities at
+`x >= 1`; `banis` produces valid affinities at `x < W - 1`.
+
+## Training
+
+Training does not crop every affinity channel to the largest common valid
+interior. It keeps prediction and target shapes unchanged and applies a
+per-channel valid mask before loss evaluation. This preserves short-range edge
+supervision while excluding convention-dependent padded borders.
+
+Mixed affinity modes in one stacked label tensor are rejected. If multiple
+affinity target groups are ever needed, they must share the same
+`affinity_mode`.
+
+## Inference And Decoding
+
+Test-time affinity crops use the same mode as training:
+
+```python
+compute_affinity_crop_pad(offsets, affinity_mode="deepem")
+compute_affinity_crop_pad(offsets, affinity_mode="banis")
+```
+
+The crop is resolved after `inference.select_channel` and output-head target
+slices. If decoding keeps only short-range channels from a larger affinity
+target stack, the automatic affinity crop must use only those selected offsets.
+
+`decode_affinity_cc` has an independent `edge_offset` knob for the numba
+backend:
+
+| Target mode | `decode_affinity_cc.kwargs.edge_offset` |
+| --- | --- |
+| `deepem` | `1` |
+| `banis` | `0` |
+
+The `cc3d` backend ignores directed edge placement and only thresholds
+foreground connectivity, so this matters mainly for `backend: numba`.
+
+## Config Policy
+
+Use `affinity_mode: deepem` for DeepEM/SNEMI/LiConn-style configs.
+
+Use `affinity_mode: banis` for BANIS/NISB reproduction configs and any config
+whose target should match `lib/banis/data.py::comp_affinities`.

.claude/refactor/training.md

@@ -63,7 +63,7 @@ All 12 issues across 4 priority levels have been resolved. 190/191 unit tests pa
 #### P3.2: Affinity decoupling in orchestrator.py
 - **File**: `training/loss/orchestrator.py`
 - **Issue**: Direct import of `data.process.affinity` created tight cross-package coupling
-- **Fix**: Dependency injection via constructor parameters (`affinity_crop_enabled_fn`, `crop_spatial_fn`, `resolve_affinity_offsets_fn`) with lazy-import bridge functions as defaults. No behavioral change
+- **Fix**: Dependency injection via constructor parameters (`resolve_affinity_mode_fn`, `resolve_affinity_offsets_fn`) with lazy-import bridge functions as defaults. Affinity target handling now routes through explicit `affinity_mode`.
 
 #### P3.3: Logging migration (print -> logging)
 - **Files**: All files in `training/` except `debugging.py`

.claude/reference/snemi_old.md

@@ -0,0 +1,65 @@
+# snemi_old: Segmentation Post-Processing Reference
+
+Summary of useful functions from `lib/snemi_old/*.py` for improving segmentation.
+
+## Key Post-Processing Strategies
+
+### 1. Segment Classification (T_pytc_v2.py, T_snemi220416.py)
+- **Border-touching**: `bb[:,1::2] == 0` or `bb[:,2] == D-1` etc. Segments touching volume boundary are unreliable for merge analysis.
+- **Interior segments**: `num_border == 0` — candidates for orphan merge.
+- **Singletons**: `bb[:,1] == bb[:,2]` — single-slice segments. SNEMI test had 1059/1651 singletons. High error rate.
+- **Disconnected components**: cc3d.connected_components per segment — remove non-largest component.
+
+### 2. Orphan Detection & Merge (T_snemi220416.py opt=='0.32', T_pytc_v2.py opt=='2.22')
+Criteria:
+1. Segment touches ≤1 boundary
+2. Not connected across z-slices
+3. Has single dominant neighbor in z±1
+4. Size-based IoU > 0.6 with neighbor
+
+### 3. Oracle Merge Analysis (T_pytc_v2.py opt=='2.211')
+- Map each predicted segment to best GT match via max IoU
+- Group predicted segments by GT label
+- Segments mapped to same GT = should be merged
+- Typical result: 190 oracle merges, ARE 0.048 → 0.025
+
+### 4. Morphological Refinement (T_pytc.py)
+- `seg_postprocess()`: 2D constrained watershed (mahotas.cwatershed) per slice
+- Optional Sobel edge guidance from raw image
+- ~0.008-0.015 error reduction
+
+### 5. Multi-Stage Waterz (T_snemi220416.py, T_waterz.py)
+Best parameters found:
+- `merge_function: aff85_his256`
+- `aff_threshold: [0.1, 0.9]`
+- `threshold: 0.4-0.7`
+- `dust_merge_size: 800 * rr²` (resolution-dependent)
+- `dust_merge_affinity: 0.3-0.5`
+
+### 6. Consistency Checking (T_consistency.py)
+- Track segment IDs across z-slices
+- Count max consecutive occurrences
+- Segments with ≤2 consecutive slices = likely noise
+- Abrupt size changes = potential errors
+
+### 7. Skeleton Analysis (T_yulun_skel.py, T_skel.py)
+- kimimaro TEASAR: `scale=4, const=500, anisotropy=(30,6,6)`
+- Cable length filtering: long axons (≥5000µm) vs short fragments (<1000µm)
+- ERL (skeleton-based) metric as alternative to pixel-based ARE
+- Oracle skeletonization bridges false splits
+
+## Practical Improvement Hierarchy
+
+1. **Remove single-slice dust** — low risk, removes noise
+2. **cc3d disconnect removal** — keep largest component per segment
+3. **Orphan merge** — segments with bbox fully inside another
+4. **IoU-based cross-slice merge** — cautious, only at bbox endpoints
+5. **Morphological refinement** — cwatershed per slice
+6. **Skeleton-guided merge** — use cable length to validate merges
+
+## Key Files
+- `T_pytc_v2.py` — comprehensive pipeline (merge, split, oracle)
+- `T_snemi220416.py` — waterz params, orphan detection
+- `T_consistency.py` — cross-slice tracking
+- `T_yulun_iou.py` — IoU computation, adapted_rand
+- `T_yulun_skel.py` — skeleton analysis

connectomics/config/pipeline/config_io.py

@@ -20,7 +20,7 @@
 from ...data.processing.build import count_stacked_label_transform_channels
 from ...models.architectures.registry import get_architecture_info
 from ...utils.channel_slices import infer_min_required_channels
-from ...utils.model_outputs import resolve_configured_output_head
+from ...utils.model_outputs import resolve_configured_output_head, resolve_output_heads
 from ..schema import Config
 from ..schema.root import MergeContext
 from .profile_engine import _YAML_PROFILE_ENGINE
@@ -329,11 +329,19 @@ def validate_config(cfg: Config) -> None:
                 f"model.primary_head='{primary_head}' is not present in model.heads "
                 f"({sorted(model_heads.keys())})."
             )
-        if inference_head is not None and inference_head not in model_heads:
-            raise ValueError(
-                f"inference.head='{inference_head}' is not present in model.heads "
-                f"({sorted(model_heads.keys())})."
+        if inference_head is not None:
+            # Accept comma-separated lists (merged-heads inference); each name must exist.
+            inference_head_names = (
+                [h.strip() for h in inference_head.split(",") if h.strip()]
+                if isinstance(inference_head, str) and "," in inference_head
+                else [inference_head]
             )
+            missing = [h for h in inference_head_names if h not in model_heads]
+            if missing:
+                raise ValueError(
+                    f"inference.head={inference_head_names} references unknown heads {missing}; "
+                    f"available: {sorted(model_heads.keys())}."
+                )
         if (
             visualization_head is not None
             and visualization_head != "all"
@@ -365,6 +373,33 @@ def validate_config(cfg: Config) -> None:
     if cfg.data.dataloader.batch_size <= 0:
         raise ValueError("data.dataloader.batch_size must be positive")
 
+    strategy = str(getattr(cfg.inference, "strategy", "whole_volume")).lower()
+    if strategy not in {"whole_volume", "chunked"}:
+        raise ValueError("inference.strategy must be 'whole_volume' or 'chunked'")
+    chunking_cfg = getattr(cfg.inference, "chunking", None)
+    chunking_enabled = bool(getattr(chunking_cfg, "enabled", False)) or strategy == "chunked"
+    if chunking_enabled:
+        if len(cfg.data.dataloader.patch_size) != 3:
+            raise ValueError("inference.chunking requires 3D data.dataloader.patch_size")
+        axes = str(getattr(chunking_cfg, "axes", "all")).lower()
+        if axes not in {"all", "z"}:
+            raise ValueError("inference.chunking.axes must be 'all' or 'z'")
+        chunk_size = getattr(chunking_cfg, "chunk_size", None)
+        if not chunk_size or len(chunk_size) != 3:
+            raise ValueError("inference.chunking.chunk_size must be a length-3 ZYX list")
+        if any(int(v) <= 0 for v in chunk_size):
+            raise ValueError("inference.chunking.chunk_size values must be positive")
+        halo = getattr(chunking_cfg, "halo", None)
+        if halo is None or len(halo) != 3:
+            raise ValueError("inference.chunking.halo must be a length-3 ZYX list")
+        if any(int(v) < 0 for v in halo):
+            raise ValueError("inference.chunking.halo values must be non-negative")
+        stitching = getattr(chunking_cfg, "stitching", None)
+        if stitching is not None:
+            min_contact = int(getattr(stitching, "min_contact", 1))
+            if min_contact <= 0:
+                raise ValueError("inference.chunking.stitching.min_contact must be positive")
+
     # Optimizer validation
     if cfg.optimization.optimizer.lr <= 0:
         raise ValueError("optimization.optimizer.lr must be positive")
@@ -583,22 +618,26 @@ def _validate_label_channel_capacity(selector_value: Any, *, path: str) -> None:
                 break
 
         if model_heads and decode_has_channel_selection:
-            decode_output_head = resolve_configured_output_head(
-                cfg,
-                purpose="decode channel selection",
-                allow_none=True,
-            )
-            if len(model_heads) > 1 and decode_output_head is None:
+            decode_heads = resolve_output_heads(cfg, purpose="decode channel selection")
+            if len(model_heads) > 1 and not decode_heads:
                 raise ValueError(
                     "Cross-section validation failed: decode channel selectors require "
                     "inference.head or model.primary_head when model.heads has multiple "
                     f"entries ({sorted(model_heads.keys())})."
                 )
-            if decode_output_head in model_heads:
-                decode_available_channels = int(
-                    getattr(model_heads[decode_output_head], "out_channels", out_channels)
+            if len(decode_heads) > 1:
+                decode_available_channels = sum(
+                    int(getattr(model_heads[h], "out_channels", 0)) for h in decode_heads
                 )
-                decode_channel_scope = f"head '{decode_output_head}'"
+                decode_channel_scope = f"merged heads {decode_heads}"
+                decode_output_head = decode_heads[0]
+            elif decode_heads:
+                decode_output_head = decode_heads[0]
+                if decode_output_head in model_heads:
+                    decode_available_channels = int(
+                        getattr(model_heads[decode_output_head], "out_channels", out_channels)
+                    )
+                    decode_channel_scope = f"head '{decode_output_head}'"
 
         for i, decode_step in enumerate(decoding_cfg):
             kwargs = getattr(decode_step, "kwargs", None)
@@ -624,39 +663,35 @@ def _validate_label_channel_capacity(selector_value: Any, *, path: str) -> None:
                         continue
                     required_output_channels.append((path, min_channels))
 
-    # 2e) TTA channel selectors
+    # 2e) Inference channel selectors
     tta_cfg = getattr(cfg.inference, "test_time_augmentation", None)
     channel_activations = getattr(tta_cfg, "channel_activations", None) if tta_cfg else None
-    select_channel = getattr(tta_cfg, "select_channel", None) if tta_cfg else None
-    tta_has_channel_selection = bool(channel_activations) or select_channel is not None
-    tta_output_head = (
-        resolve_configured_output_head(
-            cfg,
-            purpose="TTA channel selection",
-            allow_none=True,
-        )
-        if model_heads
-        else None
+    select_channel = getattr(cfg.inference, "select_channel", None)
+    inference_has_channel_selection = bool(channel_activations) or select_channel is not None
+    tta_heads = (
+        resolve_output_heads(cfg, purpose="inference channel selection") if model_heads else []
     )
-    if (
-        model_heads
-        and len(model_heads) > 1
-        and tta_has_channel_selection
-        and tta_output_head is None
-    ):
+    tta_output_head = tta_heads[0] if tta_heads else None
+    if model_heads and len(model_heads) > 1 and inference_has_channel_selection and not tta_heads:
         raise ValueError(
-            "Cross-section validation failed: TTA channel selectors require inference.head "
+            "Cross-section validation failed: inference channel selectors require inference.head "
             "or model.primary_head when model.heads has multiple entries "
             f"({sorted(model_heads.keys())})."
         )
-    tta_available_channels = (
-        int(getattr(model_heads[tta_output_head], "out_channels", out_channels))
-        if tta_output_head in model_heads
-        else out_channels
-    )
-    tta_channel_scope = (
-        f"head '{tta_output_head}'" if tta_output_head in model_heads else "model output"
-    )
+    if len(tta_heads) > 1:
+        tta_available_channels = sum(
+            int(getattr(model_heads[h], "out_channels", 0)) for h in tta_heads
+        )
+        tta_channel_scope = f"merged heads {tta_heads}"
+    else:
+        tta_available_channels = (
+            int(getattr(model_heads[tta_output_head], "out_channels", out_channels))
+            if tta_output_head in model_heads
+            else out_channels
+        )
+        tta_channel_scope = (
+            f"head '{tta_output_head}'" if tta_output_head in model_heads else "model output"
+        )
 
     def _validate_tta_channel_capacity(selector_value: Any, *, path: str) -> None:
         min_selector_channels = infer_min_required_channels(
@@ -692,7 +727,7 @@ def _validate_tta_channel_capacity(selector_value: Any, *, path: str) -> None:
             )
     _validate_tta_channel_capacity(
         select_channel,
-        path="inference.test_time_augmentation.select_channel",
+        path="inference.select_channel",
     )
 
     if required_output_channels: