feat: stalign with a fit apply pattern.

.github/workflows/test.yaml

@@ -155,6 +155,40 @@ jobs:
           DISPLAY: :42
         run: uvx hatch run ${{ matrix.env.name }}:run -v --color=yes -n auto
 
+      # Numerical comparison against the vendored STalign reference (scverse/squidpy#1243).
+      # Scheduled runs only: it needs float64 and is far slower than the wiring tests.
+      #
+      # Deliberately confined to one cell. jax is only installed in the -jax env, and the
+      # tolerances are calibrated on Linux -- macOS Accelerate and Linux OpenBLAS reduce in
+      # different orders, which shows up over a few hundred gradient steps.
+      #
+      # HATCH_TEST_ARGS lands after the `-m "not reference"` in hatch.toml's run script,
+      # and pytest's -m is last-wins, so this selects exactly the reference tests.
+      - name: run STalign reference comparison
+        if: >-
+          github.event_name == 'schedule'
+          && matrix.os == 'ubuntu-latest'
+          && matrix.env.name == 'hatch-test.py3.13-stable-jax'
+        env:
+          HATCH_TEST_ARGS: " -m reference"
+          JAX_ENABLE_X64: "1"
+          JAX_PLATFORMS: cpu
+          MKL_NUM_THREADS: "1"
+          OMP_NUM_THREADS: "1"
+        run: >-
+          uvx hatch run ${{ matrix.env.name }}:run -v --color=yes
+          --junitxml=stalign-reference.xml
+
+      # Every comparison records its measured relative error via record_property, so this
+      # artifact is a running log of how far the port has drifted -- including on green runs.
+      - name: Archive STalign reference measurements
+        if: always() && github.event_name == 'schedule' && matrix.env.name == 'hatch-test.py3.13-stable-jax'
+        uses: actions/upload-artifact@v4
+        with:
+          name: stalign_reference_${{ runner.os }}
+          path: ${{ github.workspace }}/stalign-reference.xml
+          if-no-files-found: ignore
+
       - name: Archive figures generated during testing
         if: always()
         uses: actions/upload-artifact@v4

docs/api.md

@@ -140,6 +140,25 @@ 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` 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_samples._stalign.fit_stalign_image
+    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 +169,8 @@ 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.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"),

docs/release/notes-dev.md

@@ -3,6 +3,63 @@
 ## Features
 
 - {func}`squidpy.experimental.im.calculate_image_features` now featurizes tiles on a shared dask engine: `n_jobs > 1` runs worker processes via a `dask.distributed.LocalCluster` (or an active `Client`), and per-tile BLAS/OpenMP threads are pinned to avoid oversubscription. This also speeds up the serial path. {func}`squidpy.experimental.tl.calculate_tiling_qc` shares the same engine. Adds `distributed` and `threadpoolctl` as dependencies.
+- Add a numerical reference suite for the experimental STalign port, comparing it against the
+  original PyTorch implementation. Deselected by default; runs on the scheduled job. Reference
+  values are generated out of band by
+  [theislab/squidpy-ports](https://github.com/theislab/squidpy-ports), so `torch` is not a squidpy
+  dependency.
+  [#1243](https://github.com/scverse/squidpy/issues/1243)
+- **Breaking (experimental):** {func}`squidpy.experimental.tl.align` now locates data with `in_`
+  and `out` paths instead of a stack of key arguments. `ref_key`, `query_key`, `spatial_key`,
+  `key_added` and `output_mode` are replaced by `in_` (e.g. `"obsm/spatial"`,
+  `"tables/slice1/obsm/spatial"`, `"images/he"`, `"shapes/landmarks"`), `out`, and `copy`.
+  `out=None` (the default) returns the fitted alignment and writes nothing. This follows the shape
+  proposed for scanpy in [scanpy#4007](https://github.com/scverse/scanpy/issues/4007).
+- **Breaking (experimental):** `squidpy.experimental.tl.align_by_landmarks` is folded into
+  {func}`squidpy.experimental.tl.align` as `by="landmarks"`. `on` becomes `by`, gaining a
+  `"landmarks"` value alongside `"obs"` and `"images"`; `in_` then names the correspondences and
+  `apply_to` names what moves. Writing to `out="cs/<name>"` registers the fitted affine on a whole
+  SpatialData coordinate system instead of materialising anything; because that moves every element
+  registered there, it refuses when the reference shares the query's coordinate system.
+- **Breaking (experimental):** the `align_samples`, `align_images` and `align_landmarks` registries
+  collapse into a single `ALIGN` registry of `AlignMethod` records, each declaring which modalities
+  it implements. Asking for one a method does not support now fails immediately and says what it
+  does support. `fit_stalign` is renamed `fit_stalign_obs` for symmetry with `fit_stalign_image`.
+- {func}`squidpy.experimental.tl.align` can now align on images. The fitted diffeomorphism cannot
+  be expressed as a SpatialData transformation, so writing to an `images/...` path materialises the
+  warped image rather than registering it lazily. Adds an `align_images` method family and
+  `squidpy.experimental.methods.align_samples.fit_stalign_image`.
+- The experimental STalign solver now runs its whole gradient descent as a single compiled
+  `lax.while_loop` instead of a Python loop around a jitted step, about **4.6x faster** per
+  iteration (2.20 to 0.46 ms on the reference fixture, so `niter=5000` drops from ~11s to ~2.4s).
+  Numerically unchanged: the reference suite still matches the original implementation at 1, 5, 50
+  and 500 iterations.
+- The reference suite now covers `fit_stalign_image`, the last part of the experimental
+  STalign port without one. It reproduces the original to ~4e-12 on the affine and ~2e-12 on the
+  velocity field over a full trajectory.
+- The experimental STalign solver returns the per-iteration `energies` trace and `n_iter`, and
+  accepts optional `tol` / `patience` early stopping. Off by default. Note the objective changes
+  definition at iteration 50, when the mixture-weight E step engages, so the convergence window
+  deliberately never spans that point.
+- **Breaking (experimental):** {func}`squidpy.experimental.tl.align_by_landmarks` takes `in_` /
+  `out` / `copy` in place of `spatial_key` / `key_added` / `output_mode`, and its coordinate-system
+  arguments are renamed `cs_ref` / `cs_query`. Because `out` is always named explicitly, the guard
+  that refused to overwrite an auto-derived key is gone.
+
+## Bugfixes
+
+- The experimental STalign estimator no longer differentiates through the contrast-transform ridge
+  solve. That solve is an expectation-maximisation M step and must be held constant; treating it as
+  part of the objective changed the search direction. Gradients now agree with the original
+  implementation to ~1e-15, previously ~1e-3.
+- **Breaking (experimental):** the rasterisation and velocity grids in the experimental STalign
+  estimator were one sample longer per axis than intended, and their length varied with
+  floating-point rounding. Output shapes change accordingly.
+- The experimental STalign rasteriser now deposits each point bilinearly rather than snapping it to
+  the nearest cell, and conserves mass at the image border. Relative error against the original
+  implementation drops from 6.2 %/2.0 %/6.0 % to 4.1 %/0.8 %/2.9 % across the default blur scales.
+- `lddmm(niter=0)` no longer raises `UnboundLocalError`.
+
 - Fix {func}`squidpy.tl.var_by_distance` behaviour when providing {mod}`numpy` arrays of coordinates as anchor point.
 - Update :attr:`squidpy.pl.var_by_distance` to show multiple variables on same plot.
   [@LLehner](https://github.com/LLehner)

hatch.toml

@@ -17,13 +17,22 @@ 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}"
+#
+# The `-m "not reference"` is the actual gate on the STalign reference suite: a CLI flag
+# works on every pytest, whereas pyproject's `[tool.pytest]` table is only read by
+# pytest >= 9. HATCH_TEST_ARGS is interpolated after it and pytest's `-m` is last-wins,
+# so the scheduled job re-enables the suite with HATCH_TEST_ARGS=' -m reference'.
+scripts.run = 'pytest -m "not reference"{env:HATCH_TEST_ARGS:} -p no:cov {args}'
+scripts.run-cov = 'coverage run -m pytest -m "not reference"{env:HATCH_TEST_ARGS:} -p no:cov {args}'
 scripts.cov-combine = ["coverage combine"]
 scripts.cov-report = [
     "coverage report",

pyproject.toml

@@ -78,6 +78,9 @@ dependencies = [
   "xarray>=2024.10",
   "zarr>=3",
 ]
+optional-dependencies.jax = [
+  "jax>=0.4",
+]
 optional-dependencies.leiden = [
   "leidenalg",
   "spatialleiden>=0.4",
@@ -275,9 +278,18 @@ python_files = [ "test_*.py" ]
 testpaths = [ "tests/" ]
 addopts = [
   "--ignore=docs",
+  # Local convenience only. The real exclusion is the `-m "not reference"` baked into
+  # hatch.toml's run/run-cov, because that is a CLI flag and so survives any pytest that
+  # does not read this table. Keeping both means a bare `pytest` is also quiet.
+  "-m",
+  "not reference",
 ]
 markers = [
   "internet: tests that require internet",
+  """\
+  reference: numerical comparison against the vendored STalign reference (slow, needs JAX_ENABLE_X64=1, scheduled runs \
+  only)\
+  """,
   # the below markers are just there because we import from anndata.tests.helpers,
   # which means pytest sees them despite us not using them.
   "gpu: tests that require GPU",

src/squidpy/experimental/methods/__init__.py

@@ -0,0 +1,35 @@
+"""In-memory model-fitting core for experimental methods.
+
+The :mod:`.registry` subpackage holds the registry machinery and the :data:`ALIGN`
+registry; 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. 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.
+from squidpy.experimental.methods import align_landmarks, align_samples  # noqa: F401
+from squidpy.experimental.methods.registry import (
+    ALIGN,
+    AlignImagesFn,
+    AlignLandmarksFn,
+    AlignMethod,
+    AlignObsFn,
+    AlignResult,
+    Modality,
+    Registry,
+)
+
+__all__ = [
+    "ALIGN",
+    "AlignImagesFn",
+    "AlignLandmarksFn",
+    "AlignMethod",
+    "AlignObsFn",
+    "AlignResult",
+    "Modality",
+    "Registry",
+]

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`. 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",
+]