feat: stalign with a fit apply pattern.

docs/api.md

@@ -140,6 +140,24 @@ See the {doc}`extensibility guide </extensibility>` for how to implement a custo
     gr.neighbors.GridBuilder
 ```
 
+The ``method=`` argument of {func}`~squidpy.experimental.tl.align` and
+{func}`~squidpy.experimental.tl.align_by_landmarks` dispatches to a registered
+fitting function. The method-specific arguments are documented on each function
+below; the fitted maps are returned (with ``output_mode="object"``) as the
+result types listed alongside.
+
+```{eval-rst}
+.. currentmodule:: squidpy
+.. autosummary::
+    :toctree: api
+
+    experimental.methods.align_samples._stalign.fit_stalign
+    experimental.methods.align_landmarks._landmark.fit_similarity
+    experimental.methods.align_landmarks._landmark.fit_affine
+    experimental.methods.align_samples._stalign.StalignResult
+    experimental.methods.align_landmarks._landmark.AffineFitResult
+```
+
 ## Experimental
 ```{eval-rst}
 .. module:: squidpy.experimental
@@ -150,6 +168,9 @@ See the {doc}`extensibility guide </extensibility>` for how to implement a custo
     experimental.im.calculate_image_features
     experimental.tl.calculate_tiling_qc
     experimental.tl.TilingQCParams
+    experimental.tl.align
+    experimental.tl.align_by_landmarks
+    experimental.tl.AlignResult
     experimental.tl.assign_stitch_groups
     experimental.tl.StitchParams
     experimental.im.make_stitched_labels

docs/conf.py

@@ -103,6 +103,7 @@
 autodoc_member_order = "groupwise"
 autodoc_typehints = "signature"
 autodoc_docstring_signature = True
+autodoc_mock_imports = ["jax"]
 napoleon_google_docstring = False
 napoleon_numpy_docstring = True
 napoleon_include_init_with_doc = False
@@ -147,6 +148,8 @@
     ("py:class", "NDArray"),
     ("py:class", "np.number"),
     ("py:class", "csr_matrix"),
+    # optional dep mocked at build time (see autodoc_mock_imports), so no resolvable target
+    ("py:class", "jax.Array"),
     # no idea why those aren’t exported
     ("py:class", "squidpy._constants._constants.SpatialAutocorr"),
     ("py:class", "squidpy._constants._constants.CoordType"),

hatch.toml

@@ -17,10 +17,14 @@ extra-dependencies = ["diff-cover"]
 matrix = [
   { deps = ["stable"], python = ["3.12", "3.13", "3.14"] },
   { deps = ["pre"], python = ["3.14"] },
+    { 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

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

src/squidpy/experimental/methods/__init__.py

@@ -0,0 +1,32 @@
+"""In-memory model-fitting core for experimental methods.
+
+The :mod:`.registry` subpackage holds the registry machinery and the family
+registries; each family subpackage (e.g. :mod:`.align_samples`,
+:mod:`.align_landmarks`) holds the estimator implementations. Importing this
+package imports those subpackages so the estimators register themselves into the
+(public) family registries. Each subpackage stays cheap to import -- heavy or
+optional dependencies (e.g. JAX) are pulled in lazily, only when an estimator
+actually runs.
+"""
+
+from __future__ import annotations
+
+# Import for side effects: populates ALIGN_SAMPLES / ALIGN_LANDMARKS.
+from squidpy.experimental.methods import align_landmarks, align_samples  # noqa: F401
+from squidpy.experimental.methods.registry import (
+    ALIGN_LANDMARKS,
+    ALIGN_SAMPLES,
+    AlignLandmarksFn,
+    AlignResult,
+    AlignSamplesFn,
+    Registry,
+)
+
+__all__ = [
+    "Registry",
+    "AlignResult",
+    "AlignSamplesFn",
+    "AlignLandmarksFn",
+    "ALIGN_SAMPLES",
+    "ALIGN_LANDMARKS",
+]

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

@@ -0,0 +1,21 @@
+"""``align_landmarks`` family: closed-form alignment from paired landmarks.
+
+Importing this package registers the family's estimators into
+:data:`~squidpy.experimental.methods.registry.ALIGN_LANDMARKS`. Only the
+implementations are re-exported here; the registry itself lives in (and is public
+from) :mod:`squidpy.experimental.methods`.
+"""
+
+from __future__ import annotations
+
+from squidpy.experimental.methods.align_landmarks._landmark import (
+    AffineFitResult,
+    fit_affine,
+    fit_similarity,
+)
+
+__all__ = [
+    "AffineFitResult",
+    "fit_affine",
+    "fit_similarity",
+]

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

@@ -0,0 +1,155 @@
+"""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.registry 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(
+    ref: np.ndarray,
+    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(ref, name="ref")
+    query = _validate_landmarks(query, name="query")
+    if ref.shape != query.shape:
+        raise ValueError(f"`ref` and `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(
+    ref: np.ndarray,
+    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
+    ----------
+    ref, 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(
+        ref,
+        query,
+        method="similarity",
+        solve_fn=_fit_similarity,
+        source_cs=source_cs,
+        target_cs=target_cs,
+    )
+
+
+@ALIGN_LANDMARKS.register("affine")
+def fit_affine(
+    ref: np.ndarray,
+    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
+    ----------
+    ref, 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(
+        ref,
+        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.registry.ALIGN_SAMPLES`. It stays
+cheap -- JAX is pulled in lazily, only when an estimator's ``fit`` runs. Only the
+implementations are re-exported here; the registry itself lives in (and is public
+from) :mod:`squidpy.experimental.methods`.
+"""
+
+from __future__ import annotations
+
+from squidpy.experimental.methods.align_samples._stalign import StalignResult, fit_stalign
+
+__all__ = ["fit_stalign", "StalignResult"]