Address review findings on alignment skeleton

- Replace private `sdata._gen_elements()` with public `gen_elements()`
- Replace dict-style `sdata[key]` lookup with explicit element-type search
- Add subprocess timeout (30s) to lazy-import hygiene test
- Document shallow X sharing in `materialise_obs` docstring
- Document JAX array retention in stalign metadata comment
- Document camelCase convention in STalignRegistrationConfig docstring
- Broaden landmark type hints to accept Sequence and np.ndarray
- Remove stale TODO comment from _stalign_helpers.py

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timtreis

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

@@ -8,8 +8,11 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Literal
 
+import numpy as np
+
 from squidpy.experimental.tl._align._backends import get_backend
 from squidpy.experimental.tl._align._io import (
     apply_affine_to_cs,
@@ -160,8 +163,8 @@ def align_by_landmarks(
     cs_name_query: str | None = None,
     scale_ref: str | None = None,
     scale_query: str | None = None,
-    landmarks_ref: tuple[tuple[float, float], ...] | None = None,
-    landmarks_query: tuple[tuple[float, float], ...] | None = None,
+    landmarks_ref: Sequence[tuple[float, float]] | np.ndarray | None = None,
+    landmarks_query: Sequence[tuple[float, float]] | np.ndarray | None = None,
     *,
     model: Literal["similarity", "affine"] = "similarity",
     output_mode: Literal["affine", "return"] = "affine",

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

@@ -82,9 +82,11 @@ def align_obs(
             ),
             metadata={
                 "flavour": "stalign",
-                # Escape hatch for power users who want the diffeomorphic
-                # part (velocity field, velocity grid, affine init) rather
-                # than just the materialised displacement.
+                # Escape hatch: the full STalignResult (velocity field,
+                # velocity grid, affine init) for power users who need
+                # the diffeomorphic map.  This keeps the JAX arrays alive
+                # in memory -- callers who only need the displacement
+                # should drop this key or use ``output_mode='obs'``.
                 "stalign_result": stalign_result,
             },
         )

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

@@ -56,10 +56,6 @@ def extract_landmarks(adata: AnnData, key: str) -> np.ndarray:
     raise KeyError(f"Key `{key}` not found in `adata.obsm` or `adata.uns`.")
 
 
-# TODO: are these duplicated? I would imagine its
-# better to keep image transform functions under some place
-
-
 def to_row_col(points: np.ndarray, *, point_order: PointOrder) -> np.ndarray:
     """Convert coordinates to row-column order."""
     arr = _validate_points(points, name="points")

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

@@ -54,6 +54,13 @@ class STalignPreprocessConfig:
 
 @dataclass(slots=True)
 class STalignRegistrationConfig:
+    """LDDMM registration hyperparameters.
+
+    Field names (``sigmaM``, ``epL``, etc.) preserve the conventions from
+    the STalign paper and reference implementation to keep them
+    recognisable when cross-referencing the literature.
+    """
+
     a: float = 500.0
     p: float = 2.0
     expand: float = 2.0

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

@@ -213,7 +213,7 @@ def apply_affine_to_cs(
 
     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]
+        element = _get_element(sdata, pair.query_element_key)
         set_transformation(element, affine.to_spatialdata(), to_coordinate_system=target_cs)
         return None if inplace else sdata
 
@@ -222,7 +222,9 @@ def apply_affine_to_cs(
         moving_cs = pair.query_cs
         sd_affine = affine.to_spatialdata()
         touched_any = False
-        for _etype, _name, element in sdata._gen_elements(include_tables=False):
+        for _etype, _name, element in sdata.gen_elements():
+            if isinstance(element, AnnData):
+                continue
             element_transforms = get_transformation(element, get_all=True)
             if moving_cs not in element_transforms:
                 continue
@@ -258,6 +260,13 @@ def materialise_obs(
     add the deltas.  When the source query lives inside a SpatialData, the new
     AnnData is registered as ``sdata.tables[key_added]``; otherwise it is
     returned directly.
+
+    .. note::
+
+       The returned AnnData **shares** ``X`` and ``var`` with the source
+       query by reference to avoid copying potentially-large expression
+       matrices.  Mutating one will affect the other.  Call
+       ``.copy()`` on the result if you need full independence.
     """
     if not isinstance(pair.query, AnnData):
         raise TypeError("materialise_obs only works for `align_obs`; `pair.query` must be an AnnData.")
@@ -297,13 +306,21 @@ def materialise_obs(
     return new_adata
 
 
+def _get_element(sdata: SpatialData, key: str) -> object:
+    """Look up a spatial element by name across all element types."""
+    for attr in ("images", "labels", "points", "shapes", "tables"):
+        store = getattr(sdata, attr)
+        if key in store:
+            return store[key]
+    raise KeyError(f"Element {key!r} not found in the SpatialData object.")
+
+
 def _shallow_copy_sdata(sdata: SpatialData) -> SpatialData:
     """Shallow copy of a SpatialData object for ``inplace=False`` writeback paths.
 
     Uses :meth:`SpatialData.subset` over every element so tables and
     ``attrs`` propagate the same way spatialdata's own subsetting handles
     them, rather than reconstructing via the ``__init__`` constructor.
     """
-    element_names = [name for _, name, _ in sdata._gen_elements(include_tables=True)]
+    element_names = [name for _, name, _ in sdata.gen_elements()]
     return sdata.subset(element_names, filter_tables=False, include_orphan_tables=True)
-

tests/experimental/tl/test_align_skeleton.py

@@ -132,6 +132,7 @@ def test_optional_deps_not_imported_at_import_time() -> None:
             ),
         ],
         text=True,
+        timeout=30,
     ).strip()
     assert out == "", f"Optional deps imported by `import squidpy`: {out}"