affinity_mask: spec-only fast path; legacy load is fallback

_maybe_apply_affinity_mask used to read the full (X, Y, Z) uint8 mask
volume, allocate a same-shape boolean zero_idx, and zero each prediction
channel via fancy indexing. On large volumes (12+ G voxels here) that
dominated decoding runtime.

build_affinity_mask already writes the full mask spec (low_z, high_z,
border_width, bg_thresh, axis_order, source_image) to the dataset attrs,
so we can reconstruct the mask implicitly:

- Z-gate: two contiguous slice writes per channel-axis end.
- Border + intensity: only the outer XY ring of width border_width is
  read from the source zarr; voxels with intensity <= bg_thresh are
  zeroed across all channels. Corner overlap removed by slicing the Y
  strips between [bw, X-bw] on the X axis.

apply_affinity_mask_from_spec(predictions, spec) is the new in-place
helper, exported from connectomics.decoding.qc. Stage code first tries
_read_affinity_mask_spec(mask_path); when all five expected attrs are
present and axis_order=="XYZ" it takes the fast path, otherwise it
falls back to the legacy full-volume load. Output log prefixes the
legacy path with "(legacy load)" to make the active path obvious.

Verified parity on a synthetic (12, 14, 8) fixture: spec-only apply
produces byte-identical predictions to the legacy load (n_zero matches,
np.array_equal == True). All 40 existing affinity/QC unit tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: Donglai Wei

connectomics/decoding/qc/__init__.py

@@ -4,6 +4,7 @@
     AffinityQCAccumulator,
     AffinityQCParams,
     AffinityQCReport,
+    apply_affinity_mask_from_spec,
     begin_streaming_qc,
     build_affinity_mask,
     finish_streaming_qc,
@@ -16,6 +17,7 @@
     "AffinityQCAccumulator",
     "AffinityQCParams",
     "AffinityQCReport",
+    "apply_affinity_mask_from_spec",
     "begin_streaming_qc",
     "build_affinity_mask",
     "finish_streaming_qc",

connectomics/decoding/qc/affinity.py

@@ -506,6 +506,92 @@ def read_block(z0: int, z1: int) -> np.ndarray:
     )
 
 
+def apply_affinity_mask_from_spec(predictions: np.ndarray, spec: dict) -> tuple[int, int]:
+    """Apply an affinity mask in-place using only the spec, skipping the full mask load.
+
+    Equivalent to building a mask via :func:`build_affinity_mask` with the same
+    parameters and then zeroing every prediction voxel where the mask is 0, but
+    without materializing the (X, Y, Z) mask array. ``predictions`` is modified
+    in place; the (zero, total) prediction-voxel counts are returned for logging.
+
+    ``spec`` requires keys ``low_z``, ``high_z``, ``border_width``, ``bg_thresh``,
+    ``image_path`` (a zarr 3D/4D source for the border-intensity check).
+    ``axis_order`` is optional and must be ``"XYZ"`` if present.
+
+    The two rules from :func:`build_affinity_mask` are applied:
+
+    - Z-gate: contiguous slice writes zero everything outside ``[low_z, high_z)``.
+    - Border + intensity: only the outer XY ring of width ``border_width`` is
+      read from the source image; voxels with intensity <= ``bg_thresh`` are
+      zeroed across all channels. Corners are zeroed twice (idempotent).
+    """
+    if predictions.ndim != 4:
+        raise ValueError(
+            f"apply_affinity_mask_from_spec expects 4D predictions, got {predictions.shape}"
+        )
+    axis_order = str(spec.get("axis_order", "XYZ"))
+    if axis_order != "XYZ":
+        raise ValueError(f"unsupported axis_order={axis_order!r}; only 'XYZ' is supported")
+
+    _C, X, Y, Z = predictions.shape
+    low_z = int(spec["low_z"])
+    high_z = int(spec["high_z"])
+    border_width = int(spec["border_width"])
+    bg_thresh = int(spec["bg_thresh"])
+    image_path = spec.get("image_path") or spec.get("source_image") or ""
+    if not (0 <= low_z <= high_z <= Z):
+        raise ValueError(f"invalid z range [{low_z}, {high_z}) for Z={Z}")
+
+    n_total_pred = int(predictions.size)
+    n_zero_voxels = 0  # counts in mask-volume coords (X*Y*Z), not predictions
+
+    # Z-gate: contiguous slice writes.
+    if low_z > 0:
+        predictions[:, :, :, :low_z] = 0
+        n_zero_voxels += X * Y * low_z
+    if high_z < Z:
+        predictions[:, :, :, high_z:] = 0
+        n_zero_voxels += X * Y * (Z - high_z)
+
+    # Border + intensity.
+    if border_width > 0 and bg_thresh > 0 and image_path and high_z > low_z:
+        bw = border_width
+        if bw * 2 > X or bw * 2 > Y:
+            raise ValueError(
+                f"border_width={bw} too large for shape (X={X}, Y={Y})"
+            )
+        import zarr
+
+        img = zarr.open(image_path, mode="r")
+        if img.ndim == 4:
+            def read_strip(xs: slice, ys: slice) -> np.ndarray:
+                return np.asarray(img[xs, ys, low_z:high_z, 0])
+        elif img.ndim == 3:
+            def read_strip(xs: slice, ys: slice) -> np.ndarray:
+                return np.asarray(img[xs, ys, low_z:high_z])
+        else:
+            raise ValueError(f"image must be 3D or 4D, got shape {img.shape}")
+
+        strips = [
+            (slice(0, bw), slice(None, None)),                       # left X edge
+            (slice(X - bw, X), slice(None, None)),                   # right X edge
+            (slice(bw, X - bw), slice(0, bw)),                       # left Y edge (no corner dup)
+            (slice(bw, X - bw), slice(Y - bw, Y)),                   # right Y edge (no corner dup)
+        ]
+        for sx, sy in strips:
+            img_strip = read_strip(sx, sy)
+            zero_strip = img_strip <= bg_thresh
+            if not zero_strip.any():
+                continue
+            # pred_view is a basic-slice view of predictions, so the fancy-index
+            # assignment below writes through to predictions.
+            pred_view = predictions[:, sx, sy, low_z:high_z]
+            pred_view[:, zero_strip] = 0
+            n_zero_voxels += int(zero_strip.sum())
+
+    return n_zero_voxels, X * Y * Z
+
+
 def _cfg_get(cfg: Any, name: str, default: Any = None) -> Any:
     if cfg is None:
         return default

connectomics/decoding/stage.py

@@ -116,13 +116,51 @@ def apply_decoding_postprocessing(cfg: Any, data: np.ndarray) -> np.ndarray:
     return output
 
 
+_AFFINITY_MASK_SPEC_ATTRS = ("low_z", "high_z", "border_width", "bg_thresh", "source_image")
+
+
+def _read_affinity_mask_spec(mask_path: str) -> dict | None:
+    """Return the affinity-mask spec from the h5 file's attrs, or None.
+
+    Files produced by :func:`connectomics.decoding.qc.build_affinity_mask` carry
+    the mask spec on the ``main`` dataset's ``attrs``. When all required attrs
+    are present (and ``axis_order`` is the expected XYZ), the spec is enough to
+    apply the mask without reading the full mask volume.
+    """
+    import h5py
+
+    with h5py.File(mask_path, "r") as f:
+        if "main" not in f:
+            return None
+        attrs = dict(f["main"].attrs)
+    if not all(k in attrs for k in _AFFINITY_MASK_SPEC_ATTRS):
+        return None
+    axis_order = str(attrs.get("axis_order", "XYZ"))
+    if axis_order != "XYZ":
+        return None
+    return {
+        "low_z": int(attrs["low_z"]),
+        "high_z": int(attrs["high_z"]),
+        "border_width": int(attrs["border_width"]),
+        "bg_thresh": int(attrs["bg_thresh"]),
+        "image_path": str(attrs["source_image"]),
+        "axis_order": axis_order,
+    }
+
+
 def _maybe_apply_affinity_mask(cfg: Any, predictions: np.ndarray) -> np.ndarray:
     """Zero affinity channels at masked-out voxels per ``decoding.affinity_mask_path``.
 
-    The mask file must contain a single 3D ``uint8`` dataset matching the spatial
-    shape of ``predictions[0]`` (the affinity volume). 0 = drop, 1 = keep. Mask is
-    broadcast across the channel dimension; voxels where ``mask==0`` get zeroed
-    so that connected-components decoders sever all outgoing edges there.
+    Two code paths:
+
+    - **Spec fast path** (preferred): if the mask h5 was produced by
+      :func:`build_affinity_mask`, its dataset ``attrs`` carry the spec
+      (``low_z``, ``high_z``, ``border_width``, ``bg_thresh``, ``source_image``).
+      We apply the rules directly via slice writes + a tiny border-strip read
+      from the source image, skipping the multi-GB mask load entirely.
+    - **Legacy slow path**: if the spec attrs are missing (older files or
+      external producers), fall back to reading the full ``(X, Y, Z) uint8``
+      mask and zeroing per channel via boolean fancy indexing.
     """
     decoding_cfg = _cfg_get(cfg, "decoding", None)
     mask_path = _cfg_get(decoding_cfg, "affinity_mask_path", "") or ""
@@ -133,6 +171,28 @@ def _maybe_apply_affinity_mask(cfg: Any, predictions: np.ndarray) -> np.ndarray:
             "affinity_mask_path requires 4D predictions (C, *spatial); "
             f"got shape {predictions.shape}"
         )
+
+    spec = _read_affinity_mask_spec(mask_path)
+    if spec is not None:
+        from .qc.affinity import apply_affinity_mask_from_spec
+
+        if spec["high_z"] - spec["low_z"] > predictions.shape[-1]:
+            raise ValueError(
+                f"affinity_mask spec z range exceeds prediction z extent: "
+                f"[{spec['low_z']}, {spec['high_z']}) vs Z={predictions.shape[-1]}"
+            )
+        t0 = time.time()
+        n_zero, n_total = apply_affinity_mask_from_spec(predictions, spec)
+        logger.info(
+            "Applied affinity mask via spec %s in %.2fs: z=[%d,%d) "
+            "border_width=%d bg_thresh=%d zeroed=%d/%d (%.2f%%) across %d channels",
+            mask_path, time.time() - t0,
+            spec["low_z"], spec["high_z"], spec["border_width"], spec["bg_thresh"],
+            n_zero, n_total, 100.0 * n_zero / max(n_total, 1),
+            predictions.shape[0],
+        )
+        return predictions
+
     import h5py
 
     with h5py.File(mask_path, "r") as f:
@@ -155,7 +215,8 @@ def _maybe_apply_affinity_mask(cfg: Any, predictions: np.ndarray) -> np.ndarray:
     n_zero = int(zero_idx.sum())
     n_total = int(zero_idx.size)
     logger.info(
-        "Applying affinity mask %s: zeroing %d/%d voxels (%.2f%%) across %d channels",
+        "Applying affinity mask (legacy load) %s: zeroing %d/%d voxels "
+        "(%.2f%%) across %d channels",
         mask_path,
         n_zero,
         n_total,