Flatten stalign subpackage, switch to xy convention, wire landmark backend, add blobs e2e tests

Four user-driven changes on top of the previous integration commit:

1. Flatten src/squidpy/experimental/tl/_align/_stalign/ into siblings
   under _align/_backends/ (_stalign_core, _stalign_helpers, _stalign_tools).
   The lifted Selman code now lives next to its only consumer
   (_backends/_stalign.py) and there's no extra package layer.

2. Switch AffineTransform / ObsDisplacement / to_spatialdata() to xy
   convention to match spatialdata's points API (which asserts axes ==
   ('x', 'y') in get_transformation_between_landmarks) and obsm['spatial'].
   Drops the latent yx-vs-xy gotcha from materialise_obs.

3. Wire fit_landmark_affine to real solvers:
   - model='similarity' -> spatialdata.transformations.get_transformation_between_landmarks
   - model='affine'     -> skimage.transform.estimate_transform('affine', ...)
   Both produce 3x3 homogeneous matrices in xy. apply_affine_to_cs grows a
   third writeback path (cs-keyed: walk every element registered in the
   moving cs and set_transformation on each) so align_by_landmarks can
   register the fit without an explicit element key.

4. Add tests/experimental/tl/test_align_blobs_e2e.py: four end-to-end tests
   on the spatialdata blobs() fixture (200 points). The closed-form
   landmark fits recover a known 12 deg rotation to 1e-6 max residual; the
   stalign LDDMM run reduces the residual below the no-op baseline.

Tests: 33/33 passing. Lint clean.

Author: claude

src/squidpy/experimental/tl/__init__.py

@@ -6,7 +6,7 @@
 from squidpy.experimental.tl._tiling_qc import calculate_tiling_qc
 
 if TYPE_CHECKING:
-    from squidpy.experimental.tl._align._stalign._tools import (
+    from squidpy.experimental.tl._align._backends._stalign_tools import (
         STalignConfig,
         STalignPreprocessConfig,
         STalignPreprocessResult,
@@ -47,7 +47,7 @@ def __getattr__(name: str) -> Any:
     """
     if name in _STALIGN_REEXPORTS:
         try:
-            from squidpy.experimental.tl._align._stalign import _tools
+            from squidpy.experimental.tl._align._backends import _stalign_tools as _tools
         except ModuleNotFoundError as e:
             if e.name == "jax":
                 raise ImportError(

src/squidpy/experimental/tl/_align/_backends/_landmark.py

@@ -1,11 +1,16 @@
 """Closed-form landmark fit.
 
-Pure NumPy — no JAX, no optional deps.  Used directly by
-:func:`squidpy.experimental.tl.align_by_landmarks`; in a later PR the same
-fit will also serve as the initial guess for the STalign backend.
+Two models, both pure NumPy / no JAX:
 
-The actual math (Umeyama for ``model='similarity'``, linear least-squares for
-``model='affine'``) lands together with the STalign body in the follow-up PR.
+- ``"similarity"`` (4 DOF: rotation + uniform scale + translation, plus an
+  optional reflection check) - delegated to
+  :func:`spatialdata.transformations.get_transformation_between_landmarks`.
+- ``"affine"`` (6 DOF: rotation + non-uniform scale + shear + translation) -
+  delegated to :func:`skimage.transform.estimate_transform`, the same
+  least-squares solver spatialdata uses internally.
+
+Useful as a one-shot alignment when you already have corresponding landmarks,
+and as a sanity-check baseline for the much heavier STalign LDDMM solver.
 """
 
 from __future__ import annotations
@@ -28,12 +33,74 @@ def fit_landmark_affine(
 ) -> AffineTransform:
     """Fit a 2D affine that maps ``landmarks_query`` onto ``landmarks_ref``.
 
-    Both inputs are ``(N, 2)`` arrays in ``(y, x)`` convention.  The returned
-    :class:`~squidpy.experimental.tl._align._types.AffineTransform` is a
-    ``(3, 3)`` homogeneous matrix in the same convention.
+    Both inputs are ``(N, 2)`` ``(x, y)`` arrays of corresponding landmarks
+    (the ``i``-th row of ``landmarks_query`` matches the ``i``-th row of
+    ``landmarks_ref``).  ``N`` must be at least 3.
+
+    Parameters
+    ----------
+    landmarks_ref, landmarks_query
+        Corresponding landmark coordinates in ``(x, y)`` convention.
+    model
+        ``"similarity"`` (4 DOF, via spatialdata) or ``"affine"`` (6 DOF,
+        via skimage).
+    source_cs, target_cs
+        Optional coordinate-system labels stamped onto the returned
+        :class:`AffineTransform` for traceability.
     """
-    raise NotImplementedError(
-        "Landmark fit: TODO. Skeleton landed; the closed-form Umeyama / "
-        "least-squares solver will land in the next PR alongside the STalign "
-        "backend body."
-    )
+    from squidpy.experimental.tl._align._types import AffineTransform
+
+    ref = np.asarray(landmarks_ref, dtype=float)
+    query = np.asarray(landmarks_query, dtype=float)
+
+    if model == "similarity":
+        matrix = _fit_similarity_via_spatialdata(ref, query)
+    elif model == "affine":
+        matrix = _fit_affine_via_skimage(ref, query)
+    else:
+        raise ValueError(f"Unknown landmark `model={model!r}`; expected 'similarity' or 'affine'.")
+
+    return AffineTransform(matrix=matrix, source_cs=source_cs, target_cs=target_cs)
+
+
+def _fit_similarity_via_spatialdata(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_via_skimage(ref_xy: np.ndarray, query_xy: np.ndarray) -> np.ndarray:
+    """Full 6-DOF affine fit, delegated to skimage's least-squares solver.
+
+    This is what :func:`spatialdata.transformations.get_transformation_between_landmarks`
+    uses under the hood before collapsing to a similarity; for the affine
+    model we keep the raw matrix instead.
+    """
+    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.
+
+    :func:`get_transformation_between_landmarks` may return either a single
+    :class:`spatialdata.transformations.Affine` or a
+    :class:`spatialdata.transformations.Sequence` of two affines (when a
+    reflection is detected and rolled into the fit).  Use
+    ``to_affine_matrix`` to collapse either back to a single 3x3.
+    """
+    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/tl/_align/_backends/_stalign.py

@@ -41,7 +41,7 @@ def align_obs(
         # _jax.require_jax.
         require_jax(device)
 
-        from squidpy.experimental.tl._align._stalign._tools import stalign_points
+        from squidpy.experimental.tl._align._backends._stalign_tools import stalign_points
         from squidpy.experimental.tl._align._types import AlignResult, ObsDisplacement
 
         if not isinstance(pair.ref, AnnData) or not isinstance(pair.query, AnnData):

…experimental/tl/_align/_stalign/_core.py …tal/tl/_align/_backends/_stalign_core.pysrc/squidpy/experimental/tl/_align/_stalign/_core.py renamed to src/squidpy/experimental/tl/_align/_backends/_stalign_core.py src/squidpy/experimental/tl/_align/_stalign/_core.py renamed to src/squidpy/experimental/tl/_align/_backends/_stalign_core.py

@@ -1,7 +1,6 @@
 """Core JAX implementation for experimental STalign point registration.
 
 Lifted byte-for-byte from scverse/squidpy#1150 (Selman Özleyen).
-See ``_align/_stalign/__init__.py`` for the lift note.
 """
 
 from __future__ import annotations

…erimental/tl/_align/_stalign/_helpers.py …/tl/_align/_backends/_stalign_helpers.pysrc/squidpy/experimental/tl/_align/_stalign/_helpers.py renamed to src/squidpy/experimental/tl/_align/_backends/_stalign_helpers.py src/squidpy/experimental/tl/_align/_stalign/_helpers.py renamed to src/squidpy/experimental/tl/_align/_backends/_stalign_helpers.py

@@ -1,7 +1,6 @@
 """Helpers for experimental STalign point-cloud registration.
 
 Lifted byte-for-byte from scverse/squidpy#1150 (Selman Özleyen).
-See ``_align/_stalign/__init__.py`` for the lift note.
 """
 
 from __future__ import annotations

…xperimental/tl/_align/_stalign/_tools.py …al/tl/_align/_backends/_stalign_tools.pysrc/squidpy/experimental/tl/_align/_stalign/_tools.py renamed to src/squidpy/experimental/tl/_align/_backends/_stalign_tools.py src/squidpy/experimental/tl/_align/_stalign/_tools.py renamed to src/squidpy/experimental/tl/_align/_backends/_stalign_tools.py

@@ -1,9 +1,8 @@
 """Low-level point-cloud tools for experimental STalign.
 
 Lifted from scverse/squidpy#1150 (Selman Özleyen). Only the two import paths
-on lines 12-20 were rewritten to point at the relocated submodules; the rest
-of the file is byte-for-byte identical to the upstream PR. See
-``_align/_stalign/__init__.py`` for the lift note.
+below were rewritten to point at the sibling lifted modules; the rest of the
+file is byte-for-byte identical to the upstream PR.
 """
 
 from __future__ import annotations
@@ -15,8 +14,8 @@
 import numpy as np
 from anndata import AnnData
 
-from squidpy.experimental.tl._align._stalign._core import JAX_DTYPE, lddmm, transform_points_row_col
-from squidpy.experimental.tl._align._stalign._helpers import (
+from squidpy.experimental.tl._align._backends._stalign_core import JAX_DTYPE, lddmm, transform_points_row_col
+from squidpy.experimental.tl._align._backends._stalign_helpers import (
     PointOrder,
     affine_from_points,
     extract_points,

src/squidpy/experimental/tl/_align/_io.py

@@ -192,23 +192,49 @@ def apply_affine_to_cs(
     *,
     inplace: bool,
 ) -> SpatialData | AnnData | None:
-    """Register ``affine`` as a transformation on the query element.
-
-    The transformation goes onto the parent element node (so all scales of a
-    multiscale image inherit it) under the *reference* coordinate-system name.
-    For plain-AnnData inputs there is no SpatialData element to attach to —
-    we apply the affine to ``query.obsm['spatial']`` in place.
+    """Register ``affine`` on the query side of the pair.
+
+    Three writeback paths, in order of specificity:
+
+    1. **Element-keyed**: ``pair.query_container`` and ``pair.query_element_key``
+       are both set (e.g. ``align_obs`` / ``align_images`` resolved an explicit
+       table or image). Register the transform on that single element so all
+       scales / sibling tables that share its parent element node inherit it.
+    2. **Cs-keyed**: only ``pair.query_cs`` is set (e.g. ``align_by_landmarks``
+       resolved a coordinate system but no specific element). Walk every
+       element that has the moving cs in its transformation graph and register
+       the transform on each, mapping into the reference cs.
+    3. **Plain AnnData**: no spatialdata container at all - warp
+       ``query.obsm['spatial']`` directly.
     """
-    from spatialdata.transformations import set_transformation
+    from spatialdata.transformations import get_transformation, set_transformation
+
+    target_cs = affine.target_cs or pair.ref_cs or "aligned"
 
     if pair.query_container is not None and pair.query_element_key is not None:
         sdata = pair.query_container if inplace else _shallow_copy_sdata(pair.query_container)
         element = sdata[pair.query_element_key]
-        target_cs = affine.target_cs or pair.ref_cs or "aligned"
         set_transformation(element, affine.to_spatialdata(), to_coordinate_system=target_cs)
         return None if inplace else sdata
 
-    # Plain AnnData query -> warp obsm["spatial"].
+    if pair.query_container is not None and pair.query_cs is not None:
+        sdata = pair.query_container if inplace else _shallow_copy_sdata(pair.query_container)
+        moving_cs = pair.query_cs
+        sd_affine = affine.to_spatialdata()
+        touched_any = False
+        for _etype, _name, element in sdata._gen_elements(include_tables=False):
+            element_transforms = get_transformation(element, get_all=True)
+            if moving_cs not in element_transforms:
+                continue
+            set_transformation(element, sd_affine, to_coordinate_system=target_cs)
+            touched_any = True
+        if not touched_any:
+            raise KeyError(
+                f"No elements in the query SpatialData are registered to coordinate "
+                f"system {moving_cs!r}; nothing to attach the alignment to."
+            )
+        return None if inplace else sdata
+
     if isinstance(pair.query, AnnData):
         adata = pair.query if inplace else pair.query.copy()
         if "spatial" not in adata.obsm:
@@ -239,14 +265,6 @@ def materialise_obs(
         raise KeyError("Source AnnData has no `obsm['spatial']` to warp.")
 
     src_coords = np.asarray(pair.query.obsm["spatial"])
-    # CONVENTION NOTE: anndata's ``obsm['spatial']`` is conventionally (x, y)
-    # in squidpy, while ``AffineTransform.matrix`` is documented as (y, x).
-    # Today the only backend that returns a ``Transform`` for AnnData inputs
-    # is stalign, which returns an ``ObsDisplacement`` whose ``deltas`` are
-    # already in the same convention as ``obsm['spatial']`` - so the addition
-    # below is correct for the displacement path. The first backend that
-    # returns an ``AffineTransform`` for an AnnData input will need to either
-    # swap here, or normalise the matrix to xy at the backend boundary.
     new_coords = result.transform.apply(src_coords)
 
     # Slim copy: share X/var/obs structurally and only rewrite obsm so we

src/squidpy/experimental/tl/_align/_stalign/__init__.py

@@ -44,7 +44,16 @@ class AlignPair:
 
 @dataclass(frozen=True)
 class AffineTransform:
-    """A ``(3, 3)`` homogeneous affine in ``(y, x)`` convention."""
+    """A ``(3, 3)`` homogeneous affine in ``(x, y)`` convention.
+
+    This matches the coordinate axis order spatialdata uses for points -
+    ``spatialdata.transformations.get_transformation_between_landmarks``
+    asserts ``axes == ("x", "y")`` - and the order squidpy / scanpy use for
+    ``adata.obsm["spatial"]``. Image elements are stored ``(c, y, x)`` in
+    spatialdata, so when registering an ``AffineTransform`` on an *image*
+    element you may need a separate matrix; this skeleton currently only
+    deals with point coordinates.
+    """
 
     matrix: np.ndarray
     source_cs: str | None = None
@@ -60,24 +69,25 @@ def to_spatialdata(self) -> Affine:
 
         return Affine(
             self.matrix,
-            input_axes=("y", "x"),
-            output_axes=("y", "x"),
+            input_axes=("x", "y"),
+            output_axes=("x", "y"),
         )
 
     def apply(self, coords: np.ndarray) -> np.ndarray:
-        """Apply the affine to an ``(N, 2)`` ``(y, x)`` coordinate array."""
+        """Apply the affine to an ``(N, 2)`` ``(x, y)`` coordinate array."""
         if coords.ndim != 2 or coords.shape[1] != 2:
             raise ValueError(f"Expected an (N, 2) coordinate array, got shape {coords.shape}.")
         return coords @ self.matrix[:2, :2].T + self.matrix[:2, 2]
 
 
 @dataclass(frozen=True)
 class ObsDisplacement:
-    """Per-obs ``(N, 2)`` displacement field.
+    """Per-obs ``(N, 2)`` ``(x, y)`` displacement field.
 
     Used by non-affine fits (e.g. LDDMM) where a single matrix cannot
     represent the deformation.  Displacements are added to the source
-    observation coordinates to obtain the aligned coordinates.
+    observation coordinates (also ``(x, y)``) to obtain the aligned
+    coordinates.
     """
 
     deltas: np.ndarray

src/squidpy/experimental/tl/_align/_types.py

@@ -62,11 +62,14 @@ def validate_landmarks(
     Parameters
     ----------
     landmarks_ref, landmarks_query
-        Sequences of ``(y, x)`` tuples.  Must have the same length.
+        Sequences of ``(x, y)`` tuples.  Must have the same length and at
+        least 3 entries (the closed-form solvers under both ``model``
+        choices need at least 3 corresponding points).
     model
-        ``"similarity"`` (≥ 2 landmarks) or ``"affine"`` (≥ 3 landmarks).
+        ``"similarity"`` (4 DOF, via spatialdata) or ``"affine"`` (6 DOF,
+        via skimage's least-squares estimator).
     cs_ref_extent, cs_query_extent
-        Optional ``(y_min, x_min, y_max, x_max)`` bounds of the named
+        Optional ``(x_min, y_min, x_max, y_max)`` bounds of the named
         coordinate system at the requested scale.  When provided, every
         landmark must fall inside.  Catches the "I extracted these from
         scale0 but asked for scale2" footgun.
@@ -75,17 +78,18 @@ def validate_landmarks(
     query = np.asarray(landmarks_query, dtype=float)
 
     if ref.ndim != 2 or ref.shape[1] != 2:
-        raise ValueError(f"`landmarks_ref` must be a sequence of (y, x) pairs, got shape {ref.shape}.")
+        raise ValueError(f"`landmarks_ref` must be a sequence of (x, y) pairs, got shape {ref.shape}.")
     if query.ndim != 2 or query.shape[1] != 2:
-        raise ValueError(f"`landmarks_query` must be a sequence of (y, x) pairs, got shape {query.shape}.")
+        raise ValueError(f"`landmarks_query` must be a sequence of (x, y) pairs, got shape {query.shape}.")
     if len(ref) != len(query):
         raise ValueError(
             f"`landmarks_ref` and `landmarks_query` must have the same length; got {len(ref)} and {len(query)}."
         )
 
-    min_landmarks = 3 if model == "affine" else 2
-    if len(ref) < min_landmarks:
-        raise ValueError(f"`model={model!r}` needs at least {min_landmarks} landmark pairs, got {len(ref)}.")
+    if len(ref) < 3:
+        raise ValueError(
+            f"`model={model!r}` needs at least 3 landmark pairs (spatialdata requirement), got {len(ref)}."
+        )
 
     if cs_ref_extent is not None:
         _check_in_extent(ref, cs_ref_extent, name="landmarks_ref")
@@ -101,13 +105,13 @@ def _check_in_extent(
     *,
     name: str,
 ) -> None:
-    y_min, x_min, y_max, x_max = extent
-    out_of_bounds = (points[:, 0] < y_min) | (points[:, 0] > y_max) | (points[:, 1] < x_min) | (points[:, 1] > x_max)
+    x_min, y_min, x_max, y_max = extent
+    out_of_bounds = (points[:, 0] < x_min) | (points[:, 0] > x_max) | (points[:, 1] < y_min) | (points[:, 1] > y_max)
     if out_of_bounds.any():
         bad = points[out_of_bounds]
         raise ValueError(
             f"{name}: {int(out_of_bounds.sum())} landmark(s) fall outside the coordinate-system "
-            f"extent (y in [{y_min}, {y_max}], x in [{x_min}, {x_max}]). "
+            f"extent (x in [{x_min}, {x_max}], y in [{y_min}, {y_max}]). "
             f"This usually means the landmarks were extracted at a different scale than the "
             f"one requested. First out-of-bounds point: {tuple(bad[0])}."
         )