pipekit: named arrays, one Signature, kernels written once

pipekit (core)
  Operator · Sequential · Graph · Signature (the only one)
        │
        ▼
pipekit-array
  array_namespace          (exists)
  NamedArray protocol      new: dims, data, isel, rewrap
  DimArray                 new: the plain named carrier (no georeference)
  dim-level functions      new: isel, reduce, stack, map_over, …
  KernelSpec · KernelOp    new: one kernel, every carrier
  adapters: raw arrays, xarray ([xarray] extra)
        │
        ▼
carrier libraries register an adapter
  geoarray: GeoArray · GeoSwath · GeoPoints (+ grid checks)
        │
        ▼
operator libraries
  geoops (geometry on geoarray carriers)
  xrops  (value analysis on xarray carriers)
  L1 kernels lifted with KernelOp, usable in both

class NamedArrayAdapter(Protocol):
    # ("band", "y", "x")
    def dims(self, x: Any) -> tuple[str, ...]: ...

    # the raw array: NumPy, JAX or dask
    def data(self, x: Any) -> Any: ...

    # static labels of one dim, or None (band names, level values)
    def labels(self, x: Any, dim: str) -> Any | None: ...

    # select by position on named dims: the adapter slices data and every
    # coordinate, label and geometry leaf on those dims together (band labels,
    # point positions, transform origins, per-patch transform rows)
    def isel(self, x: Any, indexers: dict[str, int | slice | Any]) -> Any: ...

    # a carrier of the same type around new data whose kept dims are unchanged
    # in length and order: untouched dims keep their metadata, dropped dims lose
    # it, new dims get none; never used for selection
    def rewrap(self, x: Any, data: Any, dims: tuple[str, ...]) -> Any: ...

    # raise if a and b cannot be combined element by element
    # (geoarray: same grid; xarray: exactly equal coordinates)
    def check_compatible(self, a: Any, b: Any) -> None: ...

    # named dims and dtype for summaries and checks
    def signature(self, x: Any) -> Signature: ...


# registration in pipekit core's registry, looked up along the carrier's MRO
xarray_adapter: NamedArrayAdapter = XarrayAdapter()
pipekit_ops.register_adapter(xr.DataArray, xarray_adapter)

@dataclass(frozen=True)
class Signature:
    # (("band", 2), ("y", 600), ("x", 600)); a size may be None
    dims: tuple[tuple[str, int | None], ...]
    dtype: str | None = None

    # new: does this signature fit a dims pattern?
    # "band y x" exact; "... y x" any leading dims
    def matches(self, pattern: str) -> bool: ...

    # new: a sizeless signature from a pattern, for annotations
    @classmethod
    def from_pattern(cls, pattern: str) -> "Signature": ...

one L1 kernel + "(c)->()"
        │
        ▼
KernelOp(kernel, "(c)->()", core={"c": "band"})
        │ dispatch on the carrier's adapter
        ├─ raw array      → call the kernel
        ├─ geoarray, NumPy → move band last, call, rewrap
        ├─ geoarray, JAX   → same, traced under jit
        └─ xarray          → apply_ufunc(input_core_dims=[["band"]],
                             dask="parallelized"), exact join

@dataclass(frozen=True)
class KernelSpec:
    # gufunc signature over abstract core dims: "(c)->()", "(),()->()", "(l),(l),(l),()->()"
    signature: str
    # backends the kernel is written for
    backends: tuple[str, ...] = ("numpy", "jax")
    # True if the kernel only handles core dims and must be vectorized over the rest
    vectorize: bool = False

class KernelOp(Operator):
    def __init__(
        self,
        kernel: Callable[..., Array],
        spec: KernelSpec | str,
        core: dict[str, str],
        geometry: dict[str, Callable[[Any], Array]] | None = None,
        **params: Any,
    ) -> None:
        # the L1 function
        self.kernel: Callable[..., Array] = kernel
        # the gufunc signature and backends
        self.spec: KernelSpec = (
            spec if isinstance(spec, KernelSpec) else KernelSpec(spec)
        )
        # abstract core dim → carrier dim, e.g. {"c": "band"}
        self.core: dict[str, str] = core
        # derived geometry: kernel argument → function of the carrier returning arrays,
        # e.g. {"grid": cgrid_from_carrier}; empty when the only geometry is a dim name
        self.geometry: dict[str, Callable[[Any], Array]] = geometry or {}
        # keyword parameters passed to every call
        self.params: dict[str, Any] = params

    def _apply(self, *carriers: Any) -> Any:
        # derived geometry from the first carrier, as arrays
        derived: dict[str, Array] = {
            name: fn(carriers[0]) for name, fn in self.geometry.items()
        }
        # every keyword input to the kernel
        kwargs: dict[str, Any] = self.params | derived
        # one dispatch for every carrier family, through its adapter
        return run_kernel(self.kernel, self.spec, self.core, carriers, kwargs)

    def compute_output_signature(self, sig: Signature) -> Signature:
        # input dims minus consumed core dims, plus produced core dims; no execution
        return apply_gufunc_signature(sig, self.spec, self.core)

# L1: element-wise, "(),()->()"
def normalized_difference(
    a: Float[Array, "*s"], b: Float[Array, "*s"], eps: float
) -> Float[Array, "*s"]:
    # a − b and a + b + ε
    num: Float[Array, "*s"] = a - b
    den: Float[Array, "*s"] = a + b + eps
    # ND = (a − b) / (a + b + ε)
    return num / den


# L1: consumes the band axis, "(c)->()"
def matched_filter_score(
    x: Float[Array, "... C"],
    mu: Float[Array, "C"],
    sigma: Float[Array, "C C"],
    t: Float[Array, "C"],
) -> Float[Array, "..."]:
    # r = x − μ
    r: Float[Array, "... C"] = einx.subtract("... c, c", x, mu)
    # w = Σ⁻¹t, solved, never inverted
    cov: lx.MatrixLinearOperator = lx.MatrixLinearOperator(
        sigma, lx.positive_semidefinite_tag
    )
    w: Float[Array, "C"] = lx.linear_solve(cov, t).value
    # rᵀw per pixel, and tᵀw
    num: Float[Array, "..."] = einx.dot("... c, c -> ...", r, w)
    norm: Float[Array, ""] = einx.dot("c, c ->", t, w)
    # α = (x − μ)ᵀΣ⁻¹t / (tᵀΣ⁻¹t)
    return num / norm


# wrapped once
nd: KernelOp = KernelOp(normalized_difference, "(),()->()", core={}, eps=1e-10)
score: KernelOp = KernelOp(
    matched_filter_score, "(c)->()", core={"c": "band"}, mu=mu, sigma=sigma, t=t
)

# a geoarray carrier, eager; scene: GeoArray["band y x"]
alpha_map: GeoArray["y x"] = score(scene)

# a lazy xarray cube, one task per chunk; cube: DataArray["time band y x"]
alpha_cube: DataArray["time y x"] = score(cube)

# under jit on a JAX carrier; μ, Σ and t become leaves (next section)
jscore: JaxOperator = pipekit_jax.to_jax(score)
score_jit: Callable = eqx.filter_jit(jscore)
alpha_jax: GeoArray["y x"] = score_jit(jscene)

# inside KernelOp, for xarray inputs (pipekit-array[xarray])
def _call_xarray(
    kernel: Callable,
    spec: KernelSpec,
    core: dict[str, str],
    kwargs: dict[str, Any],  # KernelOp's params | derived geometry, from _apply
    *xs: DataArray["..."],
) -> DataArray["..."]:
    # "(c)->()" with core = {"c": "band"}: band moves last and is consumed
    in_dims: list[list[str]] = [[core[d] for d in s] for s in spec.inputs]  # [["band"]]
    out_dims: list[list[str]] = [[core[d] for d in s] for s in spec.outputs]  # [[]]
    # coordinates must be equal; never inner-joined
    adapter.check_compatible(*xs)
    # apply_ufunc unwraps .data, calls kernel(*blocks, **kwargs) per chunk, and
    # rewraps coords and attrs; kwargs (eps; μ, Σ, t; a grid) are passed unmapped
    return xr.apply_ufunc(
        kernel,
        *xs,
        input_core_dims=in_dims,
        output_core_dims=out_dims,
        kwargs=kwargs,
        vectorize=spec.vectorize,
        dask="parallelized",
        join="exact",
    )

# NumPy and xarray: a plain Operator
score: KernelOp = KernelOp(
    matched_filter_score, "(c)->()", core={"c": "band"}, mu=mu, sigma=sigma, t=t
)

# JAX: the same kernel and spec as a JaxOperator (an eqx.Module);
# array parameters are leaves, spec and dim names are static
jscore: JaxOperator = pipekit_jax.to_jax(score)

# compile with Equinox's filter, since the module holds arrays; jscene: GeoArray["band y x"]
score_jit: Callable = eqx.filter_jit(jscore)
alpha: GeoArray["y x"] = score_jit(jscene)

# gradients flow to the parameters, e.g. a learned target signature: ∂ℒ/∂t
loss_grad: Callable = eqx.filter_grad(loss)
dloss: JaxOperator = loss_grad(jscore, jscene)

# estimate μ and Σ from the scene's background pixels, eagerly, in pre-processing
mf: MatchedFilter = MatchedFilter(target=t, dim="band")
fitted: MatchedFilter = mf.fit(scene)  # scene: GeoArray["band y x"]

# apply: a KernelOp underneath, any carrier
alpha_map: GeoArray["y x"] = fitted(scene)