decoding: decoder graph (DAG) generalizing the linear steps pipeline

Add a decoder-graph execution model to connectomics/decoding/: named nodes with
typed inputs (raw prediction channels and/or upstream node outputs), multi-parent
combine nodes, and one declared saved output. The linear decoding.steps chain is
lowered to a chain graph and run by the same run_decode_graph executor (single code
path; byte-parity by construction). Adds channel_gate and combine_split (background-
preserving coarsest common refinement) ops and a strict decoding.graph schema with
graph-xor-steps validation. Built via CCC (Claude plan / Codex code, Claude review);
35 tests pass.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: Donglai Wei

connectomics/config/schema/decoding.py

@@ -56,6 +56,25 @@ class DecodeModeConfig:
     kwargs: Dict[str, Any] = field(default_factory=dict)
 
 
+@dataclass
+class GraphNodeConfig:
+    """Single node in a decoder graph."""
+
+    enabled: bool = True
+    name: str = ""
+    op: str = ""
+    inputs: List[str] = field(default_factory=list)
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class DecodeGraphConfig:
+    """Decoder graph configuration."""
+
+    nodes: List[GraphNodeConfig] = field(default_factory=list)
+    output: str = ""
+
+
 @dataclass
 class TuningParameterConfig:
     """Single tunable decoding/postprocessing parameter."""
@@ -154,6 +173,7 @@ class DecodingConfig:
     save_path: str = ""
     # Optional user-controlled filename suffix appended to decoded outputs.
     save_suffix: str = ""
+    graph: Optional[DecodeGraphConfig] = None
     steps: List[DecodeModeConfig] = field(default_factory=list)
     postprocessing: PostprocessingConfig = field(default_factory=PostprocessingConfig)
     # Optional explicit raw-prediction file (.h5). If set, pipeline loads
@@ -162,3 +182,7 @@ class DecodingConfig:
     affinity_mask_path: str = ""
     affinity_qc: AffinityQCConfig = field(default_factory=AffinityQCConfig)
     tuning: Optional[DecodingTuningConfig] = None
+
+    def __post_init__(self) -> None:
+        if self.graph is not None and self.steps:
+            raise ValueError("decoding.graph cannot be set together with non-empty decoding.steps.")

connectomics/decoding/base.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Any, Dict, Protocol
+from typing import Any, Dict, List, Protocol
 
 import numpy as np
 
@@ -14,10 +14,35 @@ class DecodeFunction(Protocol):
     def __call__(self, predictions: np.ndarray, **kwargs: Any) -> np.ndarray: ...
 
 
+class GraphOp(Protocol):
+    """Callable decoder-graph operation signature used by the registry."""
+
+    def __call__(self, inputs: List[np.ndarray], **kwargs: Any) -> np.ndarray: ...
+
+
 @dataclass
 class DecodeStep:
     """Single step in a decoding pipeline."""
 
     enabled: bool = True
     name: str = ""
     kwargs: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class DecodeNode:
+    """Single named node in a decoding graph."""
+
+    enabled: bool = True
+    name: str = ""
+    op: str = ""
+    inputs: List[str] = field(default_factory=list)
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class DecodeGraph:
+    """Decoder graph with one declared output node."""
+
+    nodes: List[DecodeNode] = field(default_factory=list)
+    output: str = ""

connectomics/decoding/decoders/combine.py

@@ -0,0 +1,81 @@
+"""Multi-input decoder-graph combine operations."""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import numpy as np
+
+
+def _validate_label_input(arr: np.ndarray, *, name: str) -> None:
+    if not np.issubdtype(arr.dtype, np.integer):
+        raise TypeError(f"combine_split {name} must be an integer label array.")
+    if np.issubdtype(arr.dtype, np.signedinteger) and arr.size and int(arr.min()) < 0:
+        raise ValueError(f"combine_split {name} must not contain negative labels.")
+
+
+def _check_key_space(max_a: int, max_b: int) -> int:
+    max_uint64 = int(np.iinfo(np.uint64).max)
+    if max_b >= max_uint64:
+        raise OverflowError("combine_split pair-key base exceeds uint64 range.")
+    base = max_b + 1
+    if max_a > (max_uint64 - max_b) // base:
+        raise OverflowError("combine_split pair keys would overflow uint64.")
+    return base
+
+
+def combine_split(
+    inputs: Sequence[np.ndarray],
+    *,
+    output_dtype: str | np.dtype = "uint32",
+) -> np.ndarray:
+    """Return the background-preserving coarsest common refinement of two labels."""
+    if len(inputs) != 2:
+        raise ValueError(f"combine_split expects exactly two inputs, got {len(inputs)}.")
+
+    a = np.asarray(inputs[0])
+    b = np.asarray(inputs[1])
+    if a.shape != b.shape:
+        raise ValueError(
+            f"combine_split inputs must have matching shapes, got {a.shape} and {b.shape}."
+        )
+    _validate_label_input(a, name="input 0")
+    _validate_label_input(b, name="input 1")
+
+    dtype = np.dtype(output_dtype)
+    if not np.issubdtype(dtype, np.integer):
+        raise TypeError(f"combine_split output_dtype must be an integer dtype, got {dtype}.")
+
+    out = np.zeros(a.shape, dtype=dtype)
+    fg = (a != 0) & (b != 0)
+    if not bool(fg.any()):
+        return out
+
+    a_fg = a[fg]
+    b_fg = b[fg]
+    max_a = int(a_fg.max())
+    max_b = int(b_fg.max())
+    base = _check_key_space(max_a, max_b)
+
+    key = a_fg.astype(np.uint64, copy=False)
+    if key is a_fg:
+        key = key.copy()
+    key *= np.uint64(base)
+    np.add(key, b_fg.astype(np.uint64, copy=False), out=key)
+
+    _, inv = np.unique(key, return_inverse=True)
+    n_labels = int(inv.max()) + 1 if inv.size else 0
+    if n_labels >= 2**32:
+        raise OverflowError("combine_split produced too many labels for uint32 output.")
+    dtype_info = np.iinfo(dtype)
+    if n_labels > int(dtype_info.max):
+        raise OverflowError(
+            f"combine_split produced {n_labels} labels, exceeding output dtype {dtype}."
+        )
+
+    labels = np.arange(1, n_labels + 1, dtype=dtype)
+    out[fg] = labels[inv]
+    return out
+
+
+__all__ = ["combine_split"]

connectomics/decoding/decoders/transforms.py

@@ -0,0 +1,40 @@
+"""Pure array transforms for decoder graphs."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from ...utils.channel_slices import resolve_channel_indices
+
+
+def channel_gate(
+    predictions: np.ndarray,
+    *,
+    signal_channels,
+    gate_channel,
+) -> np.ndarray:
+    """Multiply selected signal channels by a single gate channel."""
+    arr = np.asarray(predictions)
+    if arr.ndim < 1:
+        raise ValueError("channel_gate expects an array with a channel axis.")
+
+    signal_indices = resolve_channel_indices(
+        signal_channels,
+        num_channels=int(arr.shape[0]),
+        context="channel_gate.signal_channels",
+    )
+    gate_indices = resolve_channel_indices(
+        gate_channel,
+        num_channels=int(arr.shape[0]),
+        context="channel_gate.gate_channel",
+    )
+    if len(gate_indices) != 1:
+        raise ValueError(
+            f"channel_gate.gate_channel must resolve to one channel, got {gate_indices}."
+        )
+
+    gated = arr[signal_indices] * arr[gate_indices[0] : gate_indices[0] + 1]
+    return gated.astype(arr.dtype, copy=False)
+
+
+__all__ = ["channel_gate"]