feat: stalign with a fit apply pattern. Also foundation for imputation

docs/api.md

@@ -149,6 +149,9 @@ See the {doc}`extensibility guide </extensibility>` for how to implement a custo
 
     experimental.tl.calculate_tiling_qc
     experimental.tl.TilingQCParams
+    experimental.tl.align
+    experimental.tl.align_by_landmarks
+    experimental.tl.AlignResult
     experimental.pl.tiling_qc
     experimental.im.fit_stain_reference
     experimental.im.apply_stain_normalization

hatch.toml

@@ -17,10 +17,14 @@ extra-dependencies = ["diff-cover"]
 matrix = [
   { deps = ["stable"], python = ["3.11", "3.12", "3.13"] },
   { deps = ["pre"], python = ["3.13"] },
+  { deps = ["stable"], python = ["3.13"], extras = ["jax"] },
 ]
 overrides.matrix.deps.env-vars = [
   { key = "UV_PRERELEASE", value = "allow", if = ["pre"] },
 ]
+overrides.matrix.extras.features = [
+  { value = "jax", if = ["jax"] },
+]
 # default commands (only `cov-report` is overridden)
 scripts.run = "pytest{env:HATCH_TEST_ARGS:} -p no:cov {args}"
 scripts.run-cov = "coverage run -m pytest{env:HATCH_TEST_ARGS:} -p no:cov {args}"

pyproject.toml

@@ -76,6 +76,9 @@ dependencies = [
   "xarray>=2024.10",
   "zarr>=3",
 ]
+optional-dependencies.jax = [
+  "jax",
+]
 optional-dependencies.leiden = [
   "leidenalg",
   "spatialleiden>=0.4",

src/squidpy/experimental/_methods/__init__.py

@@ -0,0 +1,8 @@
+"""In-memory model-fitting core for experimental methods."""
+
+from __future__ import annotations
+
+from squidpy.experimental._methods._protocols import AlignLandmarksFn, AlignResult, AlignSamplesFn
+from squidpy.experimental._methods._registry import Registry
+
+__all__ = ["Registry", "AlignResult", "AlignSamplesFn", "AlignLandmarksFn"]

src/squidpy/experimental/_methods/_families.py

@@ -0,0 +1,14 @@
+"""Estimator family registries."""
+
+from __future__ import annotations
+
+from squidpy.experimental._methods._protocols import AlignLandmarksFn, AlignSamplesFn
+from squidpy.experimental._methods._registry import Registry
+
+#: Sample-to-sample alignment estimators -- ref/query point clouds in, transform out.
+#: Consumed by ``squidpy.experimental.tl.align``.
+ALIGN_SAMPLES: Registry[AlignSamplesFn] = Registry("align_samples")
+
+#: Closed-form landmark alignment estimators -- paired landmarks in, affine out.
+#: Consumed by ``squidpy.experimental.tl.align_by_landmarks``.
+ALIGN_LANDMARKS: Registry[AlignLandmarksFn] = Registry("align_landmarks")

src/squidpy/experimental/_methods/_protocols.py

@@ -0,0 +1,59 @@
+"""Structural contracts shared across the alignment estimator families.
+
+These :class:`~typing.Protocol` types are what the public API and the registries
+are typed against, so the orchestration layer never names a concrete estimator
+result (e.g. ``StalignResult``). A new estimator only has to satisfy
+:class:`AlignResult` -- a ``transform`` that maps points into the reference
+frame -- to plug into :func:`squidpy.experimental.tl.align`.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+import numpy.typing as npt
+
+from squidpy._utils import NDArrayA
+
+if TYPE_CHECKING:
+    from squidpy.experimental._methods.align_landmarks._landmark import AffineFitResult
+
+__all__ = ["AlignResult", "AlignSamplesFn", "AlignLandmarksFn"]
+
+
+@runtime_checkable
+class AlignResult(Protocol):
+    """A fitted alignment that maps ``(N, 2)`` ``(x, y)`` points into the reference frame.
+
+    This is the only thing the public ``align*`` functions require of an
+    estimator's result, so ``output_mode="object"`` is agnostic to the method
+    that produced it.
+    """
+
+    def transform(self, points: npt.ArrayLike, /) -> NDArrayA:
+        """Map an ``(N, 2)`` ``(x, y)`` array into the reference frame."""
+        ...
+
+
+class AlignSamplesFn(Protocol):
+    """Calling convention for ``align_samples`` estimators.
+
+    Two point clouds in (passed by keyword as ``ref`` / ``query`` so the
+    direction can never be silently swapped), one :class:`AlignResult` out.
+    Solver-specific options arrive through ``**kwargs``.
+    """
+
+    def __call__(self, ref: npt.ArrayLike, query: npt.ArrayLike, **kwargs: Any) -> AlignResult: ...
+
+
+class AlignLandmarksFn(Protocol):
+    """Calling convention for ``align_landmarks`` estimators: paired landmarks in, affine out."""
+
+    def __call__(
+        self,
+        landmarks_ref: npt.ArrayLike,
+        landmarks_query: npt.ArrayLike,
+        *,
+        source_cs: str | None = ...,
+        target_cs: str | None = ...,
+    ) -> AffineFitResult: ...

src/squidpy/experimental/_methods/_registry.py

@@ -0,0 +1,73 @@
+"""A flat registry mapping method names to fitting functions."""
+
+from __future__ import annotations
+
+import functools
+import importlib.util
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar
+
+#: The calling convention a family's registry advertises (returned by :meth:`Registry.get`).
+F = TypeVar("F", bound=Callable[..., Any])
+#: The concrete function being registered. Kept separate from ``F`` so an estimator may
+#: declare specific keyword parameters (e.g. ``config=``) without having to structurally
+#: match the family's open-ended ``**kwargs`` calling convention.
+RegisteredT = TypeVar("RegisteredT", bound=Callable[..., Any])
+
+
+class Registry(Generic[F]):
+    """A flat ``name -> function`` registry for one *family* of methods.
+
+    One :class:`Registry` is created per family (e.g. ``align``, ``impute``),
+    so keys are plain method names -- there is no ``(method, mode)`` compound
+    key, because the family already pins the rest.
+
+    The type parameter ``F`` is the family's calling convention (a callable
+    :class:`~typing.Protocol`); :meth:`get` returns it, so dispatch sites are
+    typed against the family contract rather than ``Callable[..., Any]``.
+    """
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self._registry: dict[str, F] = {}
+
+    def register(self, key: str, *, requires: tuple[str, ...] = ()) -> Callable[[RegisteredT], RegisteredT]:
+        """Return a decorator registering a method/function under ``key``."""
+
+        def decorator(func: RegisteredT) -> RegisteredT:
+            if key in self._registry:
+                raise ValueError(f"Method {key!r} is already registered in the {self.name!r} registry.")
+
+            if requires:
+
+                @functools.wraps(func)
+                def wrapped(*args: Any, **kwargs: Any) -> Any:
+                    missing = [pkg for pkg in requires if importlib.util.find_spec(pkg) is None]
+                    if missing:
+                        verb = "is" if len(missing) == 1 else "are"
+                        names = ", ".join(repr(p) for p in missing)
+                        extras = ",".join(missing)
+                        raise ImportError(
+                            f"Method {key!r} requires {names}, which {verb} not installed. "
+                            f'Install with `pip install "squidpy[{extras}]"`.'
+                        )
+                    return func(*args, **kwargs)
+
+                self._registry[key] = wrapped  # type: ignore[assignment]
+                return wrapped  # type: ignore[return-value]
+            else:
+                self._registry[key] = func  # type: ignore[assignment]
+                return func
+
+        return decorator
+
+    def get(self, key: str) -> F:
+        """Return the function registered under ``key``."""
+        try:
+            return self._registry[key]
+        except KeyError:
+            raise ValueError(f"Unknown {self.name} method {key!r}. Available: {sorted(self._registry)}.") from None
+
+    def keys(self) -> tuple[str, ...]:
+        """Return the registered method names."""
+        return tuple(self._registry)

src/squidpy/experimental/_methods/align_landmarks/__init__.py

@@ -0,0 +1,17 @@
+"""``align_landmarks`` family: closed-form alignment from paired landmarks."""
+
+from __future__ import annotations
+
+from squidpy.experimental._methods._families import ALIGN_LANDMARKS
+from squidpy.experimental._methods.align_landmarks._landmark import (
+    AffineFitResult,
+    fit_affine,
+    fit_similarity,
+)
+
+__all__ = [
+    "ALIGN_LANDMARKS",
+    "AffineFitResult",
+    "fit_affine",
+    "fit_similarity",
+]

src/squidpy/experimental/_methods/align_landmarks/_landmark.py

@@ -0,0 +1,157 @@
+"""Closed-form landmark alignment estimators."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+from squidpy._utils import NDArrayA
+from squidpy.experimental._methods._families import ALIGN_LANDMARKS
+
+
+@dataclass
+class AffineFitResult:
+    """A fitted ``(3, 3)`` homogeneous affine mapping query onto ref, in ``(x, y)``."""
+
+    matrix: np.ndarray
+    source_cs: str | None = None
+    target_cs: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.matrix.shape != (3, 3):
+            raise ValueError(f"Expected a (3, 3) homogeneous matrix, found shape {self.matrix.shape}.")
+
+    def transform(self, x: npt.ArrayLike) -> NDArrayA:
+        """Apply the affine to an ``(N, 2)`` ``(x, y)`` coordinate array."""
+        coords = np.asarray(x, dtype=float)
+        if coords.ndim != 2 or coords.shape[1] != 2:
+            raise ValueError(f"Expected an (N, 2) coordinate array, found shape {coords.shape}.")
+        return coords @ self.matrix[:2, :2].T + self.matrix[:2, 2]
+
+
+def _fit_landmark_relation(
+    landmarks_ref: np.ndarray,
+    landmarks_query: np.ndarray,
+    *,
+    method: str,
+    solve_fn: Callable[[np.ndarray, np.ndarray], np.ndarray],
+    source_cs: str | None = None,
+    target_cs: str | None = None,
+) -> AffineFitResult:
+    ref = _validate_landmarks(landmarks_ref, name="landmarks_ref")
+    query = _validate_landmarks(landmarks_query, name="landmarks_query")
+    if ref.shape != query.shape:
+        raise ValueError(
+            f"`landmarks_ref` and `landmarks_query` must have the same shape; got {ref.shape} and {query.shape}."
+        )
+    if ref.shape[0] < 3:
+        raise ValueError(f"`{method}` needs at least 3 landmark pairs, got {ref.shape[0]}.")
+
+    matrix = solve_fn(ref, query)
+    return AffineFitResult(
+        matrix=matrix,
+        source_cs=source_cs,
+        target_cs=target_cs,
+        metadata={"method": method},
+    )
+
+
+@ALIGN_LANDMARKS.register("similarity")
+def fit_similarity(
+    landmarks_ref: np.ndarray,
+    landmarks_query: np.ndarray,
+    *,
+    source_cs: str | None = None,
+    target_cs: str | None = None,
+) -> AffineFitResult:
+    """4-DOF similarity fit (rotation + uniform scale + translation), via spatialdata.
+
+    Parameters
+    ----------
+    landmarks_ref, landmarks_query
+        Pre-paired ``(N, 2)`` ``(x, y)`` landmark arrays (``N >= 3``).
+    source_cs, target_cs
+        Optional coordinate-system labels stamped onto the result for
+        traceability; they do not affect the fit.
+    """
+    return _fit_landmark_relation(
+        landmarks_ref,
+        landmarks_query,
+        method="similarity",
+        solve_fn=_fit_similarity,
+        source_cs=source_cs,
+        target_cs=target_cs,
+    )
+
+
+@ALIGN_LANDMARKS.register("affine")
+def fit_affine(
+    landmarks_ref: np.ndarray,
+    landmarks_query: np.ndarray,
+    *,
+    source_cs: str | None = None,
+    target_cs: str | None = None,
+) -> AffineFitResult:
+    """6-DOF affine fit (rotation + non-uniform scale + shear + translation), via skimage.
+
+    Parameters
+    ----------
+    landmarks_ref, landmarks_query
+        Pre-paired ``(N, 2)`` ``(x, y)`` landmark arrays (``N >= 3``).
+    source_cs, target_cs
+        Optional coordinate-system labels stamped onto the result for
+        traceability; they do not affect the fit.
+    """
+    return _fit_landmark_relation(
+        landmarks_ref,
+        landmarks_query,
+        method="affine",
+        solve_fn=_fit_affine,
+        source_cs=source_cs,
+        target_cs=target_cs,
+    )
+
+
+def _validate_landmarks(points: np.ndarray, *, name: str) -> np.ndarray:
+    arr = np.asarray(points, dtype=float)
+    if arr.ndim != 2 or arr.shape[1] != 2:
+        raise ValueError(f"`{name}` must be a sequence of (x, y) pairs, got shape {arr.shape}.")
+    if not np.all(np.isfinite(arr)):
+        raise ValueError(f"`{name}` must contain only finite values.")
+    return arr
+
+
+def _fit_similarity(ref_xy: np.ndarray, query_xy: np.ndarray) -> np.ndarray:
+    """4-DOF similarity fit, delegated to spatialdata."""
+    from spatialdata.models import PointsModel
+    from spatialdata.transformations import get_transformation_between_landmarks
+
+    refs_pts = PointsModel.parse(ref_xy)
+    moving_pts = PointsModel.parse(query_xy)
+    sd_transform = get_transformation_between_landmarks(refs_pts, moving_pts)
+    return _extract_affine_matrix(sd_transform)
+
+
+def _fit_affine(ref_xy: np.ndarray, query_xy: np.ndarray) -> np.ndarray:
+    """Full 6-DOF affine fit, delegated to skimage's least-squares estimator."""
+    from skimage.transform import estimate_transform
+
+    model_obj = estimate_transform("affine", src=query_xy, dst=ref_xy)
+    return np.asarray(model_obj.params)
+
+
+def _extract_affine_matrix(sd_transform: object) -> np.ndarray:
+    """Pull a ``(3, 3)`` homogeneous matrix out of a spatialdata transformation."""
+    from spatialdata.transformations import Affine as SDAffine
+    from spatialdata.transformations import Sequence as SDSequence
+
+    if isinstance(sd_transform, SDAffine):
+        return np.asarray(sd_transform.matrix)
+    if isinstance(sd_transform, SDSequence):
+        return np.asarray(sd_transform.to_affine_matrix(input_axes=("x", "y"), output_axes=("x", "y")))
+    raise TypeError(f"Unexpected transformation type from spatialdata: {type(sd_transform).__name__}.")

src/squidpy/experimental/_methods/align_samples/__init__.py

@@ -0,0 +1,14 @@
+"""``align_samples`` family: align two samples' point clouds (STalign).
+
+Importing this package registers the family's estimators into
+:data:`~squidpy.experimental._methods._families.ALIGN_SAMPLES`. It stays cheap --
+JAX is pulled in lazily, only when an estimator's ``fit`` runs.
+"""
+
+from __future__ import annotations
+
+from squidpy.experimental._methods._families import ALIGN_SAMPLES
+from squidpy.experimental._methods.align_samples._stalign import fit_stalign
+from squidpy.experimental._methods.align_samples._stalign_impl._tools import StalignResult
+
+__all__ = ["ALIGN_SAMPLES", "fit_stalign", "StalignResult"]