diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index a38c81478..b49fada8b 100644 --- a/.github/workflows/test.yaml +++ b/.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 diff --git a/docs/api.md b/docs/api.md index f8f159b87..f6c84d2b8 100644 --- a/docs/api.md +++ b/docs/api.md @@ -140,6 +140,25 @@ See the {doc}`extensibility guide ` 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 ` 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 diff --git a/docs/conf.py b/docs/conf.py index 0f413dfac..7f13ca59e 100644 --- a/docs/conf.py +++ b/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"), diff --git a/docs/release/notes-dev.md b/docs/release/notes-dev.md index 221023e5e..b6b1103f2 100644 --- a/docs/release/notes-dev.md +++ b/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/"` 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) diff --git a/hatch.toml b/hatch.toml index 1614be834..82defb133 100644 --- a/hatch.toml +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index e1ed678e1..face49285 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/squidpy/experimental/methods/__init__.py b/src/squidpy/experimental/methods/__init__.py new file mode 100644 index 000000000..c72dccb7e --- /dev/null +++ b/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", +] diff --git a/src/squidpy/experimental/methods/align_landmarks/__init__.py b/src/squidpy/experimental/methods/align_landmarks/__init__.py new file mode 100644 index 000000000..0a905e026 --- /dev/null +++ b/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", +] diff --git a/src/squidpy/experimental/methods/align_landmarks/_landmark.py b/src/squidpy/experimental/methods/align_landmarks/_landmark.py new file mode 100644 index 000000000..1f0a5dcf1 --- /dev/null +++ b/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 + + +@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.register("similarity", "landmarks") +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.register("affine", "landmarks") +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__}.") diff --git a/src/squidpy/experimental/methods/align_samples/__init__.py b/src/squidpy/experimental/methods/align_samples/__init__.py new file mode 100644 index 000000000..280fc8450 --- /dev/null +++ b/src/squidpy/experimental/methods/align_samples/__init__.py @@ -0,0 +1,18 @@ +"""``align_samples`` family: align two samples' point clouds (STalign). + +Importing this package registers the family's estimators into +:data:`~squidpy.experimental.methods.registry.ALIGN`. 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_image, + fit_stalign_obs, +) + +__all__ = ["fit_stalign_obs", "fit_stalign_image", "StalignResult"] diff --git a/src/squidpy/experimental/methods/align_samples/_stalign.py b/src/squidpy/experimental/methods/align_samples/_stalign.py new file mode 100644 index 000000000..af08e8b2b --- /dev/null +++ b/src/squidpy/experimental/methods/align_samples/_stalign.py @@ -0,0 +1,501 @@ +"""STalign estimator: JAX LDDMM point-cloud registration. + +Holds both the estimator adapter :func:`fit_stalign_obs` and its result type +:class:`StalignResult`; the pure numerics live under :mod:`._stalign_impl`. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal + +import numpy.typing as npt + +from squidpy.experimental.methods.registry import ALIGN + +if TYPE_CHECKING: + import jax + + JaxArray = jax.Array +else: # pragma: no cover - typing only + JaxArray = Any + + +@dataclass(slots=True) +class StalignResult: + """A fitted STalign diffeomorphism, ready to transform arbitrary points. + + :meth:`transform` works in ``(x, y)``; ``aligned_points`` is the fitted query + cloud already mapped into the reference frame. + """ + + affine: JaxArray + velocity: JaxArray + velocity_grid: tuple[JaxArray, JaxArray] + aligned_points: JaxArray + #: Row-col axes of the query and reference rasters the fit ran on, when it ran on + #: images. ``None`` for point-cloud fits, where no raster survives the call. + query_axes: tuple[JaxArray, JaxArray] | None = None + ref_axes: tuple[JaxArray, JaxArray] | None = None + match_weights: JaxArray | None = None + artifact_weights: JaxArray | None = None + background_weights: JaxArray | None = None + energies: JaxArray | None = None + n_iter: int | None = None + + def deformation_grid( + self, + *, + direction: Literal["forward", "backward"] = "forward", + query_axes: tuple[JaxArray, JaxArray] | None = None, + ref_axes: tuple[JaxArray, JaxArray] | None = None, + ) -> JaxArray: + """Return a dense row-column coordinate transform for visualisation. + + ``direction="forward"`` evaluates the query grid in the reference frame; + ``"backward"`` evaluates the reference grid in the query frame. The returned + array has shape ``(2, rows, columns)``. + """ + import jax.numpy as jnp + + from ._stalign_impl._core import _grid_points, _transform_grid_backward, transform_points_row_col + + source_axes = query_axes if query_axes is not None else self.query_axes + target_axes = ref_axes if ref_axes is not None else self.ref_axes + if source_axes is None or target_axes is None: + raise ValueError( + "This result was fitted on point clouds and carries no raster axes. " + "Pass both `query_axes=` and `ref_axes=`, or fit with " + "`align(in_='images/', ...)`." + ) + if direction == "backward": + return _transform_grid_backward(target_axes, self.velocity_grid, self.velocity, self.affine) + if direction != "forward": + raise ValueError(f"Expected `direction` to be 'forward' or 'backward', found {direction!r}.") + + grid = _grid_points(source_axes) + points = jnp.moveaxis(grid, 0, -1).reshape((-1, 2)) + transformed = transform_points_row_col( + self.velocity_grid, self.velocity, self.affine, points, direction="forward" + ) + return jnp.moveaxis(transformed.reshape((*grid.shape[1:], 2)), -1, 0) + + def warp_image( + self, + image: JaxArray, + *, + direction: Literal["forward", "backward"] = "forward", + query_axes: tuple[JaxArray, JaxArray] | None = None, + ref_axes: tuple[JaxArray, JaxArray] | None = None, + ) -> JaxArray: + """Resample an image through the fitted transformation. + + A diffeomorphism cannot be expressed as a SpatialData transformation -- the + available types are affine at most -- so an aligned image has to be materialised + rather than registered. ``direction="forward"`` maps a query-frame image onto + the reference grid; ``"backward"`` maps a reference-frame image onto the query + grid. Explicit axes allow results fitted from point clouds to warp their density + rasters without pretending those rasters are original image elements. + """ + import jax.numpy as jnp + + from ._stalign_impl._core import _interp, jax_dtype + + arr = jnp.asarray(image, dtype=jax_dtype()) + if arr.ndim == 2: + arr = arr[None] + if direction not in {"forward", "backward"}: + raise ValueError(f"Expected `direction` to be 'forward' or 'backward', found {direction!r}.") + source_axes = query_axes if query_axes is not None else self.query_axes + target_axes = ref_axes if ref_axes is not None else self.ref_axes + grid = self.deformation_grid( + direction="backward" if direction == "forward" else "forward", + query_axes=source_axes, + ref_axes=target_axes, + ) + sampling_axes = source_axes if direction == "forward" else target_axes + if sampling_axes is None: # guarded by deformation_grid; keeps the type checker honest + raise AssertionError("missing sampling axes") + return _interp(sampling_axes, arr, grid) + + def transform( + self, + points: JaxArray, + *, + direction: Literal["forward", "backward"] = "forward", + ) -> JaxArray: + """Map ``(N, 2)`` ``(x, y)`` points with the fitted diffeomorphism.""" + import jax.numpy as jnp + + from ._stalign_impl._core import jax_dtype, transform_points_row_col + + pts = jnp.asarray(points, dtype=jax_dtype()) + if pts.ndim != 2 or pts.shape[1] != 2: + raise ValueError(f"Expected an (N, 2) `(x, y)` array, found shape {pts.shape}.") + transformed_rc = transform_points_row_col( + self.velocity_grid, + self.velocity, + self.affine, + pts[:, ::-1], + direction=direction, + ) + return transformed_rc[:, ::-1] + + +@ALIGN.register("stalign", "obs", requires=("jax",)) +def fit_stalign_obs( + ref: npt.ArrayLike, + query: npt.ArrayLike, + *, + landmarks_source: npt.ArrayLike | None = None, + landmarks_target: npt.ArrayLike | None = None, + initial_affine: npt.ArrayLike | None = None, + initial_velocity: npt.ArrayLike | None = None, + velocity_grid: tuple[npt.ArrayLike, npt.ArrayLike] | None = None, + # rasterization + dx: float = 30.0, + blur: float | Sequence[float] = (2.0, 1.0, 0.5), + raster_expand: float = 1.1, + # LDDMM registration + a: float = 500.0, + p: float = 2.0, + expand: float = 2.0, + nt: int = 3, + niter: int = 5000, + diffeo_start: int = 0, + epL: float = 2e-8, + epT: float = 2e-1, + epV: float = 2e3, + sigmaM: float = 1.0, + sigmaB: float = 2.0, + sigmaA: float = 5.0, + sigmaR: float = 5e5, + sigmaP: float = 2e1, + muA: npt.ArrayLike | None = None, + muB: npt.ArrayLike | None = None, + tol: float | None = None, + patience: int = 25, +) -> StalignResult: + """Fit a deformation mapping ``query`` onto ``ref``. + + Parameters + ---------- + ref, query + ``(N, 2)`` / ``(M, 2)`` reference and query point clouds in ``(x, y)`` + order; the query is aligned onto the reference. Both are plain in-memory + arrays -- extracting them from an ``AnnData`` / ``SpatialData`` is the + caller's responsibility. + landmarks_source, landmarks_target + Optional corresponding ``(x, y)`` landmark arrays used to initialise the + affine. Must be provided together. + initial_affine + Optional homogeneous ``(3, 3)`` affine in public ``(x, y)`` coordinates. + Mutually exclusive with landmark initialisation. + initial_velocity, velocity_grid + Optional continuation state. The velocity has shape ``(nt, rows, columns, 2)``; + its components and the two grid axes use the solver's row-column convention. + dx, blur, raster_expand + Rasterization of the point clouds into density images: grid spacing, + Gaussian blur scale(s), and field-of-view padding factor. + a, p, expand, nt, niter, diffeo_start + LDDMM controls: kernel width ``a``, regularisation power ``p``, + velocity-grid padding ``expand``, number of integration time steps + ``nt``, iterations ``niter``, and the iteration at which the + diffeomorphic (non-affine) part starts updating ``diffeo_start``. + epL, epT, epV + Gradient-descent step sizes for the linear part, translation, and + velocity field. + sigmaM, sigmaB, sigmaA, sigmaR, sigmaP + Noise scales for the matching, background, artifact, regularisation, and + landmark-point terms of the objective. + muA, muB + Optional per-channel artifact and background means. ``None`` estimates the + corresponding mean during fitting, matching upstream STalign's default. + tol, patience + Stop once the objective's relative improvement over the last ``patience`` + iterations falls below ``tol``. ``tol=None`` (default) always runs ``niter``. + + Returns + ------- + A :class:`StalignResult` whose :meth:`~StalignResult.transform` maps + ``(x, y)`` points into the reference frame; ``aligned_points`` is the fitted + ``query`` already mapped. + + Notes + ----- + Runs in JAX's active float precision, which is **single** unless x64 is enabled. + The original STalign is double throughout, so results differ correspondingly. For + ``niter`` in the thousands, or a large ``sigmaR``, enable double precision before + importing JAX:: + + import jax + + jax.config.update("jax_enable_x64", True) + """ + # Import the JAX-backed solver only after the registry's requirements check + # passes, so callers without JAX get a clean ImportError rather than a + # confusing failure from a module-level `import jax`. + import jax.numpy as jnp + + from ._stalign_impl._core import jax_dtype, lddmm, transform_points_row_col + from ._stalign_impl._helpers import affine_from_points, rasterize_cloud, validate_points + + if (landmarks_source is None) != (landmarks_target is None): + raise ValueError("Expected both landmark arrays to be provided together.") + if initial_affine is not None and landmarks_source is not None: + raise ValueError("`initial_affine` is mutually exclusive with landmark initialisation.") + + # The solver runs internally in row-col (y, x); inputs are (x, y) -- swap at the boundary. + source_rc = validate_points(query, name="query")[:, ::-1] + target_rc = validate_points(ref, name="ref")[:, ::-1] + source_grid, source_image = rasterize_cloud(source_rc, dx=dx, blur=blur, expand=raster_expand) + target_grid, target_image = rasterize_cloud(target_rc, dx=dx, blur=blur, expand=raster_expand) + + dtype = jax_dtype() + if initial_affine is not None: + affine_xy = jnp.asarray(initial_affine, dtype=dtype) + if affine_xy.shape != (3, 3): + raise ValueError(f"Expected `initial_affine` to have shape (3, 3), found {affine_xy.shape}.") + swap = jnp.asarray([[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=dtype) + affine_rc = swap @ affine_xy @ swap + linear, translation = affine_rc[:2, :2], affine_rc[:2, 2] + src_lm = tgt_lm = None + elif landmarks_source is None: + linear, translation = jnp.eye(2, dtype=dtype), jnp.zeros(2, dtype=dtype) + src_lm = tgt_lm = None + else: + src_lm = validate_points(landmarks_source, name="landmarks_source")[:, ::-1] + tgt_lm = validate_points(landmarks_target, name="landmarks_target")[:, ::-1] + linear_np, translation_np = affine_from_points(src_lm, tgt_lm) + linear, translation = jnp.asarray(linear_np, dtype=dtype), jnp.asarray(translation_np, dtype=dtype) + + result = lddmm( + source_grid, + source_image, + target_grid, + target_image, + L=linear, + T=translation, + initial_velocity=initial_velocity, + velocity_grid=velocity_grid, + points_source=src_lm, + points_target=tgt_lm, + a=a, + p=p, + expand=expand, + nt=nt, + niter=niter, + diffeo_start=diffeo_start, + epL=epL, + epT=epT, + epV=epV, + sigmaM=sigmaM, + sigmaB=sigmaB, + sigmaA=sigmaA, + sigmaR=sigmaR, + sigmaP=sigmaP, + muA=muA, + muB=muB, + tol=tol, + patience=patience, + ) + aligned_rc = transform_points_row_col(result["xv"], result["v"], result["A"], source_rc, direction="forward") + return StalignResult( + affine=result["A"], + velocity=result["v"], + velocity_grid=result["xv"], + aligned_points=aligned_rc[:, ::-1], + match_weights=result["WM"], + artifact_weights=result["WA"], + background_weights=result["WB"], + energies=result["energies"], + n_iter=int(result["n_iter"]), + # No raster axes: the grids here are the internal density rasters at `dx` + # resolution, not a frame any real image lives on. Offering `warp_image` off them + # would quietly resample the caller's image onto a coarse, unrelated grid. + ) + + +@ALIGN.register("stalign", "images", requires=("jax",)) +def fit_stalign_image( + ref: npt.ArrayLike, + query: npt.ArrayLike, + *, + ref_scale: tuple[float, float] = (1.0, 1.0), + query_scale: tuple[float, float] = (1.0, 1.0), + ref_axes: tuple[npt.ArrayLike, npt.ArrayLike] | None = None, + query_axes: tuple[npt.ArrayLike, npt.ArrayLike] | None = None, + initial_affine: npt.ArrayLike | None = None, + initial_velocity: npt.ArrayLike | None = None, + velocity_grid: tuple[npt.ArrayLike, npt.ArrayLike] | None = None, + # LDDMM registration + a: float = 20.0, + p: float = 2.0, + expand: float = 2.0, + nt: int = 3, + niter: int = 200, + diffeo_start: int = 100, + epL: float = 2e-8, + epT: float = 2e-1, + epV: float = 1.0, + sigmaM: float = 1.0, + sigmaB: float = 2.0, + sigmaA: float = 5.0, + sigmaR: float = 5e5, + sigmaP: float = 2e1, + muA: npt.ArrayLike | None = None, + muB: npt.ArrayLike | None = None, + tol: float | None = None, + patience: int = 25, +) -> StalignResult: + """Fit a deformation mapping the ``query`` image onto the ``ref`` image. + + Parameters + ---------- + ref, query + Channels-first ``(c, y, x)`` rasters (a bare ``(y, x)`` array is promoted). The + query is aligned onto the reference; they need not share a shape. + ref_scale, query_scale + Physical size of one pixel as ``(y, x)``. Defaults to pixel units. Pass the + element's scale when the two images have different resolutions, otherwise the + fit is done in mismatched coordinates. + ref_axes, query_axes + Optional explicit physical row and column axes. Both pairs must be supplied; + they are mutually exclusive with non-unit ``ref_scale``/``query_scale``. + initial_affine + Optional homogeneous ``(3, 3)`` affine in public ``(x, y)`` coordinates. + initial_velocity, velocity_grid + Optional continuation state in the solver's row-column convention. + a, p, expand, nt, niter, diffeo_start + LDDMM controls, as in :func:`fit_stalign_obs`. Note ``a`` is a length in the *same* + units as ``ref_scale`` -- the default of 20 suits pixel units, where + :func:`fit_stalign_obs`'s 500 would exceed most images. ``diffeo_start`` defaults to + half of ``niter`` so the affine settles before the deformable part switches on; + starting both at once lets the velocity field absorb what is really a + translation, and it fits it worse than the affine would have. + epL, epT, epV + Gradient-descent step sizes for the linear part, translation, and velocity field. + These are **scale dependent**: they are tuned here for images in pixel units, so + a non-unit ``ref_scale`` will need them rescaled to match. ``epV`` is the one to + reach for first -- too large and the deformation overwhelms the affine. + sigmaM, sigmaB, sigmaA, sigmaR, sigmaP + Noise scales for the matching, background, artifact, regularisation, and + landmark-point terms of the objective. + muA, muB + Optional per-channel artifact and background means. ``None`` estimates the + corresponding mean during fitting, matching upstream STalign's default. + tol, patience + Stop once the objective's relative improvement over the last ``patience`` + iterations falls below ``tol``. ``tol=None`` (default) always runs ``niter``. + + Returns + ------- + A :class:`StalignResult`. Its :meth:`~StalignResult.transform` maps ``(x, y)`` points + in query pixel coordinates into the reference frame, and + :meth:`~StalignResult.warp_image` resamples a query image onto the reference grid. + """ + import jax.numpy as jnp + + from ._stalign_impl._core import jax_dtype, lddmm + + dtype = jax_dtype() + + def as_chw(image: npt.ArrayLike, name: str) -> JaxArray: + # Not `jnp.atleast_3d`: it appends the new axis, turning a (y, x) image into + # (y, x, 1) -- y channels of x by 1 -- instead of a single (1, y, x) channel. + arr = jnp.asarray(image, dtype=dtype) + if arr.ndim == 2: + return arr[None] + if arr.ndim != 3: + raise ValueError(f"Expected `{name}` to be a `(y, x)` or `(c, y, x)` image, found shape {arr.shape}.") + return arr + + source_image = as_chw(query, "query") + target_image = as_chw(ref, "ref") + if source_image.shape[0] != target_image.shape[0]: + raise ValueError( + f"Expected `ref` and `query` to have the same number of channels, found " + f"{target_image.shape[0]} and {source_image.shape[0]}." + ) + + def axes(image: JaxArray, scale: tuple[float, float]) -> tuple[JaxArray, JaxArray]: + # Row-col physical coordinates, centred so the affine initialises near identity. + rows, cols = image.shape[1], image.shape[2] + return ( + (jnp.arange(rows, dtype=dtype) - (rows - 1) / 2.0) * scale[0], + (jnp.arange(cols, dtype=dtype) - (cols - 1) / 2.0) * scale[1], + ) + + if (query_axes is None) != (ref_axes is None): + raise ValueError("Expected both `query_axes` and `ref_axes` to be provided together.") + + def explicit_axes(value: tuple[npt.ArrayLike, npt.ArrayLike], image: JaxArray, name: str): + resolved = (jnp.asarray(value[0], dtype=dtype), jnp.asarray(value[1], dtype=dtype)) + expected = image.shape[1:] + if resolved[0].ndim != 1 or resolved[1].ndim != 1 or tuple(map(len, resolved)) != expected: + raise ValueError(f"Expected `{name}` lengths {expected}, found {tuple(map(len, resolved))}.") + if len(resolved[0]) < 2 or len(resolved[1]) < 2: + raise ValueError(f"Expected each `{name}` axis to contain at least two coordinates.") + return resolved + + if query_axes is None: + source_grid = axes(source_image, query_scale) + target_grid = axes(target_image, ref_scale) + else: + if query_scale != (1.0, 1.0) or ref_scale != (1.0, 1.0): + raise ValueError("Explicit axes are mutually exclusive with non-unit image scales.") + source_grid = explicit_axes(query_axes, source_image, "query_axes") + target_grid = explicit_axes(ref_axes, target_image, "ref_axes") + + if initial_affine is None: + linear, translation = jnp.eye(2, dtype=dtype), jnp.zeros(2, dtype=dtype) + else: + affine_xy = jnp.asarray(initial_affine, dtype=dtype) + if affine_xy.shape != (3, 3): + raise ValueError(f"Expected `initial_affine` to have shape (3, 3), found {affine_xy.shape}.") + swap = jnp.asarray([[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=dtype) + affine_rc = swap @ affine_xy @ swap + linear, translation = affine_rc[:2, :2], affine_rc[:2, 2] + + result = lddmm( + source_grid, + source_image, + target_grid, + target_image, + L=linear, + T=translation, + initial_velocity=initial_velocity, + velocity_grid=velocity_grid, + a=a, + p=p, + expand=expand, + nt=nt, + niter=niter, + diffeo_start=diffeo_start, + epL=epL, + epT=epT, + epV=epV, + sigmaM=sigmaM, + sigmaB=sigmaB, + sigmaA=sigmaA, + sigmaR=sigmaR, + sigmaP=sigmaP, + muA=muA, + muB=muB, + tol=tol, + patience=patience, + ) + return StalignResult( + affine=result["A"], + velocity=result["v"], + velocity_grid=result["xv"], + aligned_points=jnp.zeros((0, 2), dtype=dtype), + query_axes=source_grid, + ref_axes=target_grid, + match_weights=result["WM"], + artifact_weights=result["WA"], + background_weights=result["WB"], + energies=result["energies"], + n_iter=int(result["n_iter"]), + ) diff --git a/src/squidpy/experimental/methods/align_samples/_stalign_impl/__init__.py b/src/squidpy/experimental/methods/align_samples/_stalign_impl/__init__.py new file mode 100644 index 000000000..26c72f7e2 --- /dev/null +++ b/src/squidpy/experimental/methods/align_samples/_stalign_impl/__init__.py @@ -0,0 +1,6 @@ +"""Ported STalign JAX LDDMM solver + +Pure numerics only and these functions are gated by the jax requirement. +""" + +from __future__ import annotations diff --git a/src/squidpy/experimental/methods/align_samples/_stalign_impl/_core.py b/src/squidpy/experimental/methods/align_samples/_stalign_impl/_core.py new file mode 100644 index 000000000..8100c9244 --- /dev/null +++ b/src/squidpy/experimental/methods/align_samples/_stalign_impl/_core.py @@ -0,0 +1,629 @@ +"""Core JAX implementation for experimental STalign point registration.""" + +from __future__ import annotations + +from functools import partial +from typing import Any, Literal + +import jax +import jax.numpy as jnp +import jax.scipy as jsp +import numpy as np + +__all__ = ["jax_dtype", "lddmm", "transform_points_row_col"] + +#: Iteration at which the mixture-weight E step switches on (STalign.py:1233). Before +#: this the weights are frozen at their initial values, so the objective changes +#: definition here and its value jumps discontinuously. +MIXTURE_E_STEP_START = 50 + + +def jax_dtype() -> jnp.dtype: + """Resolve the active JAX float dtype at call time, not import time.""" + return jnp.float64 if jax.config.x64_enabled else jnp.float32 + + +def _to_affine(linear: jax.Array, translation: jax.Array) -> jax.Array: + return jnp.array( + [ + [linear[0, 0], linear[0, 1], translation[0]], + [linear[1, 0], linear[1, 1], translation[1]], + [0.0, 0.0, 1.0], + ], + dtype=linear.dtype, + ) + + +def _grid_points(x: tuple[jax.Array, jax.Array]) -> jax.Array: + yy, xx = jnp.meshgrid(x[0], x[1], indexing="ij") + return jnp.stack((yy, xx)) + + +def _interp( + x: tuple[jax.Array, jax.Array], + image: jax.Array, + phii: jax.Array, + *, + mode: str = "nearest", +) -> jax.Array: + """Interpolate a channels-first image on physical row-column coordinates.""" + arr = jnp.asarray(image) + coords = jnp.asarray(phii) + if coords.shape[0] != 2: + raise ValueError(f"Expected interpolation coordinates to have leading axis of size 2, found `{coords.shape}`.") + + if arr.ndim == 2: + arr = arr[None, ...] + + row_step = x[0][1] - x[0][0] + col_step = x[1][1] - x[1][0] + row_idx = (coords[0] - x[0][0]) / row_step + col_idx = (coords[1] - x[1][0]) / col_step + idx = jnp.stack((row_idx.reshape(-1), col_idx.reshape(-1))) + + def _sample(channel: jax.Array) -> jax.Array: + values = jsp.ndimage.map_coordinates(channel, idx, order=1, mode=mode) + return values.reshape(coords.shape[1:]) + + return jax.vmap(_sample)(arr) + + +def transform_points_row_col( + xv: tuple[jax.Array, jax.Array], + velocity: jax.Array, + affine: jax.Array, + points: np.ndarray | jax.Array, + *, + direction: Literal["forward", "backward"] = "forward", +) -> jax.Array: + pts = jnp.asarray(points) + n_steps = velocity.shape[0] + time_steps = range(n_steps) + flow_sign = 1.0 + if direction == "backward": + affine = jnp.linalg.inv(affine) + pts = pts @ affine[:2, :2].T + affine[:2, -1] + flow_sign = -1.0 + time_steps = reversed(time_steps) + + for t in time_steps: + disp = _interp( + xv, + jnp.moveaxis(flow_sign * velocity[t], -1, 0), + pts.T[:, :, None], + mode="nearest", + )[:, :, 0].T + pts = pts + disp / n_steps + + if direction == "forward": + pts = pts @ affine[:2, :2].T + affine[:2, -1] + + return pts + + +def _transform_grid_backward( + x_target: tuple[jax.Array, jax.Array], + xv: tuple[jax.Array, jax.Array], + velocity: jax.Array, + affine: jax.Array, +) -> jax.Array: + target_grid = _grid_points(x_target) + affine_inv = jnp.linalg.inv(affine) + source_grid = jnp.einsum("ij,jhw->ihw", affine_inv[:2, :2], target_grid) + affine_inv[:2, -1][:, None, None] + + for t in range(velocity.shape[0] - 1, -1, -1): + disp = _interp(xv, jnp.moveaxis(-velocity[t], -1, 0), source_grid, mode="nearest") + source_grid = source_grid + disp / velocity.shape[0] + + return source_grid + + +def _contrast_transform(source_image: jax.Array, target_image: jax.Array, weights: jax.Array) -> jax.Array: + """Weighted ridge fit mapping source intensities onto target intensities. + + The coefficients are held constant with respect to the optimisation. This is an + expectation-maximisation M step, solved exactly at the current estimate, not a + quantity to descend on -- differentiating through the solve would silently turn the + alternating minimisation into a joint one and change the search direction. + """ + flat_source = source_image.reshape(source_image.shape[0], -1) + flat_target = target_image.reshape(target_image.shape[0], -1) + flat_weights = weights.reshape(-1) + + design = jnp.concatenate((jnp.ones((1, flat_source.shape[1]), dtype=source_image.dtype), flat_source), axis=0) + weighted_design = design * flat_weights[None, :] + design_cov = weighted_design @ design.T + target_cov = weighted_design @ flat_target.T + regularized = design_cov + 0.1 * jnp.eye(design_cov.shape[0], dtype=design_cov.dtype) + coefficients = jax.lax.stop_gradient(jnp.linalg.solve(regularized, target_cov)) + return (coefficients.T @ design).reshape(target_image.shape) + + +def _axis(start: float, stop: float, step: float) -> jax.Array: + """``step``-spaced samples covering ``[start, stop)``, with a stable length. + + ``jnp.arange(start, stop, step)`` on floats derives its length from the arguments by + floating-point division, so a ``stop`` that is itself a sum of floats can yield one + more or one fewer sample than intended. Taking the count first makes the length a + function of the interval alone. + """ + count = max(int(np.ceil((stop - start) / step)), 1) + return start + step * jnp.arange(count, dtype=jax_dtype()) + + +def _build_velocity_grid( + x_source: tuple[jax.Array, jax.Array], *, a: float, expand: float +) -> tuple[jax.Array, jax.Array]: + minimum = np.array([x_source[0][0], x_source[1][0]], dtype=float) + maximum = np.array([x_source[0][-1], x_source[1][-1]], dtype=float) + center = (minimum + maximum) / 2.0 + half_width = (maximum - minimum) * expand / 2.0 + step = a * 0.5 + return ( + _axis(center[0] - half_width[0], center[0] + half_width[0], step), + _axis(center[1] - half_width[1], center[1] + half_width[1], step), + ) + + +def _build_regularizer( + xv: tuple[jax.Array, jax.Array], + *, + a: float, + p: float, +) -> tuple[jax.Array, jax.Array, float | jax.Array]: + dv = jnp.array([xv[0][1] - xv[0][0], xv[1][1] - xv[1][0]]) + shape = (xv[0].shape[0], xv[1].shape[0]) + fy = jnp.arange(shape[0], dtype=xv[0].dtype) / (shape[0] * dv[0]) + fx = jnp.arange(shape[1], dtype=xv[1].dtype) / (shape[1] * dv[1]) + frequency_grid = jnp.stack(jnp.meshgrid(fy, fx, indexing="ij"), axis=-1) + ll = (1.0 + 2.0 * a**2 * jnp.sum((1.0 - jnp.cos(2.0 * np.pi * frequency_grid * dv)) / (dv**2), axis=-1)) ** ( + 2.0 * p + ) + kernel = 1.0 / ll + dv_prod = jnp.prod(dv) + return kernel, ll, dv_prod + + +def _update_mixture_weights( + transformed_source: jax.Array, + target_image: jax.Array, + match_weights: jax.Array, + artifact_weights: jax.Array, + background_weights: jax.Array, + *, + sigmaM: float, + sigmaA: float, + sigmaB: float, + estimate_muA: bool, + estimate_muB: bool, + muA: jax.Array, + muB: jax.Array, + iteration: int, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + if estimate_muA: + muA = jnp.sum(artifact_weights * target_image, axis=(-1, -2)) / jnp.maximum(jnp.sum(artifact_weights), 1e-12) + if estimate_muB: + muB = jnp.sum(background_weights * target_image, axis=(-1, -2)) / jnp.maximum( + jnp.sum(background_weights), 1e-12 + ) + + def _e_step() -> tuple[jax.Array, jax.Array, jax.Array]: + weights = jnp.stack((match_weights, artifact_weights, background_weights)) + mixing = jnp.sum(weights, axis=(1, 2)) + mixing = mixing + jnp.max(mixing) * 1e-6 + mixing = mixing / jnp.sum(mixing) + + n_channels = target_image.shape[0] + norm_match = (2.0 * np.pi * sigmaM**2) ** (n_channels / 2.0) + norm_artifact = (2.0 * np.pi * sigmaA**2) ** (n_channels / 2.0) + norm_background = (2.0 * np.pi * sigmaB**2) ** (n_channels / 2.0) + + match = mixing[0] * jnp.exp(-jnp.sum((transformed_source - target_image) ** 2, axis=0) / (2.0 * sigmaM**2)) + match = match / norm_match + artifact = mixing[1] * jnp.exp(-jnp.sum((muA[:, None, None] - target_image) ** 2, axis=0) / (2.0 * sigmaA**2)) + artifact = artifact / norm_artifact + background = mixing[2] * jnp.exp(-jnp.sum((muB[:, None, None] - target_image) ** 2, axis=0) / (2.0 * sigmaB**2)) + background = background / norm_background + + total = match + artifact + background + total = total + jnp.max(total) * 1e-6 + return match / total, artifact / total, background / total + + # Before the E step switches on the weights stay at their initial 0.5/0.4/0.1, + # while muA/muB are still re-estimated every 5th iteration. + match_weights, artifact_weights, background_weights = jax.lax.cond( + iteration >= MIXTURE_E_STEP_START, + _e_step, + lambda: (match_weights, artifact_weights, background_weights), + ) + return match_weights, artifact_weights, background_weights, muA, muB + + +def _lddmm_loss( + linear: jax.Array, + translation: jax.Array, + velocity: jax.Array, + *, + x_source: tuple[jax.Array, jax.Array], + source_image: jax.Array, + x_target: tuple[jax.Array, jax.Array], + target_image: jax.Array, + xv: tuple[jax.Array, jax.Array], + match_weights: jax.Array, + ll: jax.Array, + dv_prod: float | jax.Array, + points_source: jax.Array, + points_target: jax.Array, + sigmaM: float, + sigmaR: float, + sigmaP: float, +) -> tuple[jax.Array, tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]]: + affine = _to_affine(linear, translation) + source_grid = _transform_grid_backward(x_target, xv, velocity, affine) + warped_source = _interp(x_source, source_image, source_grid, mode="nearest") + contrast_source = _contrast_transform(warped_source, target_image, match_weights) + + match_energy = jnp.sum((contrast_source - target_image) ** 2 * match_weights) / (2.0 * sigmaM**2) + fft_velocity = jnp.fft.fftn(velocity, axes=(1, 2)) + reg_energy = ( + jnp.sum(jnp.sum(jnp.abs(fft_velocity) ** 2, axis=(0, 3)) * ll) + * dv_prod + / 2.0 + / velocity.shape[1] + / velocity.shape[2] + / sigmaR**2 + ) + + transformed_points = transform_points_row_col(xv, velocity, affine, points_source, direction="forward") + if points_source.shape[0] == 0: + point_energy = jnp.array(0.0, dtype=source_image.dtype) + else: + point_energy = jnp.sum((transformed_points - points_target) ** 2) / (2.0 * sigmaP**2) + + total = match_energy + reg_energy + point_energy + return total, (contrast_source, transformed_points, match_energy, reg_energy, point_energy) + + +@partial( + jax.jit, + static_argnames=( + "niter", + "diffeo_start", + "epL", + "epT", + "epV", + "sigmaM", + "sigmaA", + "sigmaB", + "sigmaR", + "sigmaP", + "tol", + "patience", + "estimate_muA", + "estimate_muB", + ), +) +def _lddmm_run( + linear, + translation, + velocity, + match_weights, + artifact_weights, + background_weights, + muA, + muB, + *, + x_source, + source_image, + x_target, + target_image, + xv, + kernel, + ll, + dv_prod, + source_landmarks, + target_landmarks, + niter, + diffeo_start, + epL, + epT, + epV, + sigmaM, + sigmaA, + sigmaB, + sigmaR, + sigmaP, + tol, + patience, + estimate_muA, + estimate_muB, +): + """The gradient descent, as one compiled loop. + + Jitted as a whole rather than per-iteration: `lax.while_loop` outside a `jit` re-traces + its body on every call, and tracing `value_and_grad` through the interpolation and FFTs + costs about as much as a thousand iterations of actually running it. + """ + loss_and_grad = jax.value_and_grad(_lddmm_loss, argnums=(0, 1, 2), has_aux=True) + + # Precomputed in Python so the two step sizes are bit-identical to `epL / 1.0` and + # `epL / 10.0` -- the `(it >= diffeo_start) * 9` scaling at STalign.py:1205-1206. + steps_before = (epL, epT) + steps_after = (epL / 10.0, epT / 10.0) + + dtype = jax_dtype() + # `niter=0` means "evaluate the initial state and stop"; the trace still needs a + # slot so the carry has a fixed shape. + energies = jnp.full((max(niter, 1),), jnp.nan, dtype=dtype) + initial = ( + jnp.asarray(0), + linear, + translation, + velocity, + match_weights, + artifact_weights, + background_weights, + muA, + muB, + jnp.asarray(jnp.nan, dtype=dtype), + source_landmarks, + energies, + ) + + def _step(carry: tuple[Any, ...]) -> tuple[Any, ...]: + iteration, linear, translation, velocity, wm, wa, wb, muA, muB, _, _, energies = carry + + (energy, aux), (grad_linear, grad_translation, grad_velocity) = loss_and_grad( + linear, + translation, + velocity, + x_source=x_source, + source_image=source_image, + x_target=x_target, + target_image=target_image, + xv=xv, + match_weights=wm, + ll=ll, + dv_prod=dv_prod, + points_source=source_landmarks, + points_target=target_landmarks, + sigmaM=sigmaM, + sigmaR=sigmaR, + sigmaP=sigmaP, + ) + contrast_source, transformed_points, _, _, _ = aux + + diffeo = iteration >= diffeo_start + step_linear = jnp.where(diffeo, steps_after[0], steps_before[0]) + step_translation = jnp.where(diffeo, steps_after[1], steps_before[1]) + linear = linear - step_linear * grad_linear + translation = translation - step_translation * grad_translation + + grad_velocity = jnp.fft.ifftn( + jnp.fft.fftn(grad_velocity, axes=(1, 2)) * kernel[None, ..., None], + axes=(1, 2), + ).real + velocity = jnp.where(diffeo, velocity - epV * grad_velocity, velocity) + + wm, wa, wb, muA, muB = jax.lax.cond( + iteration % 5 == 0, + lambda: _update_mixture_weights( + contrast_source, + target_image, + wm, + wa, + wb, + sigmaM=sigmaM, + sigmaA=sigmaA, + sigmaB=sigmaB, + estimate_muA=estimate_muA, + estimate_muB=estimate_muB, + muA=muA, + muB=muB, + iteration=iteration, + ), + lambda: (wm, wa, wb, muA, muB), + ) + return ( + iteration + 1, + linear, + translation, + velocity, + wm, + wa, + wb, + muA, + muB, + energy, + transformed_points, + energies.at[iteration].set(energy), + ) + + def _keep_going(carry: tuple[Any, ...]) -> jax.Array: + iteration, energies = carry[0], carry[-1] + if tol is None: + return iteration < niter + # Compare against `patience` iterations ago rather than the previous step: the + # weights only move every 5th iteration, so consecutive energies plateau and + # then jump, and a one-step test would stop on the plateau. + recent = energies[jnp.maximum(iteration - 1, 0)] + older = energies[jnp.maximum(iteration - 1 - patience, 0)] + improving = (older - recent) > tol * jnp.abs(older) + # The whole comparison window has to sit after the E step switches on. The + # objective changes definition at `MIXTURE_E_STEP_START` and its value jumps + # upward there, which reads as "no longer improving" to any stopping rule; the + # first energy computed with the new weights is at `MIXTURE_E_STEP_START + 1`, so + # the oldest index we may look at is that, hence the `+ 2` once the window and the + # one-step lag are accounted for. + warming_up = iteration < MIXTURE_E_STEP_START + 2 + patience + return (iteration < niter) & (warming_up | improving) + + return jax.lax.while_loop(_keep_going, _step, initial) + + +def lddmm( + xI: tuple[np.ndarray | jax.Array, np.ndarray | jax.Array], + I: np.ndarray | jax.Array, + xJ: tuple[np.ndarray | jax.Array, np.ndarray | jax.Array], + J: np.ndarray | jax.Array, + *, + L: np.ndarray | jax.Array, + T: np.ndarray | jax.Array, + initial_velocity: np.ndarray | jax.Array | None = None, + velocity_grid: tuple[np.ndarray | jax.Array, np.ndarray | jax.Array] | None = None, + points_source: np.ndarray | jax.Array | None = None, + points_target: np.ndarray | jax.Array | None = None, + a: float = 500.0, + p: float = 2.0, + expand: float = 2.0, + nt: int = 3, + niter: int = 5000, + diffeo_start: int = 0, + epL: float = 2e-8, + epT: float = 2e-1, + epV: float = 2e3, + sigmaM: float = 1.0, + sigmaB: float = 2.0, + sigmaA: float = 5.0, + sigmaR: float = 5e5, + sigmaP: float = 2e1, + muA: np.ndarray | jax.Array | None = None, + muB: np.ndarray | jax.Array | None = None, + tol: float | None = None, + patience: int = 25, +) -> dict[str, Any]: + """Fit an LDDMM registration of ``I`` onto ``J`` by gradient descent. + + The whole descent runs as a single ``lax.while_loop``, so XLA compiles the loop body + once and fuses across it instead of paying per-iteration dispatch from Python. + + Parameters + ---------- + tol + Stop once the objective's relative improvement over the last ``patience`` + iterations falls below this. ``None`` (default) always runs the full ``niter``. + patience + Window for the ``tol`` test. Compared against ``patience`` iterations ago rather + than the previous step because the mixture weights only move every 5th iteration, + so the objective plateaus and then jumps -- a one-step test would stop on a + plateau. + + Returns + ------- + A dict with the fitted ``A``/``v``/``xv``, the mixture weights, the final energy + ``E``, the per-iteration ``energies`` trace, and ``n_iter`` actually run. + """ + x_source = (jnp.asarray(xI[0]), jnp.asarray(xI[1])) + x_target = (jnp.asarray(xJ[0]), jnp.asarray(xJ[1])) + source_image = jnp.asarray(I, dtype=jax_dtype()) + target_image = jnp.asarray(J, dtype=jax_dtype()) + linear = jnp.asarray(L, dtype=jax_dtype()) + translation = jnp.asarray(T, dtype=jax_dtype()) + + if points_source is None: + source_landmarks = jnp.zeros((0, 2), dtype=jax_dtype()) + target_landmarks = jnp.zeros((0, 2), dtype=jax_dtype()) + else: + source_landmarks = jnp.asarray(points_source, dtype=jax_dtype()) + target_landmarks = jnp.asarray(points_target, dtype=jax_dtype()) + + if (initial_velocity is None) != (velocity_grid is None): + raise ValueError("Expected `initial_velocity` and `velocity_grid` to be provided together.") + if velocity_grid is None: + xv = _build_velocity_grid(x_source, a=a, expand=expand) + velocity = jnp.zeros((nt, xv[0].shape[0], xv[1].shape[0], 2), dtype=jax_dtype()) + else: + xv = (jnp.asarray(velocity_grid[0], dtype=jax_dtype()), jnp.asarray(velocity_grid[1], dtype=jax_dtype())) + velocity = jnp.asarray(initial_velocity, dtype=jax_dtype()) + if velocity.ndim != 4: + raise ValueError(f"Expected `initial_velocity` to be four-dimensional, found {velocity.shape}.") + expected = (velocity.shape[0], xv[0].shape[0], xv[1].shape[0], 2) + if velocity.shape != expected: + raise ValueError( + "Expected `initial_velocity` to have shape " + f"(nt, {xv[0].shape[0]}, {xv[1].shape[0]}, 2), found {velocity.shape}." + ) + kernel, ll, dv_prod = _build_regularizer(xv, a=a, p=p) + + match_weights = jnp.full(target_image.shape[1:], 0.5, dtype=target_image.dtype) + background_weights = jnp.full(target_image.shape[1:], 0.4, dtype=target_image.dtype) + artifact_weights = jnp.full(target_image.shape[1:], 0.1, dtype=target_image.dtype) + n_channels = target_image.shape[0] + + def mixture_mean(value: np.ndarray | jax.Array | None, *, name: str, default: jax.Array) -> jax.Array: + if value is None: + return default + mean = jnp.asarray(value, dtype=target_image.dtype) + if mean.shape != (n_channels,): + raise ValueError(f"Expected `{name}` to have shape ({n_channels},), found {mean.shape}.") + return mean + + artifact_mean = mixture_mean(muA, name="muA", default=jnp.mean(target_image, axis=(1, 2))) + background_mean = mixture_mean(muB, name="muB", default=jnp.zeros(n_channels, dtype=target_image.dtype)) + estimate_muA = muA is None + estimate_muB = muB is None + + final = _lddmm_run( + linear, + translation, + velocity, + match_weights, + artifact_weights, + background_weights, + artifact_mean, + background_mean, + x_source=x_source, + source_image=source_image, + x_target=x_target, + target_image=target_image, + xv=xv, + kernel=kernel, + ll=ll, + dv_prod=dv_prod, + source_landmarks=source_landmarks, + target_landmarks=target_landmarks, + niter=niter, + diffeo_start=diffeo_start, + epL=epL, + epT=epT, + epV=epV, + sigmaM=sigmaM, + sigmaA=sigmaA, + sigmaB=sigmaB, + sigmaR=sigmaR, + sigmaP=sigmaP, + tol=tol, + patience=patience, + estimate_muA=estimate_muA, + estimate_muB=estimate_muB, + ) + ( + completed, + linear, + translation, + velocity, + match_weights, + artifact_weights, + background_weights, + muA, + muB, + energy, + transformed_points, + energies, + ) = final + + affine = _to_affine(linear, translation) + return { + "A": affine, + "v": velocity, + "xv": xv, + "WM": match_weights, + "WB": background_weights, + "WA": artifact_weights, + "E": energy, + "points": transformed_points, + # Per-iteration objective, so a caller can tell a converged run from a diverged + # one without running it again. Trailing entries stay NaN if `tol` stopped early. + "energies": energies[:niter], + "n_iter": completed, + } diff --git a/src/squidpy/experimental/methods/align_samples/_stalign_impl/_helpers.py b/src/squidpy/experimental/methods/align_samples/_stalign_impl/_helpers.py new file mode 100644 index 000000000..b26fab863 --- /dev/null +++ b/src/squidpy/experimental/methods/align_samples/_stalign_impl/_helpers.py @@ -0,0 +1,174 @@ +"""Numeric helpers for STalign point-cloud registration.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import jax.numpy as jnp +import numpy as np + +from ._core import jax_dtype + +if TYPE_CHECKING: + import jax + + JaxArray = jax.Array +else: # pragma: no cover - typing only + JaxArray = Any + +__all__ = [ + "affine_from_points", + "rasterize", + "rasterize_cloud", + "validate_points", +] + + +def rasterize_cloud( + points_rc: JaxArray, *, dx: float, blur: float | list[float], expand: float +) -> tuple[tuple[JaxArray, JaxArray], JaxArray]: + """Rasterize a row-col cloud into a ``((grid_y, grid_x), image)`` density.""" + grid_x, grid_y, image = rasterize(points_rc[:, 1], points_rc[:, 0], dx=dx, blur=blur, expand=expand) + return (grid_y, grid_x), image + + +def validate_points(points: Any, *, name: str) -> JaxArray: + """Coerce ``points`` to a finite ``(n, 2)`` JAX array.""" + arr = jnp.asarray(points, dtype=jax_dtype()) + if arr.ndim != 2 or arr.shape[1] != 2: + raise ValueError(f"Expected `{name}` to have shape `(n, 2)`, found `{arr.shape}`.") + if not bool(jnp.all(jnp.isfinite(arr))): + raise ValueError(f"Expected `{name}` to contain only finite values.") + return arr + + +def _axis(start: float, stop: float, step: float) -> np.ndarray: + """``step``-spaced samples covering ``[start, stop)``, with a stable length. + + ``np.arange`` on floats derives its length from the arguments by floating-point + division, so a ``stop`` that is itself a sum of floats can yield one more or one + fewer sample than intended. Taking the count first makes the length a function of + the interval alone. + """ + count = max(int(np.ceil((stop - start) / step)), 1) + return start + step * np.arange(count, dtype=float) + + +def rasterize( + x: np.ndarray, + y: np.ndarray, + *, + dx: float = 30.0, + blur: float | list[float] = 1.0, + expand: float = 1.1, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Rasterize a point cloud into a multi-scale Gaussian density image. + + Each point deposits unit mass bilinearly across its four neighbouring cells of a + regular ``dx``-spaced grid; every ``blur`` scale is then an isotropic Gaussian blur + of that histogram, so each point becomes a unit-integral Gaussian. ``blur`` sets the + kernel width: ``sigma = 2 * blur`` pixels, i.e. ``2 * blur * dx`` in physical units. + Total mass is preserved exactly, including for points near the border. + """ + x = np.asarray(x, dtype=float).reshape(-1) + y = np.asarray(y, dtype=float).reshape(-1) + if x.shape != y.shape: + raise ValueError("Expected `x` and `y` to be 1D arrays with the same length.") + if x.size == 0: + raise ValueError("Expected at least one point to rasterize.") + if dx <= 0: + raise ValueError("Expected `dx` to be positive.") + if expand <= 0: + raise ValueError("Expected `expand` to be positive.") + + blur_values = np.atleast_1d(np.asarray(blur, dtype=float)) + if blur_values.ndim != 1 or np.any(blur_values <= 0): + raise ValueError("Expected `blur` to be a positive scalar or a 1D sequence of positive values.") + + min_x = float(np.min(x)) + max_x = float(np.max(x)) + min_y = float(np.min(y)) + max_y = float(np.max(y)) + + center_x = (min_x + max_x) / 2.0 + center_y = (min_y + max_y) / 2.0 + half_x = (max_x - min_x) * expand / 2.0 + half_y = (max_y - min_y) * expand / 2.0 + + grid_x = _axis(center_x - half_x, center_x + half_x, dx) + grid_y = _axis(center_y - half_y, center_y + half_y, dx) + if grid_x.size < 2 or grid_y.size < 2: + raise ValueError("Rasterized grid is too small. Increase the point spread or lower `dx`.") + + histogram = _deposit(x, y, grid_x, grid_y, dx) + out = np.stack([_blur_conserving(histogram, sigma=2.0 * float(b)) for b in blur_values]) + return grid_x, grid_y, out + + +def _deposit(x: np.ndarray, y: np.ndarray, grid_x: np.ndarray, grid_y: np.ndarray, dx: float) -> np.ndarray: + """Spread each point's unit mass bilinearly over its four neighbouring cells. + + Snapping to the nearest cell instead would quantise every position by up to half a + cell before any blurring happens, which at a typical ``dx`` is comparable to the + features being registered. ``np.bincount`` keeps this a handful of vectorised passes + rather than a Python loop over points. + """ + n_rows, n_cols = grid_y.size, grid_x.size + col_f = (x - grid_x[0]) / dx + row_f = (y - grid_y[0]) / dx + col_0 = np.floor(col_f).astype(np.intp) + row_0 = np.floor(row_f).astype(np.intp) + col_w = col_f - col_0 + row_w = row_f - row_0 + + flat = np.zeros(n_rows * n_cols, dtype=float) + for row_offset, row_weight in ((0, 1.0 - row_w), (1, row_w)): + for col_offset, col_weight in ((0, 1.0 - col_w), (1, col_w)): + rows, cols = row_0 + row_offset, col_0 + col_offset + inside = (rows >= 0) & (rows < n_rows) & (cols >= 0) & (cols < n_cols) + flat += np.bincount( + rows[inside] * n_cols + cols[inside], + weights=(row_weight * col_weight)[inside], + minlength=flat.size, + ) + return flat.reshape(n_rows, n_cols) + + +def _blur_conserving(histogram: np.ndarray, *, sigma: float) -> np.ndarray: + """Gaussian blur that keeps every point's mass on the grid. + + ``mode="constant"`` lets a kernel centred near the border spill off the edge, so + points there contribute less than one unit and the density is biased low around the + rim. The mass a point at cell ``c`` retains is ``sum_p K(p - c)``, which by symmetry + of ``K`` equals ``gaussian_filter(ones)[c]`` -- dividing by that before blurring + makes the total exactly the number of points, wherever they lie. + """ + from scipy.ndimage import gaussian_filter + + retained = gaussian_filter(np.ones_like(histogram), sigma=sigma, mode="constant") + return gaussian_filter(histogram / retained, sigma=sigma, mode="constant") + + +def affine_from_points( + points_source: JaxArray, + points_target: JaxArray, +) -> tuple[np.ndarray, np.ndarray]: + """Compute an affine initialization from corresponding landmarks.""" + source = np.asarray(points_source, dtype=float) + target = np.asarray(points_target, dtype=float) + if source.shape != target.shape: + raise ValueError( + f"Expected `points_source` and `points_target` to have the same shape, found " + f"`{source.shape}` and `{target.shape}`." + ) + + if source.shape[0] < 3: + linear = np.eye(2, dtype=float) + translation = np.mean(target, axis=0) - np.mean(source, axis=0) + return linear, translation + + from skimage.transform import estimate_transform + + model_obj = estimate_transform("affine", src=source, dst=target) + affine = np.asarray(model_obj.params) + return affine[:2, :2], affine[:2, -1] diff --git a/src/squidpy/experimental/methods/registry/__init__.py b/src/squidpy/experimental/methods/registry/__init__.py new file mode 100644 index 000000000..d67f452ab --- /dev/null +++ b/src/squidpy/experimental/methods/registry/__init__.py @@ -0,0 +1,196 @@ +"""The registry machinery and the align family it powers. + +This module holds three things that belong together: + +* :class:`Registry` -- a ``name -> AlignMethod`` map for one method *family*. +* :class:`AlignMethod` -- what one named method can align. A method fills the slots it + implements (``obs``, ``images``, ``landmarks``) and leaves the rest ``None``, so + "``stalign`` cannot align landmarks" is data rather than a special case buried in a + dispatch site. +* The structural :class:`~typing.Protocol` contracts each slot advertises, so the public + API is typed against a contract rather than a concrete 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`. + +The three slots do **not** all take the same kind of input. ``obs`` and ``images`` take +the data being aligned; ``landmarks`` takes pre-paired correspondences that annotate it. +Each public entry point only ever reads its own slot, so no caller sees a mixed +convention, but the distinction is why these are named slots rather than one callable +with a mode flag. +""" + +from __future__ import annotations + +import functools +import importlib.util +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Literal, Protocol, get_args, 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__ = [ + "ALIGN", + "AlignImagesFn", + "AlignLandmarksFn", + "AlignMethod", + "AlignObsFn", + "AlignResult", + "Modality", + "Registry", +] + +#: What an alignment can be driven by. Also the slot names on :class:`AlignMethod`, and +#: the modalities :mod:`squidpy.experimental.tl._align._paths` resolves a path to. +Modality = Literal["obs", "images", "landmarks"] +MODALITIES: tuple[Modality, ...] = get_args(Modality) + + +@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`` function requires of an estimator's + result, so ``out=None`` 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 AlignObsFn(Protocol): + """Calling convention for the ``obs`` slot. + + Two point clouds in (by keyword as ``ref`` / ``query``, so the direction can never be + silently swapped), one :class:`AlignResult` out. + """ + + def __call__(self, ref: npt.ArrayLike, query: npt.ArrayLike, **kwargs: Any) -> AlignResult: ... + + +class AlignImagesFn(Protocol): + """Calling convention for the ``images`` slot: two ``(c, y, x)`` rasters in.""" + + def __call__(self, ref: npt.ArrayLike, query: npt.ArrayLike, **kwargs: Any) -> AlignResult: ... + + +class AlignLandmarksFn(Protocol): + """Calling convention for the ``landmarks`` slot: paired correspondences in, affine out. + + Unlike the other two slots, ``ref`` / ``query`` here are *not* the data being aligned + -- they are ``(N, 2)`` landmark arrays annotating it, matched by row order. + """ + + def __call__( + self, + ref: npt.ArrayLike, + query: npt.ArrayLike, + *, + source_cs: str | None = ..., + target_cs: str | None = ..., + ) -> AffineFitResult: ... + + +@dataclass(frozen=True, slots=True) +class AlignMethod: + """One named alignment method and the modalities it implements.""" + + name: str + obs: AlignObsFn | None = None + images: AlignImagesFn | None = None + landmarks: AlignLandmarksFn | None = None + + def supports(self) -> tuple[Modality, ...]: + """Modalities this method implements.""" + return tuple(m for m in MODALITIES if getattr(self, m) is not None) + + def implementation(self, modality: Modality) -> Callable[..., Any]: + """Return the estimator for ``modality``, or explain what this method does instead.""" + fn = getattr(self, modality) + if fn is None: + supported = ", ".join(self.supports()) + raise ValueError(f"Method {self.name!r} does not support {modality} alignment. It supports: {supported}.") + return fn + + +class Registry: + """A ``name -> AlignMethod`` registry for one *family* of methods. + + Registration is per slot, so a method that gains a modality later does not have to be + declared in one place:: + + @ALIGN.register("stalign", "obs", requires=("jax",)) + def fit_stalign_obs(...): ... + + @ALIGN.register("stalign", "images", requires=("jax",)) + def fit_stalign_image(...): ... + """ + + def __init__(self, name: str) -> None: + self.name = name + self._registry: dict[str, AlignMethod] = {} + + def register( + self, + key: str, + modality: Modality, + *, + requires: tuple[str, ...] = (), + ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + """Return a decorator registering a function as ``key``'s ``modality`` estimator.""" + if modality not in MODALITIES: + raise ValueError(f"Unknown modality {modality!r}. Expected one of {', '.join(MODALITIES)}.") + + def decorator(func: Callable[..., Any]) -> Callable[..., Any]: + existing = self._registry.get(key, AlignMethod(name=key)) + if getattr(existing, modality) is not None: + raise ValueError(f"Method {key!r} already has a {modality} estimator in the {self.name!r} registry.") + + registered = func + 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) + + registered = wrapped + + self._registry[key] = replace(existing, **{modality: registered}) + return registered + + return decorator + + def get(self, key: str) -> AlignMethod: + """Return the method 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 every registered method name.""" + return tuple(self._registry) + + def supporting(self, modality: Modality) -> tuple[str, ...]: + """Return the method names implementing ``modality``, for docs and error messages.""" + return tuple(sorted(key for key, m in self._registry.items() if getattr(m, modality) is not None)) + + +#: Alignment estimators, keyed by method name. Each declares which of ``obs`` / ``images`` +#: / ``landmarks`` it can align. Consumed by ``squidpy.experimental.tl.align``. +ALIGN: Registry = Registry("align") diff --git a/src/squidpy/experimental/tl/__init__.py b/src/squidpy/experimental/tl/__init__.py index 7122bd3cd..bb7b1d0f4 100644 --- a/src/squidpy/experimental/tl/__init__.py +++ b/src/squidpy/experimental/tl/__init__.py @@ -1,6 +1,21 @@ from __future__ import annotations +# `AlignResult` is the only result type on the public surface: it is the estimator +# contract (a `transform` mapping points into the reference frame) and the declared +# return of `align`. The concrete results (`StalignResult`, +# `AffineFitResult`) stay in their home modules under `squidpy.experimental.methods` +# for callers that need raw fields -- the public API stays method-agnostic. +from squidpy.experimental.methods import AlignResult + +from ._align import align from ._tiling_qc import TilingQCParams, calculate_tiling_qc from ._tiling_stitch import StitchParams, assign_stitch_groups -__all__ = ["StitchParams", "TilingQCParams", "assign_stitch_groups", "calculate_tiling_qc"] +__all__ = [ + "align", + "calculate_tiling_qc", + "TilingQCParams", + "AlignResult", + "StitchParams", + "assign_stitch_groups", +] diff --git a/src/squidpy/experimental/tl/_align/__init__.py b/src/squidpy/experimental/tl/_align/__init__.py new file mode 100644 index 000000000..e01742a4f --- /dev/null +++ b/src/squidpy/experimental/tl/_align/__init__.py @@ -0,0 +1,7 @@ +"""Public alignment API for :mod:`squidpy.experimental.tl`.""" + +from __future__ import annotations + +from squidpy.experimental.tl._align._api import align + +__all__ = ["align"] diff --git a/src/squidpy/experimental/tl/_align/_api.py b/src/squidpy/experimental/tl/_align/_api.py new file mode 100644 index 000000000..4ee944fdf --- /dev/null +++ b/src/squidpy/experimental/tl/_align/_api.py @@ -0,0 +1,325 @@ +"""The public alignment function, built on the :mod:`squidpy.experimental.methods` core. + +A thin orchestrator: resolve ``in_`` to in-memory arrays, dispatch to a fit-core +estimator, write the result back at ``out``. Path resolution lives in :mod:`._paths` and +transformation write-back in :mod:`._io`; the estimators themselves never see a container. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, TypeVar + +import numpy as np +from anndata import AnnData +from spatialdata import SpatialData + +from squidpy.experimental.methods import ALIGN +from squidpy.experimental.methods.registry import MODALITIES +from squidpy.experimental.tl._align._io import shallow_copy_sdata, writeback_affine_sdata +from squidpy.experimental.tl._align._paths import DataPath, parse_path, read_path, write_path + +if TYPE_CHECKING: + from squidpy.experimental.methods import AlignResult, Modality, Registry + +__all__ = ["align"] + +F = TypeVar("F", bound="Callable[..., Any]") + +#: Method used when ``method`` is left unset, per modality. Point clouds and rasters both +#: go to the diffeomorphic solver; landmarks default to the more constrained of the two +#: closed-form fits, since a similarity cannot shear a sample that should not be sheared. +DEFAULT_METHOD: dict[str, str] = {"obs": "stalign", "images": "stalign", "landmarks": "similarity"} + + +def _resolve_in(in_: str | tuple[str, str]) -> tuple[DataPath, DataPath]: + """Normalise ``in_`` to a ``(ref_path, query_path)`` pair.""" + if isinstance(in_, str): + path = parse_path(in_, name="in_") + return path, path + if not (isinstance(in_, tuple | list) and len(in_) == 2): + raise ValueError(f"`in_` must be a path or a (ref_path, query_path) pair, got {in_!r}.") + return parse_path(in_[0], name="in_[0]"), parse_path(in_[1], name="in_[1]") + + +def _check_path_suits(by: Modality, path: DataPath, *, name: str) -> None: + """Reject a path whose contents the ``by`` slot could not consume. + + ``by`` says what drives the fit; the path only says where to read it from. Those are + free to differ -- ``by="landmarks"`` with ``in_="obsm/lm"`` reads correspondences out + of an ``obsm`` key, so no SpatialData is needed to hold four points. What cannot + differ is the *shape*: raster slots need rasters and coordinate slots need ``(N, 2)`` + arrays, and catching that here beats a shape error from deep inside a solver. + """ + wants_raster = by == "images" + is_raster = path.modality == "images" + if wants_raster and not is_raster: + raise ValueError(f"`by='images'` needs an image path, but `{name}={path.raw!r}` reads coordinates.") + if not wants_raster and is_raster: + raise ValueError(f"`by={by!r}` needs an (N, 2) coordinate path, but `{name}={path.raw!r}` reads an image.") + + +def _copy_for_write(container: AnnData | SpatialData, path: DataPath) -> AnnData | SpatialData: + """Duplicate just enough of ``container`` that writing at ``path`` leaves it untouched.""" + if isinstance(container, AnnData): + return container.copy() + + out = shallow_copy_sdata(container) + # `shallow_copy_sdata` shares element objects with the original, so the one element + # we are about to write through still has to be duplicated. + if path.modality == "obs" and path.element in out.tables: + out.tables[path.element] = out.tables[path.element].copy() + return out + + +def _methods_rst(registry: Registry, modality: Modality, indent: str = " " * 8) -> str: + """Render the methods supporting ``modality`` as a reST list.""" + items = [ + f"- ``{key}`` -- :func:`~{(fn := registry.get(key).implementation(modality)).__module__}.{fn.__name__}`" + for key in registry.supporting(modality) + ] + return ("\n" + indent).join(items) + + +def _document_methods(registry: Registry, **tokens: Modality) -> Callable[[F], F]: + """Fill ``{}`` docstring placeholders with each modality's method list. + + First-party and deterministic -- the registry is fully populated by import time, so + this only templates known content (nothing from optional packages). ``str.replace`` + (not ``str.format``) leaves other ``{...}`` in the docstring untouched. + """ + + def decorator(fn: F) -> F: + if fn.__doc__: + for token, modality in tokens.items(): + fn.__doc__ = fn.__doc__.replace("{" + token + "}", _methods_rst(registry, modality)) + return fn + + return decorator + + +@_document_methods(ALIGN, obs_methods="obs", image_methods="images", landmark_methods="landmarks") +def align( + data_ref: AnnData | SpatialData, + data_query: AnnData | SpatialData | None = None, + *, + in_: str | tuple[str, str], + out: str | None = None, + by: Modality = "obs", + apply_to: str | None = None, + method: str | None = None, + copy: bool = False, + **method_kwargs: Any, +) -> AlignResult | AnnData | SpatialData | None: + """Align a query sample onto a reference sample. + + Parameters + ---------- + data_ref, data_query + Both :class:`~anndata.AnnData`, or both :class:`~spatialdata.SpatialData`, or + ``data_ref`` a SpatialData with ``data_query=None`` to align two of its own + elements (distinguished by passing a pair to ``in_``). + in_ + Where to read from. One path applied to both containers, or a + ``(ref_path, query_path)`` pair. Accepted forms: + + - ``"obsm/spatial"`` -- an AnnData ``obsm`` key + - ``"tables/slice1/obsm/spatial"`` -- an ``obsm`` key of a SpatialData table + - ``"images/he"`` -- a SpatialData image + - ``"shapes/landmarks"`` -- a shapes element, as napari-spatialdata writes landmarks + + ``in_`` says only *where* to read. What the arrays mean is ``by``'s job, so + ``by="landmarks"`` can read correspondences straight out of an ``obsm`` key and + needs no SpatialData just to hold a handful of points. + out + Where to write, as a path into ``data_query``. ``None`` (default) writes nothing + and returns the fitted alignment instead -- fitting is expensive and usually worth + inspecting before it overwrites anything. + + - an ``obsm`` path writes the transformed coordinates + - an ``images`` path materialises the warped image + - ``"cs/aligned"`` registers the fit as a transformation into that coordinate + system, leaving the data untouched. Only available for methods whose fit is an + affine; a diffeomorphism has no SpatialData transformation to be expressed as. + by + What drives the alignment: + + - ``"obs"`` (default) -- the point clouds themselves + - ``"images"`` -- raster intensities + - ``"landmarks"`` -- paired correspondences, matched by row order + + This also selects which of ``method``'s slots is used, so asking for a modality a + method does not implement fails immediately and says what it does implement. + apply_to + Which array the fitted transform is applied to before writing. Defaults to ``in_`` + -- with ``by="obs"`` or ``"images"`` the thing you aligned is the thing you want + moved. ``by="landmarks"`` is the exception: ``in_`` holds correspondences rather + than data, so this must be given (or use ``out="cs/..."`` to move a whole + coordinate system at once). + method + Which method to fit with. Defaults to ``"stalign"`` for ``by="obs"`` and + ``by="images"``, and ``"similarity"`` for ``by="landmarks"``. Available per + modality: + + ``by="obs"``: + + {obs_methods} + + ``by="images"``: + + {image_methods} + + ``by="landmarks"``: + + {landmark_methods} + copy + Write into a copy of the query container and return it, instead of mutating in + place. Ignored when ``out`` is ``None``. + method_kwargs + Solver arguments, forwarded flat to the chosen ``method``. + + Returns + ------- + The fitted :class:`~squidpy.experimental.tl.AlignResult` when ``out`` is ``None``; + the modified copy when ``copy=True``; otherwise ``None``. + """ + if by not in MODALITIES: + raise ValueError(f"Unknown `by={by!r}`. Expected one of {', '.join(MODALITIES)}.") + + ref_path, query_path = _resolve_in(in_) + if ref_path.modality != query_path.modality: + raise ValueError( + f"`in_` mixes modalities: {ref_path.raw!r} is {ref_path.modality}, " + f"{query_path.raw!r} is {query_path.modality}. Both must address the same kind of data." + ) + _check_path_suits(by, ref_path, name="in_") + align_method = ALIGN.get(method if method is not None else DEFAULT_METHOD[by]) + + query_container = data_ref if data_query is None else data_query + if data_query is None and not isinstance(data_ref, SpatialData): + raise ValueError("`data_query` is required unless `data_ref` is a SpatialData holding both elements.") + + ref_array = read_path(data_ref, ref_path, name="in_") + query_array = read_path(query_container, query_path, name="in_") + + result = align_method.implementation(by)(ref=ref_array, query=query_array, **method_kwargs) + + if out is None: + return result + + out_path = parse_path(out, name="out") + if out_path.coordinate_system: + return _register_transformation(result, data_ref, query_container, ref_path, query_path, out_path, copy=copy) + + source_path = _resolve_apply_to(apply_to, query_path, by, out_path) + if out_path.modality != source_path.modality: + raise ValueError( + f"`out={out!r}` is {out_path.modality} but the data being transformed " + f"({source_path.raw!r}) is {source_path.modality}; alignment does not convert between the two." + ) + + target = _copy_for_write(query_container, out_path) if copy else query_container + if out_path.modality == "images": + value = np.asarray(result.warp_image(read_path(target, source_path, name="apply_to"))) + else: + value = np.asarray(result.transform(read_path(target, source_path, name="apply_to"))) + write_path(target, out_path, value) + return target if copy else None + + +def _resolve_apply_to( + apply_to: str | None, + query_path: DataPath, + by: Modality, + out_path: DataPath, +) -> DataPath: + """Which array the fitted transform is applied to. + + For an ``obs`` or ``images`` fit this is just ``in_`` -- the thing you aligned is the + thing you want moved. A landmark fit is different: ``in_`` holds correspondences, not + data, so the target has to be named. There is no default for that; guessing + ``obsm/spatial`` would silently transform the wrong array in any dataset that happens + to key its coordinates differently. + """ + if apply_to is not None: + return parse_path(apply_to, name="apply_to") + if by == "landmarks": + raise ValueError( + f"`out={out_path.raw!r}` needs `apply_to` when aligning by landmarks: `in_` holds the " + f"landmark correspondences, so it does not say which array to transform. Pass e.g. " + f'`apply_to="obsm/spatial"`, or use `out="cs/"` on a SpatialData to move every ' + f"element in the coordinate system at once." + ) + return query_path + + +def _register_transformation( + result: AlignResult, + data_ref: AnnData | SpatialData, + container: AnnData | SpatialData, + ref_path: DataPath, + query_path: DataPath, + out_path: DataPath, + *, + copy: bool, +) -> SpatialData | None: + """Register an affine fit into a coordinate system instead of materialising it.""" + if not isinstance(container, SpatialData): + raise TypeError(f"`out={out_path.raw!r}` names a coordinate system, which only a SpatialData has.") + if not hasattr(result, "matrix"): + raise ValueError( + f"`out={out_path.raw!r}` registers the fit as a transformation, but this method fits a " + f"deformation that SpatialData has no transformation type for -- its transformations are " + f"affine at most. Write to an `images/` or `obsm/` path to materialise it instead." + ) + + moving_cs = _coordinate_system_of(container, query_path, name="in_") + # Registering moves *everything* in `moving_cs`. If the reference sits in that same + # coordinate system of the same object, it would be dragged along with the query -- + # silently producing a wrong answer rather than failing. + if data_ref is container and _coordinate_system_of(data_ref, ref_path, name="in_") == moving_cs: + raise ValueError( + f"The reference and query are both in coordinate system {moving_cs!r}, so registering " + f"the fit there would move the reference too. Put each sample in its own coordinate " + f"system (what napari-spatialdata does when landmarks are picked per sample), or write " + f"to a data path with `apply_to` to move only the query." + ) + + return writeback_affine_sdata( + result, + container, + output_mode="copy" if copy else "inplace", + moving_cs=moving_cs, + target_cs=out_path.element, + ) + + +#: Element collections whose members carry transformations. Tables do not: they annotate +#: elements rather than sitting in space themselves. +_SPATIAL_COLLECTIONS = ("shapes", "points", "images", "labels") + + +def _coordinate_system_of(sdata: SpatialData, path: DataPath, *, name: str) -> str: + """The coordinate system the element at ``path`` is annotated in. + + Everything registered to it moves with the fit, so it has to be unambiguous. Reading + it off the element rather than taking it as an argument keeps the call site to + ``in_``/``out``, and it is the same element the user picked the landmarks on. + """ + from spatialdata.transformations import get_transformation + + collection = path.raw.strip("/").split("/")[0] + if collection not in _SPATIAL_COLLECTIONS: + raise ValueError( + f'`out="cs/..."` needs `{name}={path.raw!r}` to name a spatial element, but a table has no ' + f"coordinate system of its own. Store the landmarks as a shapes element, or write to a " + f"data path with `apply_to` to move only the query's coordinates." + ) + + systems = sorted(get_transformation(getattr(sdata, collection)[path.element], get_all=True)) + if len(systems) != 1: + raise ValueError( + f"`{name}={path.raw!r}` is registered to {len(systems)} coordinate systems ({', '.join(systems)}), " + f"so which one the alignment should move is ambiguous. Register it to exactly one." + ) + return systems[0] diff --git a/src/squidpy/experimental/tl/_align/_io.py b/src/squidpy/experimental/tl/_align/_io.py new file mode 100644 index 000000000..8a7c95e3a --- /dev/null +++ b/src/squidpy/experimental/tl/_align/_io.py @@ -0,0 +1,71 @@ +"""SpatialData write-back for the public align functions. + +Reading and writing arrays at a path lives in :mod:`._paths`. What is left here is the +one write that is *not* an array write: registering a fitted affine as a SpatialData +transformation, which changes how elements are placed rather than materialising anything. + +The fit estimators in :mod:`squidpy.experimental.methods` operate on plain arrays and +never see a container. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from anndata import AnnData +from spatialdata import SpatialData + +if TYPE_CHECKING: + from squidpy.experimental.methods.align_landmarks import AffineFitResult + +__all__ = ["shallow_copy_sdata", "writeback_affine_sdata"] + + +def writeback_affine_sdata( + result: AffineFitResult, + sdata: SpatialData, + *, + output_mode: str, + moving_cs: str | None, + target_cs: str | None, +) -> SpatialData | None: + """Register the fitted affine on every element living in ``moving_cs``. + + Non-destructive: it adds a transformation into ``target_cs`` so the whole + coordinate system inherits the alignment. Nothing is materialised. + """ + from spatialdata import deepcopy as sd_deepcopy + from spatialdata.transformations import Affine, Sequence, get_transformation, set_transformation + + if moving_cs is None or target_cs is None: + raise ValueError("`cs_query` and `cs_ref` are required to register a transform on a SpatialData.") + + out = sdata if output_mode == "inplace" else shallow_copy_sdata(sdata) + sd_affine = Affine(np.asarray(result.matrix), input_axes=("x", "y"), output_axes=("x", "y")) + touched = False + for etype, name, element in list(out.gen_elements()): + if isinstance(element, AnnData): + continue + if moving_cs not in get_transformation(element, get_all=True): + continue + if output_mode == "copy": + # `shallow_copy_sdata` shares element objects with the original; deep-copy each + # element we register a transform on so `copy=True` leaves the input untouched. + element = sd_deepcopy(element) + getattr(out, etype)[name] = element + # The fitted affine maps `moving_cs` coords into `target_cs`, not the element's + # intrinsic frame. Compose it after the element's existing intrinsic -> `moving_cs` + # transform so a non-identity placement into `moving_cs` is preserved. + existing = get_transformation(element, to_coordinate_system=moving_cs) + set_transformation(element, Sequence([existing, sd_affine]), to_coordinate_system=target_cs) + touched = True + if not touched: + raise KeyError(f"No elements in the SpatialData are registered to coordinate system {moving_cs!r}.") + return None if output_mode == "inplace" else out + + +def shallow_copy_sdata(sdata: SpatialData) -> SpatialData: + """Shallow copy of a SpatialData for ``copy=True`` (via ``subset``).""" + names = [name for _, name, _ in sdata.gen_elements()] + return sdata.subset(names, filter_tables=False, include_orphan_tables=True) diff --git a/src/squidpy/experimental/tl/_align/_paths.py b/src/squidpy/experimental/tl/_align/_paths.py new file mode 100644 index 000000000..d91a7f83a --- /dev/null +++ b/src/squidpy/experimental/tl/_align/_paths.py @@ -0,0 +1,218 @@ +"""Location paths for the ``in_`` / ``out`` arguments of the align functions. + +A path names one location inside an :class:`~anndata.AnnData` or +:class:`~spatialdata.SpatialData`: + +- ``obsm/spatial`` -- an AnnData ``obsm`` key +- ``tables/slice1/obsm/spatial`` -- an ``obsm`` key of a SpatialData table +- ``images/he`` -- a SpatialData image +- ``shapes/landmarks`` -- a SpatialData shapes element holding landmark correspondences +- ``cs/aligned`` -- a coordinate system to register a transformation into (``out`` only) + +Naming a location this way rather than through a stack of ``*_key`` arguments means one +argument covers *which element* and *which array inside it*, and the modality follows +from the path instead of a separate ``on=`` switch. This is the shape proposed for +scanpy in https://github.com/scverse/scanpy/issues/4007. + +The path also decides what drives the alignment: ``obsm``/``tables`` paths align point +clouds, ``images`` paths align raster intensities, and ``shapes`` paths align by paired +landmarks. ``cs`` is the one write-only form -- it registers a transformation rather than +materialising an array, which only an affine-representable fit can do. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +from anndata import AnnData +from spatialdata import SpatialData + +__all__ = ["DataPath", "parse_path", "read_path", "write_path"] + +from squidpy.experimental.methods.registry import Modality + +#: Element collections a path may address on a SpatialData. +_SDATA_ATTRS = ("tables", "images", "labels", "points", "shapes") +#: Accepted spellings for the coordinate-system namespace, which is not an element +#: collection: writing there registers a transformation instead of storing an array. +_COORDINATE_SYSTEM_ATTRS = ("cs", "coordinate_systems") +_EXAMPLES = '"obsm/spatial", "tables//obsm/spatial", "images/", "shapes/", or "cs/"' + + +@dataclass(frozen=True, slots=True) +class DataPath: + """A parsed location inside a container.""" + + modality: Modality + #: SpatialData element name, or ``None`` for a bare AnnData path. + element: str | None + #: ``obsm`` key for ``obs`` paths; ``None`` for every other modality. + key: str | None + #: The original string, for error messages. + raw: str + #: True for ``cs/``: a coordinate system to register a transformation into, + #: rather than a location holding an array. + coordinate_system: bool = False + + def __str__(self) -> str: + return self.raw + + +def parse_path(path: str, *, name: str) -> DataPath: + """Parse a location string, or raise :class:`ValueError` explaining the grammar.""" + if not isinstance(path, str): + raise TypeError(f"`{name}` must be a string path, got {type(path).__name__}. Expected one of {_EXAMPLES}.") + + parts = tuple(part for part in path.strip("/").split("/") if part) + if not parts: + raise ValueError(f"`{name}` is empty. Expected one of {_EXAMPLES}.") + + head = parts[0] + if head == "obsm": + if len(parts) != 2: + raise ValueError(f"`{name}={path!r}` is not a valid obsm path. Expected `obsm/`.") + return DataPath(modality="obs", element=None, key=parts[1], raw=path) + + if head == "tables": + if len(parts) != 4 or parts[2] != "obsm": + raise ValueError(f"`{name}={path!r}` is not a valid table path. Expected `tables/
/obsm/`.") + return DataPath(modality="obs", element=parts[1], key=parts[3], raw=path) + + if head == "images": + if len(parts) != 2: + raise ValueError(f"`{name}={path!r}` is not a valid image path. Expected `images/`.") + return DataPath(modality="images", element=parts[1], key=None, raw=path) + + if head == "shapes": + if len(parts) != 2: + raise ValueError(f"`{name}={path!r}` is not a valid landmark path. Expected `shapes/`.") + return DataPath(modality="landmarks", element=parts[1], key=None, raw=path) + + if head in _COORDINATE_SYSTEM_ATTRS: + if len(parts) != 2: + raise ValueError(f"`{name}={path!r}` is not a valid coordinate system. Expected `cs/`.") + return DataPath(modality="landmarks", element=parts[1], key=None, raw=path, coordinate_system=True) + + if head in _SDATA_ATTRS: + raise ValueError( + f"`{name}={path!r}` addresses {head!r}, which alignment does not read or write yet. " + f"Expected one of {_EXAMPLES}." + ) + raise ValueError(f"`{name}={path!r}` does not start with a known collection. Expected one of {_EXAMPLES}.") + + +def _table(container: AnnData | SpatialData, path: DataPath, *, name: str) -> AnnData: + """Resolve the AnnData a point path refers to.""" + if isinstance(container, AnnData): + if path.element is not None: + raise ValueError( + f"`{name}={path.raw!r}` names a SpatialData table, but the container is an AnnData. " + f"Use `obsm/{path.key}`." + ) + return container + + if not isinstance(container, SpatialData): + raise TypeError(f"Expected an AnnData or SpatialData, got {type(container).__name__}.") + if path.element is None: + raise ValueError( + f"`{name}={path.raw!r}` is ambiguous for a SpatialData, which may hold several tables. " + f"Use `tables/
/obsm/{path.key}`." + ) + if path.element not in container.tables: + raise KeyError(f"`{name}={path.raw!r}`: no table {path.element!r}. Available: {sorted(container.tables)}.") + return container.tables[path.element] + + +def read_path(container: AnnData | SpatialData, path: DataPath, *, name: str) -> np.ndarray: + """Read the array a path points at. + + Returns an ``(N, 2)`` ``(x, y)`` array for ``obs`` and ``landmarks`` paths, or a + channels-first ``(C, H, W)`` array for ``images`` paths. + """ + if path.coordinate_system: + raise ValueError( + f"`{name}={path.raw!r}` names a coordinate system, which holds transformations " + f"rather than data to align. It is only valid as an `out`." + ) + + if path.modality == "obs": + adata = _table(container, path, name=name) + if path.key not in adata.obsm: + raise KeyError(f"`{name}={path.raw!r}`: no `obsm[{path.key!r}]`. Available: {sorted(adata.obsm)}.") + coords = np.asarray(adata.obsm[path.key]) + if coords.ndim != 2 or coords.shape[1] != 2: + raise ValueError(f"`{name}={path.raw!r}` must be an (N, 2) array, found shape {coords.shape}.") + return coords + + if path.modality == "landmarks": + return _read_landmarks(container, path, name=name) + + return _read_image(container, path, name=name) + + +def _read_landmarks(container: AnnData | SpatialData, path: DataPath, *, name: str) -> np.ndarray: + """Read ``(N, 2)`` ``(x, y)`` landmark coordinates from a shapes element. + + This is the layout napari-spatialdata writes when landmarks are picked interactively, + so annotations made in the viewer are usable here without conversion. + """ + if not isinstance(container, SpatialData): + raise TypeError(f"`{name}={path.raw!r}` names a shapes element, which only a SpatialData holds.") + if path.element not in container.shapes: + raise KeyError(f"`{name}={path.raw!r}`: no shapes {path.element!r}. Available: {sorted(container.shapes)}.") + + shapes = container.shapes[path.element] + geometry = shapes.geometry + coords = np.column_stack([geometry.x.to_numpy(), geometry.y.to_numpy()]) + if coords.ndim != 2 or coords.shape[1] != 2: + raise ValueError( + f"`{name}={path.raw!r}` must hold point or circle geometries, found {geometry.geom_type[0]!r}." + ) + return coords + + +def _read_image(container: AnnData | SpatialData, path: DataPath, *, name: str) -> np.ndarray: + if not isinstance(container, SpatialData): + raise TypeError(f"`{name}={path.raw!r}` names an image, which only a SpatialData holds.") + if path.element not in container.images: + raise KeyError(f"`{name}={path.raw!r}`: no image {path.element!r}. Available: {sorted(container.images)}.") + + element = container.images[path.element] + # Multiscale images are a DataTree; the full-resolution level is the first scale. + if not hasattr(element, "dims"): + element = next(iter(element.values())) + element = element[next(iter(element.data_vars))] + + array = np.asarray(element.data) + if array.ndim == 2: + array = array[None] + if array.ndim != 3: + raise ValueError(f"`{name}={path.raw!r}` must be a 2D or (c, y, x) image, found shape {array.shape}.") + return array + + +def write_path( + container: AnnData | SpatialData, + path: DataPath, + value: np.ndarray, + *, + name: str = "out", +) -> None: + """Write ``value`` at ``path``, mutating ``container`` in place. + + Copy semantics belong to the caller: it decides what to duplicate before calling. + """ + if path.modality == "obs": + _table(container, path, name=name).obsm[path.key] = np.asarray(value) + return + + if not isinstance(container, SpatialData): + raise TypeError(f"`{name}={path.raw!r}` names an image, which only a SpatialData holds.") + + from spatialdata.models import Image2DModel + + array = np.asarray(value) + if array.ndim == 2: + array = array[None] + container.images[path.element] = Image2DModel.parse(array, dims=("c", "y", "x")) diff --git a/tests/_data/stalign_reference/converged_n500.npz b/tests/_data/stalign_reference/converged_n500.npz new file mode 100644 index 000000000..47cd3eacf Binary files /dev/null and b/tests/_data/stalign_reference/converged_n500.npz differ diff --git a/tests/_data/stalign_reference/energy.npz b/tests/_data/stalign_reference/energy.npz new file mode 100644 index 000000000..23d7fd58f Binary files /dev/null and b/tests/_data/stalign_reference/energy.npz differ diff --git a/tests/_data/stalign_reference/gradients.npz b/tests/_data/stalign_reference/gradients.npz new file mode 100644 index 000000000..4c4e4a98d Binary files /dev/null and b/tests/_data/stalign_reference/gradients.npz differ diff --git a/tests/_data/stalign_reference/image_trajectory.npz b/tests/_data/stalign_reference/image_trajectory.npz new file mode 100644 index 000000000..25ad7e717 Binary files /dev/null and b/tests/_data/stalign_reference/image_trajectory.npz differ diff --git a/tests/_data/stalign_reference/image_trajectory_matched.npz b/tests/_data/stalign_reference/image_trajectory_matched.npz new file mode 100644 index 000000000..d5ba4377b Binary files /dev/null and b/tests/_data/stalign_reference/image_trajectory_matched.npz differ diff --git a/tests/_data/stalign_reference/primitives.npz b/tests/_data/stalign_reference/primitives.npz new file mode 100644 index 000000000..eaacc8936 Binary files /dev/null and b/tests/_data/stalign_reference/primitives.npz differ diff --git a/tests/_data/stalign_reference/trajectory_n1.npz b/tests/_data/stalign_reference/trajectory_n1.npz new file mode 100644 index 000000000..eb405f3b3 Binary files /dev/null and b/tests/_data/stalign_reference/trajectory_n1.npz differ diff --git a/tests/_data/stalign_reference/trajectory_n5.npz b/tests/_data/stalign_reference/trajectory_n5.npz new file mode 100644 index 000000000..b4e1212e0 Binary files /dev/null and b/tests/_data/stalign_reference/trajectory_n5.npz differ diff --git a/tests/_data/stalign_reference/trajectory_n50.npz b/tests/_data/stalign_reference/trajectory_n50.npz new file mode 100644 index 000000000..899dbc8d0 Binary files /dev/null and b/tests/_data/stalign_reference/trajectory_n50.npz differ diff --git a/tests/experimental/methods/STALIGN_DIVERGENCES.md b/tests/experimental/methods/STALIGN_DIVERGENCES.md new file mode 100644 index 000000000..5eb2ab333 --- /dev/null +++ b/tests/experimental/methods/STALIGN_DIVERGENCES.md @@ -0,0 +1,191 @@ +# STalign port — divergence ledger + +Where squidpy's JAX port +(`src/squidpy/experimental/methods/align_samples/_stalign_impl/`) differs from the +PyTorch original it was ported from, why, and what was done about it. + +Every number below is **measured**, not estimated, by +`tests/experimental/methods/test_stalign_reference.py` against the reference bundle in +`tests/_data/stalign_reference/` (generated by +[theislab/squidpy-ports](https://github.com/theislab/squidpy-ports) from upstream commit +`b2068edc98974efa54537eca194736e177bbe11d`). Upstream line references are into that +commit's `STalign/STalign.py`. + +Each row has an id. Every `xfail(strict=True)` in the test module cites its row id, and +`test_divergences_doc_covers_all_xfails` asserts the citation resolves — so this file +cannot silently rot. + +Comparisons go through the public API wherever one exists: point transforms via +`StalignResult.transform`, the image warp via `StalignResult.warp_image`, the landmark +initialisation via `fit_stalign_obs(..., niter=0)`. What stays white-box is what has no +public route — the objective and its gradients are not observable from outside, and the +regulariser and grids are preconditions rather than results. + +## Where it stands + +The port now reproduces the original **to machine precision** on everything that defines +the optimisation: + +| | relative error | +| --- | --- | +| objective `E` | **0.0** (bit-for-bit, `131.057468`) | +| `dE/dL`, `dE/dT`, `dE/dv` | **≤ 6e-15** | +| `_interp` | 4e-16 | +| `_transform_grid_backward` | 5e-18 | +| forward point transform | 5e-17 | +| regulariser `LL` / `K` / `DV` | ≤ 1e-15 | +| raster and velocity grids | identical | +| full loop, 1 / 5 / 50 steps | ≤ 1e-10 / 1e-9 / 1e-6 on `A`, `v`, `WM`, `WA`, `WB` | +| converged (500 steps) | final `E` within 1 %; 95th-percentile point disagreement < 0.1·`dx` | + +The 50-step case matters on its own: the mixture-weight E step is gated on `it >= 50` +(`STalign.py:1233`), so a shorter run leaves that branch entirely untested. + +Three findings were real bugs and are **fixed** (D2, D3, R1). One is a deliberate +approximation, now measured and budgeted (D1). The rest are places where squidpy is +deliberately *not* bug-compatible with upstream (D4, D5, D6, D7, D9) — those stay, pinned. + +## Ledger + +| id | What | Upstream | squidpy | Measured | Status | +| --- | --- | --- | --- | --- | --- | +| **D1** | Rasterisation algorithm | exact sub-pixel Gaussian splat per point, truncated to ±`ceil(4·max(blur))` px and renormalised over that window — `:178-201` | bilinear deposit onto the grid, then one `scipy.ndimage.gaussian_filter` per scale, mass-corrected — [`_helpers.py`](../../../src/squidpy/experimental/methods/align_samples/_stalign_impl/_helpers.py) | relL2 **4.08 % / 0.81 % / 2.87 %** at blur 2.0/1.0/0.5; mass exact | **Accepted, budgeted.** Improved from 6.20/1.98/5.96 % — see R2 | +| **D2** | Grid construction off-by-one | `np.arange(lo, hi, dx)` — `:137-138`, `:1069` | was `np.arange(lo, hi + dx, dx)` | was +1 sample per axis, and *length varied with rounding* — `n+1` in 7 of 9 sweep cases, `n+2` in 2 | **Fixed.** Count derived from the interval first; grids now identical to upstream | +| **D3** | Gradient of the contrast transform | ridge coefficients solved under `torch.no_grad()` — `:1184-1188` | was differentiating through `jnp.linalg.solve` | was `dE/dL` 1.25e-3, `dE/dT` 9.98e-4, `dE/dv` 1.22e-3; now **≤ 6e-15** | **Fixed** with `jax.lax.stop_gradient`. The ridge fit is an EM **M step**; differentiating through it turned alternating minimisation into joint optimisation | +| **D4** | Returned affine lags a step | `A = to_A(L,T)` built at the *top* of the loop and returned — `:1155`, `:1308` | built after the loop | `LDDMM(n)["A"]` reflects `n-1` updates | **squidpy is right** — upstream discards its last step. Comparisons use `squidpy(n).A ↔ upstream(n+1).A` | +| **D5** | Padding when sampling outside the domain | `grid_sample` default `padding_mode='zeros'` for velocity *and* point warps (`:1163`, `:1167`); only the image warp uses `'border'` (`:1171`) | `map_coordinates(mode="nearest")` (≈`border`) everywhere | outside the domain: **25 %** from upstream-`zeros`, **2.8e-16** from upstream-`border` | **squidpy is right.** Upstream's zeros make a point that drifts off the velocity grid snap to *no* displacement — a discontinuity. Pinned `xfail` | +| **D6** | Backward point transform, time order | integrates `-v[t]` for `t` in `range(nt)` — forward order — `:1828-1843` | `reversed(range(nt))` | outputs differ by **1.03e-6**; forward∘backward round-trip error **5.79e-7 (squidpy)** vs **9.01e-7 (upstream)** | **squidpy is right** — it is the correct explicit-Euler inverse, and upstream contradicts its own image warp (`:1163`, which *does* reverse). Pinned `xfail` | +| **D7** | Landmark affine solve | normal equations for the plain least-squares fit, with explicit `np.linalg.inv` on the Gram matrix — `:897-910` | `skimage.transform.estimate_transform("affine")`, a Hartley-normalised homogeneous solve by SVD | `L` differs **6.4e-4**, `T` **7.2e-3**. Fit residual on clean landmarks: **21.7026** vs **21.6984** (1.9e-4 apart). On near-collinear landmarks: **7.4e-13** vs **5.6e+2** | **Different estimators, not the same one twice.** skimage minimises algebraic error, upstream geometric; upstream is a hair better on clean input and collapses when ill-conditioned. Keep squidpy's — see R7 | +| **D8** | `lddmm()` cannot take a precomputed `xv`/`v` | `LDDMM` accepts both — `:1060-1064` | no such parameters | — | **No change needed.** The generator forces *upstream* onto squidpy's grid instead | +| **D9** | Division guards | none | `jnp.maximum(…, 1e-12)` in `_update_mixture_weights` | inert for short runs | **squidpy is right.** Keep | +| **D10** | On-grid interpolation kink | normalises `(c-x0)/(x[-1]-x0)`, then `grid_sample(align_corners=True)` scales by `(n-1)` | `(c-x0)/(x[1]-x0)` | equal to ~1 ulp, but a sample landing exactly on a grid line can `floor()` to different neighbours. **Measured cost: 1e-12 → 1e-3** on the velocity field (see below) | Not a defect in either. Every fixture is built off-grid on purpose and asserts it | + +## Review — beyond the divergences + +Port-quality findings, not upstream comparisons. + +### R1. `lddmm(niter=0)` raised `UnboundLocalError` — fixed + +`energy` and `transformed_points` were bound only inside the loop, so the `return` +read unbound locals. `niter=0` is a reasonable request (evaluate the initial affine and +stop); both are now initialised before the loop, and +`test_lddmm_accepts_zero_iterations` covers it. + +### R2. Rasteriser accuracy — fixed, ~2× better + +The gap was two separable problems, dominating at opposite ends of the blur range. +Measured, all four combinations: + +| variant | blur 2.0 | blur 1.0 | blur 0.5 | mass (of 800) | +| --- | --- | --- | --- | --- | +| nearest-cell + leaky border (was) | 6.20 % | 1.98 % | 5.96 % | 770.3 / 795.1 / 799.5 | +| bilinear deposit only | 6.34 % | 1.25 % | 2.88 % | leaks | +| mass conservation only | 3.92 % | 1.72 % | 5.97 % | exact | +| **both (now)** | **4.08 %** | **0.81 %** | **2.87 %** | **exact** | + +1. **Sub-pixel quantisation.** `np.rint` snapped every point to a cell centre before any + blurring — up to half a cell of positional error, comparable to the features being + registered. Now deposited bilinearly across the four neighbouring cells, still fully + vectorised (`np.bincount`, no Python loop over points). +2. **Border mass loss.** `mode="constant"` let kernels near the edge spill off-grid, so + the density was biased low around the whole rim — 3 % of total mass at `blur=2.0`. The + mass a point at cell `c` retains is `sum_p K(p-c)`, which by symmetry of `K` equals + `gaussian_filter(ones)[c]`; dividing by that before blurring conserves mass exactly, + for one extra filter pass. + +Note the two are not additive: bilinear deposit alone made `blur=2.0` slightly *worse*, +because the error there was dominated by mass loss. + +### R3. `blur` docstring — fixed + +It said blur was "the kernel width in units of `2 * dx`", which reads as if `blur` were +scaled by `2·dx`. The code was always correct and matches upstream exactly; the docstring +now says σ = `2·blur` pixels = `2·blur·dx` physical. + +### R4. The default dtype is float32; upstream is float64 throughout — documented + +`jax_dtype()` returns `float32` unless `jax_enable_x64` is set, and nothing in squidpy +sets it — so the shipped default runs a 5000-iteration gradient descent with +`sigmaR=5e5` in single precision, while every published STalign result is double. +`fit_stalign` now says so and shows how to enable x64. This suite requires +`JAX_ENABLE_X64=1` and skips without it. + +### D10 in practice: how the image path was nearly mis-assessed + +`fit_stalign_image` had no reference comparison at all — the last gap in the port's public +surface. Adding one produced a **5.3e-2** disagreement on the affine and **2.6e-1** on the +velocity field, which looked like a real port bug. + +It was not. The bisect is worth recording, because the same trap is waiting for anyone who +adds a fixture without the off-grid guard: + +| step | finding | +| --- | --- | +| energy at iteration 0 | matched to **4.6e-16** — the objective is identical | +| energy at iteration 1 | already 3e-3 apart — the *first gradient step* diverges | +| conditioning check | a 1e-12 nudge in `epV` moved the answer by 1e-13, linearly — not chaos | +| grids, axes, images, every parameter | verified identical on both sides | +| first hypothesis: padding | 11 % of the target grid samples the source out of domain, because each raster is centred on its own centre. **Wrong** — padding agrees on values *and* gradients | +| actual cause | centred pixel axes are **integers**, so an identity starting affine put every interpolation sample exactly on a grid line: D10, at full strength | + +Starting the fixture from a deliberately off-grid affine (a 0.0371449 rad rotation and a +non-integer shift, mirroring `fixtures.THETA`/`SHIFT`) drops the disagreement to **3.95e-12** +on the affine and **2.43e-12** on the velocity — with the 11 % padded samples still present. + +`image_trajectory_matched` keeps the degenerate case on purpose: both rasters cropped to a +common extent, so their axes are the *same* integers and grid coincidence persists even +off-grid. It still lands at ~1e-3, and `test_on_grid_sampling_costs_six_orders_of_magnitude` +pins that gap so the cost of D10 stays a measured number rather than a warning. + +### R7. `affine_from_points` silently changes the landmark estimator — open + +Following on from D7: `estimate_transform("affine")` is not a drop-in for upstream's +normal-equations fit. It minimises **algebraic** error on Hartley-normalised coordinates; +upstream minimises **geometric** residual. Every landmark-initialised alignment therefore +starts from a measurably different affine (`L` by 6.4e-4, `T` by 7.2e-3), which for a +5000-iteration descent is a different starting point, not a rounding detail. + +skimage's choice is the safer one, so this is not a request to revert — it is a request to +*say so* in the docstring. If exact least-squares is wanted, `np.linalg.lstsq` on the +padded design matrix gives upstream's answer without upstream's conditioning problem. + +### R5. `niter=5000` ran blind — fixed + +`lddmm` now returns the per-iteration `energies` trace and the `n_iter` actually run, so +a converged run is distinguishable from a diverged one without running it again. +(Upstream accumulates the same values in a local `Esave` and throws them away.) + +It also takes an optional `tol` / `patience` early stop: it halts once the objective's +relative improvement over the last `patience` iterations falls below `tol`. Off by +default, so the shipped behaviour is unchanged. + +**The window must clear iteration 50.** The mixture-weight E step switches on at +`MIXTURE_E_STEP_START` (`STalign.py:1233`), which changes what the objective *is* — and +its value jumps **upward** there, on the reference fixture from `3.259` to `5.489`. Any +rule comparing across that boundary reads the jump as "no longer improving" and quits +immediately: with a naive guard every tolerance from `1e-3` to `1e-5` stopped at +iteration 76 regardless. The comparison window now has to sit entirely after the jump. +`test_early_stopping_never_fires_before_the_weights_switch_on` pins this. + +### R6. The Python loop forgave most of JAX's advantage — fixed + +The descent was a Python `for` loop around a jitted `value_and_grad`, paying dispatch +per iteration and blocking XLA from fusing across steps. It is now a single +`lax.while_loop` inside one `jax.jit`, with `lax.cond` for the every-5th-iteration +weight update and `jnp.where` for the `diffeo_start` gates. + +Measured on the reference fixture, warm: + +| | ms/iteration | `niter=5000` | +| --- | --- | --- | +| Python loop (was) | 2.20 | ~11.0 s | +| `lax.while_loop`, not jitted | 0.91 | ~4.6 s | +| **`lax.while_loop` inside `jit`** | **0.46** | **~2.4 s** | + +The middle row is worth keeping in mind: `lax.while_loop` outside a `jit` re-traces its +body on *every call*, and tracing `value_and_grad` through the interpolation and FFTs +costs roughly a second — about as much as a thousand iterations of running it. The loop +has to be inside the `jit` to get the win. + +Reference parity is unchanged at 1 / 5 / 50 / 500 iterations, which is what makes this +rewrite safe to make at all. diff --git a/tests/experimental/methods/__init__.py b/tests/experimental/methods/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/experimental/methods/_stalign_fixtures.py b/tests/experimental/methods/_stalign_fixtures.py new file mode 100644 index 000000000..3861cacc9 --- /dev/null +++ b/tests/experimental/methods/_stalign_fixtures.py @@ -0,0 +1,167 @@ +"""Fixture definitions shared by squidpy-ports and squidpy. + +This module is vendored **byte-identically** into squidpy as +``tests/experimental/methods/_stalign_fixtures.py``. Both sides building their inputs +from the same source is what stops the committed reference bundle and the tests that +consume it from drifting apart; a checksum of this file is stamped into every generated +``.npz`` and asserted on the squidpy side. + +Consequently it must depend on **numpy only** -- no torch, no jax, no squidpy_ports. +""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import numpy as np + +__all__ = [ + "LDDMM_PARAMS", + "N_LANDMARKS", + "N_POINTS", + "RASTER_PARAMS", + "SEED", + "SHIFT", + "THETA", + "Clouds", + "checksum", + "make_clouds", + "rotation", +] + +SEED = 20240501 +N_POINTS = 800 +N_LANDMARKS = 6 + +# Deliberately not a round angle and not an integer shift. +# +# Upstream `interp` normalises coordinates with (c - x0) / (x[-1] - x0) and then lets +# grid_sample(align_corners=True) denormalise by *(n - 1); squidpy's `_interp` uses +# (c - x0) / (x[1] - x0). Those are algebraically equal but differ by ~1 ulp, so a +# sample landing exactly on a grid line can floor() to different neighbour pairs on the +# two sides -- an O(1) disagreement that says nothing about the port. An "easy" rotation +# angle or an integer shift makes that happen everywhere at once. `test_fixture_samples_ +# are_off_grid` asserts we stayed clear of it. +THETA = 0.1503216 +SHIFT = (25.371, -15.119) + +#: Fraction of the cloud's RMS radius added as isotropic noise to the query. +JITTER_FRAC = 0.03 +#: Fraction of query points dropped, so the correspondence is not a bijection. +DROPOUT_FRAC = 0.05 + +#: Rasterisation controls. Small enough that a full LDDMM solve runs in seconds. +RASTER_PARAMS = {"dx": 30.0, "blur": [2.0, 1.0, 0.5], "expand": 1.1} + +#: LDDMM controls. `a=200` keeps the velocity grid ~14x23; everything else is upstream's +#: default so we are comparing the published configuration. +LDDMM_PARAMS = { + "a": 200.0, + "p": 2.0, + "expand": 2.0, + "nt": 3, + "diffeo_start": 0, + "epL": 2e-8, + "epT": 2e-1, + "epV": 2e3, + "sigmaM": 1.0, + "sigmaB": 2.0, + "sigmaA": 5.0, + "sigmaR": 5e5, + "sigmaP": 2e1, +} + + +class Clouds: + """A reference/query point-cloud pair with corresponding landmarks, all in ``(x, y)``.""" + + def __init__( + self, + ref: np.ndarray, + query: np.ndarray, + landmarks_ref: np.ndarray, + landmarks_query: np.ndarray, + ) -> None: + self.ref = ref + self.query = query + self.landmarks_ref = landmarks_ref + self.landmarks_query = landmarks_query + + # `torch.tensor` rejects arrays with negative strides, which a bare `[:, ::-1]` + # produces, so every row-col view is materialised contiguous. + + @property + def ref_rc(self) -> np.ndarray: + """Reference cloud in row-col order, which is what upstream works in.""" + return np.ascontiguousarray(self.ref[:, ::-1]) + + @property + def query_rc(self) -> np.ndarray: + """Query cloud in row-col order.""" + return np.ascontiguousarray(self.query[:, ::-1]) + + @property + def landmarks_ref_rc(self) -> np.ndarray: + """Reference landmarks in row-col order.""" + return np.ascontiguousarray(self.landmarks_ref[:, ::-1]) + + @property + def landmarks_query_rc(self) -> np.ndarray: + """Query landmarks in row-col order.""" + return np.ascontiguousarray(self.landmarks_query[:, ::-1]) + + +def rotation(theta: float) -> np.ndarray: + """A 2x2 rotation matrix acting on ``(x, y)`` column vectors.""" + c, s = np.cos(theta), np.sin(theta) + return np.array([[c, -s], [s, c]], dtype=float) + + +def make_clouds() -> Clouds: + """Build the deterministic reference/query pair. + + The reference is three structures -- an anisotropic Gaussian lobe, a smaller offset + lobe, and an annulus -- so the rasterised density has real structure rather than one + featureless blob, and the registration has something to lock onto. The query is that + cloud rotated, shifted, jittered and randomly thinned. + """ + rng = np.random.default_rng(SEED) + + n_lobe, n_small = 400, 200 + n_ring = N_POINTS - n_lobe - n_small + + lobe = rng.normal(size=(n_lobe, 2)) @ np.diag([150.0, 100.0]) @ rotation(0.4).T + small = rng.normal(size=(n_small, 2)) @ np.diag([80.0, 80.0]) + np.array([400.0, 250.0]) + + angle = rng.uniform(0.0, 2.0 * np.pi, size=n_ring) + radius = 300.0 + rng.normal(scale=25.0, size=n_ring) + ring = np.stack([radius * np.cos(angle), radius * np.sin(angle)], axis=1) + np.array([100.0, -200.0]) + + ref = np.concatenate([lobe, small, ring], axis=0) + + rms_radius = float(np.sqrt(np.mean(np.sum((ref - ref.mean(0)) ** 2, axis=1)))) + moved = ref @ rotation(THETA).T + np.asarray(SHIFT) + jittered = moved + rng.normal(scale=JITTER_FRAC * rms_radius, size=moved.shape) + + keep = rng.permutation(N_POINTS)[: int(round(N_POINTS * (1.0 - DROPOUT_FRAC)))] + keep.sort() + query = jittered[keep] + + # Landmarks are spread across the whole cloud rather than clustered, so the affine + # they induce is well conditioned. + lm_idx = keep[np.linspace(0, keep.size - 1, N_LANDMARKS).round().astype(int)] + landmarks_ref = ref[lm_idx] + landmarks_query = jittered[lm_idx] + + return Clouds( + ref=np.ascontiguousarray(ref, dtype=float), + query=np.ascontiguousarray(query, dtype=float), + landmarks_ref=np.ascontiguousarray(landmarks_ref, dtype=float), + landmarks_query=np.ascontiguousarray(landmarks_query, dtype=float), + ) + + +def checksum() -> str: + """SHA-256 of this file, so both copies can prove they are the same file.""" + return hashlib.sha256(Path(__file__).read_bytes()).hexdigest() diff --git a/tests/experimental/methods/test_core.py b/tests/experimental/methods/test_core.py new file mode 100644 index 000000000..ab782fab6 --- /dev/null +++ b/tests/experimental/methods/test_core.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import pytest + +from squidpy.experimental.methods.registry import AlignMethod, Registry + + +@dataclass +class _MeanShiftResult: + """Toy result: a constant per-axis offset baked into ``transform``.""" + + delta: np.ndarray + metadata: dict[str, Any] = field(default_factory=dict) + + def transform(self, x: np.ndarray) -> np.ndarray: + return np.asarray(x, dtype=float) + self.delta + + +def fit_mean_shift(ref: np.ndarray, query: np.ndarray) -> _MeanShiftResult: + """Toy estimator function: fit the offset that maps the query centroid onto the ref centroid.""" + delta = np.asarray(ref, dtype=float).mean(0) - np.asarray(query, dtype=float).mean(0) + return _MeanShiftResult(delta=delta, metadata={"method": "mean_shift"}) + + +def test_fit_then_transform_round_trip() -> None: + ref = np.array([[1.0, 1.0], [3.0, 3.0]]) # centroid (2, 2) + query = np.array([[0.0, 0.0], [2.0, 2.0]]) # centroid (1, 1) + + result = fit_mean_shift(ref, query) + + np.testing.assert_allclose(result.delta, [1.0, 1.0]) + np.testing.assert_allclose(result.transform(query), query + 1.0) + assert result.metadata == {"method": "mean_shift"} + + +def test_registry_register_get_keys() -> None: + reg = Registry("demo") + + @reg.register("mean_shift", "obs") + def _registered(ref: np.ndarray, query: np.ndarray) -> _MeanShiftResult: + return fit_mean_shift(ref, query) + + assert reg.keys() == ("mean_shift",) + assert reg.get("mean_shift").implementation("obs") is _registered + assert isinstance(reg.get("mean_shift").obs(np.ones((2, 2)), np.zeros((2, 2))), _MeanShiftResult) + + +def test_one_method_accumulates_several_modalities() -> None: + """A method is one record with a slot per modality, not one entry per modality.""" + reg = Registry("demo") + reg.register("multi", "obs")(fit_mean_shift) + reg.register("multi", "images")(fit_mean_shift) + + method = reg.get("multi") + assert reg.keys() == ("multi",) + assert method.supports() == ("obs", "images") + assert method.landmarks is None + + +def test_unsupported_modality_says_what_is_supported() -> None: + reg = Registry("demo") + reg.register("obs_only", "obs")(fit_mean_shift) + + with pytest.raises(ValueError, match="'obs_only' does not support landmarks alignment. It supports: obs"): + reg.get("obs_only").implementation("landmarks") + + +def test_supporting_filters_by_modality() -> None: + reg = Registry("demo") + reg.register("a", "obs")(fit_mean_shift) + reg.register("b", "landmarks")(fit_mean_shift) + reg.register("c", "obs")(fit_mean_shift) + + assert reg.supporting("obs") == ("a", "c") + assert reg.supporting("landmarks") == ("b",) + assert reg.supporting("images") == () + + +def test_registry_unknown_key_lists_available() -> None: + reg = Registry("demo") + reg.register("a", "obs")(fit_mean_shift) + + with pytest.raises(ValueError, match=r"Unknown demo method 'b'. Available: \['a'\]"): + reg.get("b") + + +def test_registry_rejects_duplicate_slot() -> None: + reg = Registry("demo") + reg.register("dup", "obs")(fit_mean_shift) + + with pytest.raises(ValueError, match="already has a obs estimator"): + reg.register("dup", "obs")(fit_mean_shift) + + +def test_registry_rejects_unknown_modality() -> None: + reg = Registry("demo") + with pytest.raises(ValueError, match="Unknown modality 'obsm'"): + reg.register("x", "obsm") # type: ignore[arg-type] + + +def test_align_method_supports_reports_filled_slots() -> None: + assert AlignMethod(name="empty").supports() == () + assert AlignMethod(name="x", obs=fit_mean_shift, landmarks=fit_mean_shift).supports() == ("obs", "landmarks") + + +def test_check_requirements_passes_when_none() -> None: + reg = Registry("demo") + # By default, registering without requires parameter does not wrap/check. + reg.register("mean_shift", "obs")(fit_mean_shift) + result = reg.get("mean_shift").implementation("obs")(np.ones((2, 2)), np.zeros((2, 2))) + assert isinstance(result, _MeanShiftResult) + + +def test_check_requirements_raises_for_missing_dependency() -> None: + reg = Registry("demo") + + @reg.register("needs_ghost", "obs", requires=("squidpy_nonexistent_pkg_xyz",)) + def _needs_ghost(ref: np.ndarray, query: np.ndarray) -> _MeanShiftResult: + return fit_mean_shift(ref, query) + + with pytest.raises( + ImportError, + match=r"Method 'needs_ghost' requires 'squidpy_nonexistent_pkg_xyz'.*squidpy\[squidpy_nonexistent_pkg_xyz\]", + ): + reg.get("needs_ghost").implementation("obs")(np.ones((2, 2)), np.zeros((2, 2))) diff --git a/tests/experimental/methods/test_stalign.py b/tests/experimental/methods/test_stalign.py new file mode 100644 index 000000000..cb00ef8ea --- /dev/null +++ b/tests/experimental/methods/test_stalign.py @@ -0,0 +1,334 @@ +"""Integration tests for the ported STalign estimator. + +Tiny synthetic fixtures with ``niter=1`` keep most of these fast; they verify wiring +and shapes (dispatch -> JAX LDDMM -> StalignResult), not solver quality. + +The two tests at the bottom do check solver quality, cheaply. They exist because the +full numerical comparison against the original STalign lives in +``test_stalign_reference.py``, which is deselected by default and only runs on the +scheduled job -- leaving a fortnight in which a real regression could ship green. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("jax") + +from squidpy.experimental.methods import ALIGN +from squidpy.experimental.methods.align_samples import StalignResult, fit_stalign_image, fit_stalign_obs + +# Flat solver kwargs (assembled into the config internally) -- smallest possible solve. +_TINY = {"dx": 0.5, "blur": 1.0, "a": 1.0, "expand": 1.0, "nt": 1, "niter": 1, "epV": 1.0} + + +def _points_xy() -> np.ndarray: + return np.array( + [ + [10.0, 1.0], + [12.0, 1.0], + [11.0, 2.0], + [10.0, 3.0], + [12.0, 3.0], + ] + ) + + +def test_stalign_registered_for_obs_and_images() -> None: + method = ALIGN.get("stalign") + assert method.supports() == ("obs", "images") + assert method.implementation("obs") is fit_stalign_obs + + +def test_stalign_declines_landmarks() -> None: + """A diffeomorphic solver has nothing to do with paired correspondences.""" + with pytest.raises(ValueError, match="does not support landmarks"): + ALIGN.get("stalign").implementation("landmarks") + + +def test_stalign_fit_returns_diffeomorphism() -> None: + ref, query = _points_xy(), _points_xy() + + result = fit_stalign_obs(ref, query, **_TINY) + + assert isinstance(result, StalignResult) + assert result.aligned_points.shape == query.shape + assert np.all(np.isfinite(np.asarray(result.aligned_points))) + assert result.affine.shape == (3, 3) + assert result.velocity.ndim == 4 + assert result.match_weights is not None + assert result.artifact_weights is not None + assert result.background_weights is not None + assert result.energies is not None + assert result.n_iter == 1 + + +def test_stalign_transform_matches_aligned_points() -> None: + ref, query = _points_xy(), _points_xy() + + result = fit_stalign_obs(ref, query, **_TINY) + + np.testing.assert_allclose(np.asarray(result.transform(query)), np.asarray(result.aligned_points), atol=1e-5) + + +def test_stalign_transform_accepts_arbitrary_points() -> None: + ref, query = _points_xy(), _points_xy() + result = fit_stalign_obs(ref, query, **_TINY) + + out = result.transform(np.zeros((1, 2))) + assert np.asarray(out).shape == (1, 2) + + +def test_stalign_transform_backward_inverts_forward() -> None: + ref, query = _points_xy(), _points_xy() + result = fit_stalign_obs(ref, query, **_TINY) + + forward = result.transform(query, direction="forward") + roundtrip = result.transform(forward, direction="backward") + np.testing.assert_allclose(np.asarray(roundtrip), query, atol=1e-3) + + +def test_stalign_transform_rejects_non_2d() -> None: + ref, query = _points_xy(), _points_xy() + result = fit_stalign_obs(ref, query, **_TINY) + + with pytest.raises(ValueError, match=r"Expected an \(N, 2\)"): + result.transform(np.zeros((5, 3))) + + +def test_stalign_fit_with_landmarks() -> None: + ref, query = _points_xy(), _points_xy() + landmarks = ref[:3] + + result = fit_stalign_obs(ref, query, landmarks_source=landmarks, landmarks_target=landmarks, **_TINY) + + assert result.aligned_points.shape == query.shape + + +def test_stalign_fit_with_explicit_affine_in_xy() -> None: + ref, query = _points_xy(), _points_xy() + affine = np.array([[1.0, 0.0, 3.0], [0.0, 1.0, 4.0], [0.0, 0.0, 1.0]]) + + result = fit_stalign_obs( + ref, query, initial_affine=affine, niter=0, **{k: v for k, v in _TINY.items() if k != "niter"} + ) + + np.testing.assert_allclose(np.asarray(result.aligned_points), query + np.array([3.0, 4.0]), atol=1e-5) + + +def test_stalign_image_accepts_explicit_axes_and_warps_both_directions() -> None: + image = np.arange(20, dtype=float).reshape(4, 5) + axes = (np.linspace(10.0, 13.0, 4), np.linspace(20.0, 24.0, 5)) + + result = fit_stalign_image(image, image, query_axes=axes, ref_axes=axes, niter=0, a=1.0, nt=1) + + np.testing.assert_allclose(np.asarray(result.warp_image(image, direction="forward"))[0], image, atol=1e-5) + np.testing.assert_allclose(np.asarray(result.warp_image(image, direction="backward"))[0], image, atol=1e-5) + assert result.deformation_grid(direction="forward").shape == (2, 4, 5) + assert result.deformation_grid(direction="backward").shape == (2, 4, 5) + + +def test_stalign_fit_rejects_non_2d_input() -> None: + with pytest.raises(ValueError, match=r"Expected `query` to have shape `\(n, 2\)`"): + fit_stalign_obs(_points_xy(), np.zeros((5, 3)), **_TINY) + + +def test_stalign_rejects_unknown_kwarg() -> None: + ref, query = _points_xy(), _points_xy() + with pytest.raises(TypeError, match="unexpected keyword argument"): + fit_stalign_obs(ref, query, not_a_real_param=1.0, **_TINY) + + +def test_lddmm_accepts_zero_iterations() -> None: + """``niter=0`` means "evaluate the initial state and stop", not a crash. + + ``energy`` and the transformed landmarks used to be bound only inside the loop, so + the return statement read unbound locals. + """ + from squidpy.experimental.methods.align_samples._stalign_impl._core import lddmm + from squidpy.experimental.methods.align_samples._stalign_impl._helpers import rasterize_cloud + + grid = rasterize_cloud(_points_xy()[:, ::-1], dx=0.5, blur=1.0, expand=1.1) + result = lddmm(*grid, *grid, L=np.eye(2), T=np.zeros(2), niter=0, a=1.0, nt=1) + + assert result["v"].shape[0] == 1 + assert result["A"].shape == (3, 3) + + +def test_lddmm_accepts_fixed_mixture_means() -> None: + from squidpy.experimental.methods.align_samples._stalign_impl._core import lddmm + from squidpy.experimental.methods.align_samples._stalign_impl._helpers import rasterize_cloud + + axes, image = rasterize_cloud(_points_xy()[:, ::-1], dx=0.5, blur=1.0, expand=1.1) + rgb = np.repeat(np.asarray(image), 3, axis=0) + result = lddmm( + axes, + rgb, + axes, + rgb, + L=np.eye(2), + T=np.zeros(2), + niter=1, + a=1.0, + nt=1, + muA=np.ones(3), + muB=np.zeros(3), + ) + + assert np.all(np.isfinite(np.asarray(result["A"]))) + + +def test_lddmm_rejects_wrong_mixture_mean_shape() -> None: + from squidpy.experimental.methods.align_samples._stalign_impl._core import lddmm + from squidpy.experimental.methods.align_samples._stalign_impl._helpers import rasterize_cloud + + grid = rasterize_cloud(_points_xy()[:, ::-1], dx=0.5, blur=1.0, expand=1.1) + with pytest.raises(ValueError, match=r"Expected `muA` to have shape \(1,\)"): + lddmm(*grid, *grid, L=np.eye(2), T=np.zeros(2), niter=0, a=1.0, nt=1, muA=np.ones(3)) + + +def test_default_dtype_is_unchanged() -> None: + """Guards the reference suite's blast radius. + + ``test_stalign_reference.py`` needs float64, but ``jax.config.update`` is + process-global and every xdist worker imports every module in this directory. If + that suite ever starts enabling x64 itself instead of reading ``JAX_ENABLE_X64`` + from the environment, these tests would silently stop testing float32 -- and this + assertion is what would catch it. + """ + import jax + import jax.numpy as jnp + + from squidpy.experimental.methods.align_samples._stalign_impl._core import jax_dtype + + expected = jnp.float64 if jax.config.jax_enable_x64 else jnp.float32 + assert jax_dtype() == expected + + +# --- solver quality, cheap enough to run on every PR ---------------------------------- + +# A cloud with enough structure to be alignable, and a query that is a known rigid +# transform of it. Small enough that a few hundred iterations run in about a second. +_SOLVE = {"dx": 8.0, "blur": 1.0, "a": 40.0, "nt": 3, "epV": 5e2} + + +def _blobs(seed: int = 0) -> np.ndarray: + rng = np.random.default_rng(seed) + return np.concatenate( + [ + rng.normal(scale=(18.0, 11.0), size=(120, 2)), + rng.normal(scale=8.0, size=(60, 2)) + np.array([55.0, 34.0]), + ] + ) + + +def _solve_grids(): + from squidpy.experimental.methods.align_samples._stalign_impl._helpers import rasterize_cloud + + ref = _blobs() + query = ref @ np.array([[np.cos(0.12), -np.sin(0.12)], [np.sin(0.12), np.cos(0.12)]]).T + np.array([6.0, -4.0]) + source = rasterize_cloud(query[:, ::-1], dx=_SOLVE["dx"], blur=_SOLVE["blur"], expand=1.1) + target = rasterize_cloud(ref[:, ::-1], dx=_SOLVE["dx"], blur=_SOLVE["blur"], expand=1.1) + solve = {k: v for k, v in _SOLVE.items() if k not in {"dx", "blur"}} + return source, target, {"L": np.eye(2), "T": np.zeros(2), **solve} + + +def test_lddmm_returns_an_energy_trace() -> None: + """Without this a caller cannot tell a converged run from a diverged one.""" + from squidpy.experimental.methods.align_samples._stalign_impl._core import lddmm + + source, target, common = _solve_grids() + result = lddmm(*source, *target, niter=30, **common) + + assert result["energies"].shape == (30,) + assert int(result["n_iter"]) == 30 + assert np.all(np.isfinite(np.asarray(result["energies"]))) + # The reported final energy is the last one in the trace. + np.testing.assert_allclose(float(result["E"]), float(result["energies"][-1]), rtol=1e-12) + + +def test_early_stopping_is_off_by_default() -> None: + """``tol=None`` must run every iteration, so the default is bit-for-bit unchanged.""" + from squidpy.experimental.methods.align_samples._stalign_impl._core import lddmm + + source, target, common = _solve_grids() + plain = lddmm(*source, *target, niter=40, **common) + explicit = lddmm(*source, *target, niter=40, tol=None, **common) + + assert int(plain["n_iter"]) == 40 + np.testing.assert_array_equal(np.asarray(plain["v"]), np.asarray(explicit["v"])) + + +def test_early_stopping_never_fires_before_the_weights_switch_on() -> None: + """The objective changes definition at iteration 50 and its value jumps upward. + + A stopping rule that compares across that boundary sees the jump as "no longer + improving" and quits immediately, which is what a naive implementation does. + """ + from squidpy.experimental.methods.align_samples._stalign_impl._core import MIXTURE_E_STEP_START, lddmm + + source, target, common = _solve_grids() + # A tolerance so loose it would stop at the first opportunity. + result = lddmm(*source, *target, niter=400, tol=1e9, patience=25, **common) + + assert int(result["n_iter"]) >= MIXTURE_E_STEP_START + 2 + 25 + trace = np.asarray(result["energies"]) + assert trace[MIXTURE_E_STEP_START + 1] > trace[MIXTURE_E_STEP_START], ( + "expected the objective to jump when the mixture-weight E step engages" + ) + + +def test_early_stopping_shortens_the_run() -> None: + from squidpy.experimental.methods.align_samples._stalign_impl._core import lddmm + + source, target, common = _solve_grids() + # Loose enough to fire on this fixture but not at the first opportunity: it stops + # around iteration 157 of 400, whereas tol=1e-2 keeps improving through all 400. + stopped = lddmm(*source, *target, niter=400, tol=1e-1, patience=25, **common) + + ran = int(stopped["n_iter"]) + assert 77 < ran < 400 + # Iterations that never ran stay NaN rather than reporting a bogus energy. + trace = np.asarray(stopped["energies"]) + assert np.all(np.isfinite(trace[:ran])) and np.all(np.isnan(trace[ran:])) + + +def test_lddmm_energy_decreases() -> None: + """More iterations must buy a lower objective. Catches a broken gradient or step.""" + from squidpy.experimental.methods.align_samples._stalign_impl._core import lddmm + from squidpy.experimental.methods.align_samples._stalign_impl._helpers import rasterize_cloud + + ref = _blobs() + query = ref @ np.array([[np.cos(0.12), -np.sin(0.12)], [np.sin(0.12), np.cos(0.12)]]).T + np.array([6.0, -4.0]) + + source = rasterize_cloud(query[:, ::-1], dx=_SOLVE["dx"], blur=_SOLVE["blur"], expand=1.1) + target = rasterize_cloud(ref[:, ::-1], dx=_SOLVE["dx"], blur=_SOLVE["blur"], expand=1.1) + solve = {k: v for k, v in _SOLVE.items() if k not in {"dx", "blur"}} + + common = {"L": np.eye(2), "T": np.zeros(2), **solve} + first = lddmm(*source, *target, niter=1, **common)["E"] + later = lddmm(*source, *target, niter=60, **common)["E"] + + assert float(later) < float(first), f"energy did not decrease: {float(first)} -> {float(later)}" + + +def test_lddmm_recovers_a_known_rigid_transform() -> None: + """The query is a known rotation+translation of the reference; alignment must undo it. + + Asserts on the median rather than the mean: rasterisation at ``dx=8`` quantises the + clouds, so a handful of points on the sparse outskirts stay poorly matched however + good the fit is. + """ + ref = _blobs() + angle = 0.12 + rotation = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]) + query = ref @ rotation.T + np.array([6.0, -4.0]) + + result = fit_stalign_obs(ref, query, niter=250, **_SOLVE) + residual = np.linalg.norm(np.asarray(result.aligned_points) - ref, axis=1) + + before = np.median(np.linalg.norm(query - ref, axis=1)) + after = float(np.median(residual)) + assert after < before / 2.0, f"alignment barely improved: median {before:.2f} -> {after:.2f}" + assert after < _SOLVE["dx"], f"median residual {after:.2f} exceeds one grid cell ({_SOLVE['dx']})" diff --git a/tests/experimental/methods/test_stalign_reference.py b/tests/experimental/methods/test_stalign_reference.py new file mode 100644 index 000000000..91b1b0dcd --- /dev/null +++ b/tests/experimental/methods/test_stalign_reference.py @@ -0,0 +1,892 @@ +"""Numerical comparison of the STalign port against the original implementation. + +Implements scverse/squidpy#1243. The reference values in ``tests/_data/stalign_reference`` +are produced out of band by https://github.com/theislab/squidpy-ports, which vendors the +PyTorch STalign at a pinned commit -- so torch never becomes a squidpy dependency. + +These are **excluded from normal runs**. See ``STALIGN_DIVERGENCES.md`` next to this file +for what each comparison found and why the known-divergent ones are pinned rather than +fixed here. + +Run them with:: + + JAX_ENABLE_X64=1 pytest tests/experimental/methods/test_stalign_reference.py -m reference +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") + +import jax.numpy as jnp # noqa: E402 + +from squidpy.experimental.methods.align_samples import ( # noqa: E402 + StalignResult, + fit_stalign_image, + fit_stalign_obs, +) +from squidpy.experimental.methods.align_samples._stalign_impl import _core, _helpers # noqa: E402 + +from . import _stalign_fixtures as F # noqa: E402 + +pytestmark = pytest.mark.reference + +DATA = Path(__file__).parent.parent.parent / "_data" / "stalign_reference" +LEDGER = Path(__file__).parent / "STALIGN_DIVERGENCES.md" + +# Tolerances are calibrated from the measured gaps recorded in STALIGN_DIVERGENCES.md, +# not guessed. Everything the port reproduces faithfully lands at 1e-15 or better, so a +# 1e-12 budget leaves three orders of headroom while still being three orders tighter +# than any divergence we are pinning. +EXACT = 1e-12 + + +def _skip_without_x64() -> None: + """float64 is a precondition, and it cannot be turned on from inside a test. + + ``jax.config.update`` is process-global; under ``-n auto`` every worker imports every + test module, so setting it here would silently flip the float32 tests in + ``test_stalign.py`` to float64 in the same worker. It has to come from the + environment. + """ + if not jax.config.jax_enable_x64: + pytest.skip( + "STalign reference comparison needs float64 (upstream is double throughout). " + "Set JAX_ENABLE_X64=1 in the environment -- not via jax.config.update, which " + "would corrupt the float32 tests in this directory." + ) + + +def _load(name: str) -> np.lib.npyio.NpzFile: + path = DATA / f"{name}.npz" + if not path.is_file(): + pytest.skip(f"reference bundle missing at {path}; regenerate it with theislab/squidpy-ports") + return np.load(path) + + +def rel(actual, expected) -> float: + """Relative L2 error, the single measure used throughout.""" + actual, expected = np.asarray(actual, float), np.asarray(expected, float) + denominator = np.linalg.norm(expected) + return float(np.linalg.norm(actual - expected) / max(denominator, np.finfo(float).tiny)) + + +@pytest.fixture(scope="module", autouse=True) +def _x64(): + _skip_without_x64() + + +@pytest.fixture(scope="module") +def primitives(): + return _load("primitives") + + +@pytest.fixture(scope="module") +def clouds(): + return F.make_clouds() + + +@pytest.fixture(scope="module") +def source_grid(primitives): + """The *moving* raster and its axes, as ``((y, x), image)``. + + The query moves onto the reference, so the query raster is the source -- the same + role LDDMM's ``I``/``xI`` and ``pointsI`` play. + """ + return (jnp.asarray(primitives["raster_query_y"]), jnp.asarray(primitives["raster_query_x"])), jnp.asarray( + primitives["raster_query"] + ) + + +@pytest.fixture(scope="module") +def target_grid(primitives): + """The *fixed* raster and its axes: the reference.""" + return (jnp.asarray(primitives["raster_ref_y"]), jnp.asarray(primitives["raster_ref_x"])), jnp.asarray( + primitives["raster_ref"] + ) + + +@pytest.fixture(scope="module") +def velocity_grid(primitives): + """Upstream's own velocity grid, which ``_lddmm_loss`` accepts directly.""" + return jnp.asarray(primitives["xv_upstream_0"]), jnp.asarray(primitives["xv_upstream_1"]) + + +@pytest.fixture(scope="module") +def fitted(primitives, velocity_grid): + """A :class:`StalignResult` carrying upstream's own fitted deformation. + + Point transforms are compared through the public result object rather than the + internal row-col helper, so what is pinned is what callers actually reach. + """ + return StalignResult( + affine=jnp.asarray(primitives["to_A"]), + velocity=jnp.asarray(primitives["velocity"]), + velocity_grid=velocity_grid, + aligned_points=jnp.zeros((0, 2)), + ) + + +def _transform_rc(result: StalignResult, points_rc, *, direction: str = "forward") -> np.ndarray: + """``result.transform`` on row-col points, for comparison with upstream. + + The public API speaks ``(x, y)`` and the reference speaks row-col, so the flip that + ``transform`` performs internally is undone on both sides here. Anything the flip + itself got wrong would still show up -- it is applied, not bypassed. + """ + got = result.transform(np.asarray(points_rc)[:, ::-1], direction=direction) + return np.asarray(got)[:, ::-1] + + +# -------------------------------------------------------------------------------------- +# Provenance -- the fixtures have to stay falsifiable +# -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "name", + [ + "primitives", + "energy", + "gradients", + "trajectory_n1", + "trajectory_n5", + "trajectory_n50", + "converged_n500", + "image_trajectory", + "image_trajectory_matched", + ], +) +def test_every_fixture_carries_provenance(name): + """Without this the .npz files are unfalsifiable magic numbers within a year.""" + payload = json.loads(str(_load(name)["__provenance__"])) + assert payload["upstream_sha"] == "b2068edc98974efa54537eca194736e177bbe11d" + for key in ("ports_commit", "torch", "numpy", "python", "platform", "generated_utc"): + assert payload[key], f"{name}: empty provenance field {key!r}" + + +def test_fixture_definitions_have_not_drifted(): + """The generator and these tests must build inputs from the same file, byte for byte. + + ``_stalign_fixtures.py`` is vendored from squidpy-ports. If someone edits one copy + without regenerating, every comparison below silently compares different inputs. + """ + recorded = json.loads(str(_load("primitives")["__provenance__"]))["fixtures_checksum"] + assert F.checksum() == recorded, ( + "_stalign_fixtures.py differs from the copy the reference bundle was generated " + "with. Re-sync it from squidpy-ports and regenerate the bundle." + ) + + +# -------------------------------------------------------------------------------------- +# Fixture validation -- runs before anything that depends on it +# -------------------------------------------------------------------------------------- + + +def test_fixture_samples_are_off_grid(primitives): + """No interpolation sample may sit on a grid line. See ledger row D10. + + Upstream and squidpy compute the fractional index by formulas that agree to ~1 ulp; + exactly on a grid line they can floor() to different neighbours, giving an O(1) + difference that says nothing about the port. + """ + coords = np.asarray(primitives["interp_coords"]) + axes = (np.asarray(primitives["raster_ref_y"]), np.asarray(primitives["raster_ref_x"])) + for axis, values in zip(axes, coords, strict=True): + fractional = np.abs(np.modf((values - axis[0]) / (axis[1] - axis[0]))[0]) + assert np.minimum(fractional, 1.0 - fractional).min() > 1e-6 + + +def test_fixture_stays_inside_velocity_grid(primitives, velocity_grid): + """Keeps the padding divergence (D5) out of every test that is not about it.""" + points = np.asarray(primitives["points"]) + for axis, column in zip(velocity_grid, points.T, strict=True): + assert column.min() > float(axis[0]) and column.max() < float(axis[-1]) + + +# -------------------------------------------------------------------------------------- +# Primitives that the port reproduces faithfully +# -------------------------------------------------------------------------------------- + + +def test_interp_matches_upstream(primitives, source_grid, record_property): + """``_interp`` vs ``STalign.interp(padding_mode='border')``.""" + axes, image = source_grid + got = _core._interp(axes, image, jnp.asarray(primitives["interp_coords"])) + error = rel(got, primitives["interp_border"]) + record_property("rel_error", error) + assert error < EXACT + + +def test_regularizer_matches_upstream(primitives, velocity_grid, record_property): + """``LL``/``K``/``DV`` vs STalign.py:1078-1090. + + This is a precondition, not a nicety: the regulariser sets the scale of the whole + velocity term, so if it disagreed every later comparison would be meaningless. + """ + kernel, ll, dv_prod = _core._build_regularizer(velocity_grid, a=F.LDDMM_PARAMS["a"], p=F.LDDMM_PARAMS["p"]) + for name, got, expected in ( + ("LL", ll, primitives["regularizer_LL"]), + ("K", kernel, primitives["regularizer_K"]), + ("DV", dv_prod, primitives["regularizer_DV"]), + ): + error = rel(got, expected) + record_property(f"rel_error_{name}", error) + assert error < EXACT, name + + +def test_transform_grid_backward_matches_upstream(primitives, target_grid, velocity_grid, record_property): + """``_transform_grid_backward`` vs ``STalign.build_transform(direction='b')``. + + This is the inner loop of the objective: invert the affine, then integrate ``-v`` + backwards in time. Upstream returns ``(H, W, 2)``; squidpy returns ``(2, H, W)``. + """ + axes, _ = target_grid + got = _core._transform_grid_backward( + axes, velocity_grid, jnp.asarray(primitives["velocity"]), jnp.asarray(primitives["to_A"]) + ) + error = rel(np.moveaxis(np.asarray(got), 0, -1), primitives["grid_backward"]) + record_property("rel_error", error) + assert error < EXACT + + +def test_warp_image_uses_the_upstream_grid(primitives, source_grid, target_grid, velocity_grid, record_property): + """``StalignResult.warp_image`` vs sampling on *upstream's own* backward grid. + + This is what ``align(by="images", out="images/...")`` materialises, and it had no + reference coverage: the two halves were each compared to upstream separately, but + nothing checked that the public method composes them the way upstream does. Feeding + upstream's ``grid_backward`` through the same interpolation isolates the composition. + """ + source_axes, source_image = source_grid + target_axes, _ = target_grid + result = StalignResult( + affine=jnp.asarray(primitives["to_A"]), + velocity=jnp.asarray(primitives["velocity"]), + velocity_grid=velocity_grid, + aligned_points=jnp.zeros((0, 2)), + query_axes=source_axes, + ref_axes=target_axes, + ) + upstream_grid = jnp.asarray(np.moveaxis(np.asarray(primitives["grid_backward"]), -1, 0)) + expected = _core._interp(source_axes, source_image, upstream_grid) + + error = rel(result.warp_image(source_image), expected) + record_property("rel_error", error) + assert error < EXACT + + +def test_transform_points_forward_matches_upstream(fitted, primitives, record_property): + """``StalignResult.transform`` vs ``STalign.transform_points_source_to_target``.""" + got = _transform_rc(fitted, primitives["points"]) + error = rel(got, primitives["points_forward"]) + record_property("rel_error", error) + assert error < EXACT + + +def test_to_affine_matches_upstream(primitives): + """``_to_affine`` vs ``STalign.to_A``.""" + got = _core._to_affine(jnp.asarray(primitives["to_A_linear"]), jnp.asarray(primitives["to_A_translation"])) + np.testing.assert_allclose(np.asarray(got), primitives["to_A"], rtol=0, atol=0) + + +# -------------------------------------------------------------------------------------- +# The objective itself +# -------------------------------------------------------------------------------------- + + +def _loss_arguments(primitives, source_grid, target_grid, velocity_grid, *, nt, with_points): + source_axes, source_image = source_grid + target_axes, target_image = target_grid + kernel, ll, dv_prod = _core._build_regularizer(velocity_grid, a=F.LDDMM_PARAMS["a"], p=F.LDDMM_PARAMS["p"]) + empty = jnp.zeros((0, 2)) + landmarks_source = jnp.asarray(primitives["landmarks_query"])[:, ::-1] + landmarks_target = jnp.asarray(primitives["landmarks_ref"])[:, ::-1] + return ( + kernel, + { + "x_source": source_axes, + "source_image": source_image, + "x_target": target_axes, + "target_image": target_image, + "xv": velocity_grid, + # LDDMM's initial state: uniform 0.5 matching weights (STalign.py:1102). + "match_weights": jnp.full(target_image.shape[1:], 0.5), + "ll": ll, + "dv_prod": dv_prod, + "points_source": landmarks_source if with_points else empty, + "points_target": landmarks_target if with_points else empty, + "sigmaM": F.LDDMM_PARAMS["sigmaM"], + "sigmaR": F.LDDMM_PARAMS["sigmaR"], + "sigmaP": F.LDDMM_PARAMS["sigmaP"], + }, + jnp.zeros((nt, velocity_grid[0].shape[0], velocity_grid[1].shape[0], 2)), + ) + + +@pytest.mark.parametrize("nt", [1, 3]) +@pytest.mark.parametrize("with_points", [False, True]) +@pytest.mark.parametrize("warm", [False, True], ids=["v_zero", "v_nonzero"]) +def test_energy_matches_upstream( + primitives, source_grid, target_grid, velocity_grid, nt, with_points, warm, record_property +): + """``_lddmm_loss`` vs upstream's ``E`` at iteration 0. + + The strongest single result in this suite: it says the two implementations optimise + the *same function*. Divergence D3 is inert here because no gradient is taken. + + The ``v_nonzero`` half is not redundant. LDDMM starts at ``v = 0``, where the + regularisation term ``ER`` contributes exactly nothing -- so evaluating only there + would compare equal even if ``ER`` were negated. Verified: with ``ER``'s sign flipped + in ``_lddmm_loss``, only these cases fail. + """ + energy = _load("energy") + _, kwargs, velocity = _loss_arguments( + primitives, source_grid, target_grid, velocity_grid, nt=nt, with_points=with_points + ) + key = f"E_nt{nt}_{'points' if with_points else 'nopoints'}" + if warm: + velocity = jnp.asarray(energy[f"warm_velocity_nt{nt}"]) + key += "_warmv" + + total, _ = _core._lddmm_loss(jnp.asarray(energy["L"]), jnp.asarray(energy["T"]), velocity, **kwargs) + error = rel(total, energy[key]) + record_property("rel_error", error) + assert error < EXACT, f"squidpy {float(total)!r} vs upstream {float(energy[key])!r}" + + +def test_energy_budget_is_not_vacuous(primitives, source_grid, target_grid, velocity_grid): + """A tolerance that a perturbed run also passes is decoration, not a test. + + Nudge one physically meaningful knob and the same assertion must fail. + """ + energy = _load("energy") + _, kwargs, velocity = _loss_arguments(primitives, source_grid, target_grid, velocity_grid, nt=3, with_points=True) + kwargs["sigmaM"] *= 1.0001 + total, _ = _core._lddmm_loss(jnp.asarray(energy["L"]), jnp.asarray(energy["T"]), velocity, **kwargs) + assert rel(total, energy["E_nt3_points"]) > EXACT + + +# -------------------------------------------------------------------------------------- +# Gradients -- ledger row D3 +# -------------------------------------------------------------------------------------- + + +@pytest.fixture(scope="module", params=[False, True], ids=["v_zero", "v_nonzero"]) +def measured_gradients(request, primitives, source_grid, target_grid, velocity_grid): + """Gradients of ``_lddmm_loss``, at ``v = 0`` and at a non-zero velocity. + + Both are needed for the same reason as the energy test: ``dER/dv`` vanishes at the + origin, so gradients taken only there cannot see the regularisation term. + """ + gradients = _load("gradients") + kernel, kwargs, velocity = _loss_arguments( + primitives, source_grid, target_grid, velocity_grid, nt=F.LDDMM_PARAMS["nt"], with_points=True + ) + suffix = "" + if request.param: + velocity = jnp.asarray(gradients["warm_velocity"]) + suffix = "_warmv" + + (_, _), (grad_l, grad_t, grad_v) = jax.value_and_grad(_core._lddmm_loss, argnums=(0, 1, 2), has_aux=True)( + jnp.asarray(gradients["L"]), jnp.asarray(gradients["T"]), velocity, **kwargs + ) + # Upstream stores the Sobolev-smoothed velocity gradient, since that is what actually + # drives the step (STalign.py:1215). + smoothed = jnp.fft.ifftn(jnp.fft.fftn(grad_v, axes=(1, 2)) * kernel[None, ..., None], axes=(1, 2)).real + return gradients, {"L": grad_l, "T": grad_t, "v": smoothed}, suffix + + +@pytest.mark.parametrize(("component", "key"), [("L", "grad_L"), ("T", "grad_T"), ("v", "grad_v_smoothed")]) +def test_gradients_match_upstream(measured_gradients, component, key, record_property): + """``dE/dL``, ``dE/dT`` and the smoothed ``dE/dv`` vs upstream's, at iteration 0. + + Together with the energy test this pins both halves of the optimisation: the same + objective *and* the same search direction. See ledger row D3 -- these disagreed by + ~1.2e-3 until ``_contrast_transform`` stopped differentiating through the ridge solve. + """ + gradients, computed, suffix = measured_gradients + error = rel(computed[component], gradients[key + suffix]) + record_property("rel_error", error) + assert error < EXACT + + +def test_contrast_transform_freezes_the_ridge_coefficients(record_property): + """Regression guard for ledger row D3, stated without reference to upstream. + + The ridge fit is an EM M step solved exactly at the current estimate, so its + coefficients must be constant with respect to the optimisation. Dropping the + ``stop_gradient`` would leave the *value* untouched and change only the gradient -- + invisible to any test that checks outputs rather than derivatives, which is how this + survived unnoticed in the first place. + """ + rng = np.random.default_rng(F.SEED) + warped = jnp.asarray(rng.normal(size=(3, 12, 15)) ** 2) + target = jnp.asarray(rng.normal(size=(3, 12, 15)) ** 2) + weights = jnp.asarray(rng.uniform(0.2, 0.8, size=(12, 15))) + + def live(x): + return jnp.sum(_core._contrast_transform(x, target, weights) ** 2) + + def leaky(x): + # _contrast_transform verbatim, minus the stop_gradient: the bug as it was. + flat_source = x.reshape(x.shape[0], -1) + flat_target = target.reshape(target.shape[0], -1) + design = jnp.concatenate((jnp.ones((1, flat_source.shape[1]), dtype=x.dtype), flat_source), axis=0) + weighted = design * weights.reshape(-1)[None, :] + coefficients = jnp.linalg.solve( + weighted @ design.T + 0.1 * jnp.eye(design.shape[0], dtype=x.dtype), + weighted @ flat_target.T, + ) + return jnp.sum((coefficients.T @ design).reshape(target.shape) ** 2) + + # Identical value, different gradient -- which is exactly why this needs its own test. + np.testing.assert_allclose(float(live(warped)), float(leaky(warped)), rtol=1e-12) + error = rel(jax.grad(live)(warped), jax.grad(leaky)(warped)) + record_property("rel_error_vs_leaky", error) + assert error > 1e-6, "the ridge coefficients are being differentiated through again" + + +# -------------------------------------------------------------------------------------- +# Pinned divergences +# -------------------------------------------------------------------------------------- + + +def test_rasterize_grid_matches_upstream(primitives, clouds): + """Ledger row D2, fixed: the raster axes are identical to upstream's.""" + grid_x, grid_y, _ = _helpers.rasterize(clouds.ref[:, 0], clouds.ref[:, 1], **F.RASTER_PARAMS) + np.testing.assert_allclose(grid_x, primitives["raster_ref_x"], rtol=0, atol=1e-9) + np.testing.assert_allclose(grid_y, primitives["raster_ref_y"], rtol=0, atol=1e-9) + + +def test_velocity_grid_matches_upstream(primitives, source_grid): + """Ledger row D2, fixed: same off-by-one, same fix, in ``_build_velocity_grid``.""" + axes, _ = source_grid + built = _core._build_velocity_grid(axes, a=F.LDDMM_PARAMS["a"], expand=F.LDDMM_PARAMS["expand"]) + np.testing.assert_allclose(np.asarray(built[0]), primitives["xv_upstream_0"], rtol=0, atol=1e-9) + np.testing.assert_allclose(np.asarray(built[1]), primitives["xv_upstream_1"], rtol=0, atol=1e-9) + + +@pytest.mark.parametrize("step", [30.0, 7.3, 0.017]) +def test_grid_length_is_stable_against_float_rounding(step): + """The other half of D2: ``np.arange(lo, hi + step, step)`` had an unstable length. + + Deriving the count from the interval first makes it exact for every case; the old + form emitted one extra point sometimes and two others, depending on how ``hi`` + happened to round. + """ + for start, n in ((-400.4123, 33), (0.1, 41), (-1234.567, 77)): + assert _helpers._axis(start, start + n * step, step).size == n + + +@pytest.mark.xfail( + strict=True, + reason=( + "ledger row D6: squidpy integrates the backward flow in reversed time order, " + "upstream does not (STalign.py:1828-1843). squidpy is correct -- see " + "test_backward_transform_inverts_better, which is the assertion that matters. " + "This xfail exists to pin the literal difference, and should be deleted only if " + "squidpy ever deliberately adopts upstream's ordering." + ), +) +def test_transform_points_backward_matches_upstream(fitted, primitives): + got = _transform_rc(fitted, primitives["points"], direction="backward") + assert rel(got, primitives["points_backward"]) < EXACT + + +def test_backward_transform_inverts_better(fitted, primitives, record_property): + """D6, stated usefully: squidpy's backward map is the better inverse of the forward one.""" + points = np.asarray(primitives["points"]) + roundtrip = _transform_rc(fitted, primitives["points_forward"], direction="backward") + ours = rel(roundtrip, points) + theirs = rel(primitives["points_roundtrip"], points) + record_property("roundtrip_squidpy", ours) + record_property("roundtrip_upstream", theirs) + assert ours <= theirs + + +@pytest.mark.xfail( + strict=True, + reason=( + "ledger row D5: upstream samples the velocity field with grid_sample's default " + "padding_mode='zeros' (STalign.py:1163, :1167), squidpy uses " + "map_coordinates(mode='nearest'). squidpy is correct -- zeros make a point that " + "drifts off the velocity grid snap to no displacement at all." + ), +) +def test_interp_outside_domain_matches_upstream(primitives, source_grid): + axes, image = source_grid + got = _core._interp(axes, image, jnp.asarray(primitives["interp_coords_outside"])) + assert rel(got, primitives["interp_zeros_outside"]) < EXACT + + +def test_interp_outside_domain_is_border_padding(primitives, source_grid, record_property): + """The positive half of D5: squidpy's behaviour is exactly upstream's 'border' mode.""" + axes, image = source_grid + got = _core._interp(axes, image, jnp.asarray(primitives["interp_coords_outside"])) + error = rel(got, primitives["interp_border_outside"]) + record_property("rel_error", error) + assert error < EXACT + assert rel(got, primitives["interp_zeros_outside"]) > 0.1 + + +# -------------------------------------------------------------------------------------- +# Budgeted divergences +# -------------------------------------------------------------------------------------- + +#: Per-blur relative-L2 budget for the rasteriser (ledger row D1). squidpy bins onto a +#: grid and convolves once; upstream splats an exact sub-pixel Gaussian per point and +#: renormalises it over a truncated window. Measured 4.08 % / 0.81 % / 2.87 % at blur +#: 2.0 / 1.0 / 0.5 -- 6 % leaves headroom without going vacuous. +RASTER_BUDGET = 0.06 + + +def test_rasterize_stays_within_budget(primitives, clouds, record_property): + """D1 is a deliberate speedup, so it gets a measured budget rather than equality.""" + _, _, got = _helpers.rasterize(clouds.ref[:, 0], clouds.ref[:, 1], **F.RASTER_PARAMS) + expected = primitives["raster_ref"] + got = np.asarray(got) + assert got.shape == expected.shape + + for index, blur in enumerate(F.RASTER_PARAMS["blur"]): + error = rel(got[index], expected[index]) + correlation = float(np.corrcoef(got[index].ravel(), expected[index].ravel())[0, 1]) + record_property(f"rel_error_blur{blur}", error) + assert error < RASTER_BUDGET, f"blur={blur}: relL2 {error:.4%} exceeds {RASTER_BUDGET:.0%}" + assert correlation > 0.99, f"blur={blur}: correlation {correlation:.5f}" + + +def test_rasterize_conserves_mass(clouds, record_property): + """Every point contributes exactly one unit, wherever it sits. + + Upstream renormalises each point's kernel over its (possibly clipped) window, so a + point near the border still carries unit mass. A plain ``mode="constant"`` blur does + not, and used to lose 3 % of the total at the coarsest scale -- a density biased low + around the whole rim. + """ + n_points = clouds.ref.shape[0] + _, _, got = _helpers.rasterize(clouds.ref[:, 0], clouds.ref[:, 1], **F.RASTER_PARAMS) + + for index, blur in enumerate(F.RASTER_PARAMS["blur"]): + mass = float(np.asarray(got)[index].sum()) + record_property(f"mass_blur{blur}", mass) + assert mass == pytest.approx(n_points, rel=1e-9), f"blur={blur}: {mass:.2f} of {n_points}" + + +def _residual(linear, translation, source, target) -> float: + source, target = np.asarray(source), np.asarray(target) + return float(np.linalg.norm(source @ np.asarray(linear).T + np.asarray(translation) - target)) + + +def test_affine_from_points_is_equivalent_when_well_conditioned(primitives, record_property): + """D7: the two are different estimators, not the same one implemented twice. + + Upstream solves the normal equations for the plain least-squares fit; skimage solves + a Hartley-normalised homogeneous system by SVD, minimising algebraic rather than + geometric error. So their coefficients differ by ~1e-3 even on clean landmarks, and + asserting agreement would be wrong. What must hold is that neither is meaningfully + worse at the job. + """ + source = np.asarray(primitives["landmarks_query"])[:, ::-1] + target = np.asarray(primitives["landmarks_ref"])[:, ::-1] + # Reached through the public estimator: `niter=0` fits nothing, so the returned + # affine *is* the landmark initialisation. That also pins the wiring -- that + # `landmarks_*` actually reach the solver as its starting affine. + fit = fit_stalign_obs( + primitives["ref"], + primitives["query"], + landmarks_source=primitives["landmarks_query"], + landmarks_target=primitives["landmarks_ref"], + niter=0, + dx=F.RASTER_PARAMS["dx"], + blur=F.RASTER_PARAMS["blur"], + raster_expand=F.RASTER_PARAMS["expand"], + ) + affine = np.asarray(fit.affine) + linear, translation = affine[:2, :2], affine[:2, 2] + + ours = _residual(linear, translation, source, target) + theirs = _residual(primitives["lt_well_L"], primitives["lt_well_T"], source, target) + record_property("residual_squidpy", ours) + record_property("residual_upstream", theirs) + record_property("rel_error_L", rel(linear, primitives["lt_well_L"])) + assert abs(ours - theirs) / theirs < 1e-2, f"squidpy {ours:.6f} vs upstream {theirs:.6f}" + + +def test_affine_from_points_survives_ill_conditioning(primitives, record_property): + """D7, the half that matters: upstream's ``inv(XᵀX)`` collapses, skimage does not.""" + source, target = primitives["ill_src"], primitives["ill_dst"] + linear, translation = _helpers.affine_from_points(jnp.asarray(source), jnp.asarray(target)) + + ours = _residual(linear, translation, source, target) + theirs = _residual(primitives["lt_ill_L"], primitives["lt_ill_T"], source, target) + record_property("residual_squidpy", ours) + record_property("residual_upstream", theirs) + assert ours < theirs * 1e-6, ( + f"expected upstream to lose badly on near-collinear landmarks, got squidpy {ours:.3e} vs upstream {theirs:.3e}" + ) + + +# -------------------------------------------------------------------------------------- +# The whole iteration loop +# -------------------------------------------------------------------------------------- + + +def _run_lddmm(primitives, snapshot, niter): + source_axes = (jnp.asarray(primitives["raster_query_y"]), jnp.asarray(primitives["raster_query_x"])) + target_axes = (jnp.asarray(primitives["raster_ref_y"]), jnp.asarray(primitives["raster_ref_x"])) + return _core.lddmm( + source_axes, + jnp.asarray(primitives["raster_query"]), + target_axes, + jnp.asarray(primitives["raster_ref"]), + L=jnp.asarray(snapshot["L"]), + T=jnp.asarray(snapshot["T"]), + points_source=jnp.asarray(primitives["landmarks_query"])[:, ::-1], + points_target=jnp.asarray(primitives["landmarks_ref"])[:, ::-1], + niter=niter, + **F.LDDMM_PARAMS, + ) + + +#: Relative-error budget after ``n`` gradient steps. Single steps agree to ~1e-14; the +#: allowance grows because bilinear-resampling VJPs and FFTs accumulate in different +#: orders on the two backends, and 50 steps of gradient descent amplify that. +_TRAJECTORY_BUDGET = {1: 1e-10, 5: 1e-9, 50: 1e-6} + + +@pytest.mark.parametrize("niter", [1, 5, 50]) +def test_trajectory_matches_upstream(primitives, niter, record_property): + """Run the real loop for ``n`` steps and compare every state it carries. + + ``niter=50`` matters specifically: the mixture-weight E-step is gated on + ``it >= 50`` (STalign.py:1233), so below it the weights are frozen at their + initial 0.5/0.4/0.1 and that whole branch goes untested. + + Note the affine offset. Upstream builds ``A`` at the *top* of each iteration and + returns it, so ``LDDMM(n)["A"]`` reflects ``n-1`` updates; squidpy builds it after + the loop. The fixture stores the un-lagged affine as ``A``. See ledger row D4. + """ + snapshot = _load(f"trajectory_n{niter}") + result = _run_lddmm(primitives, snapshot, niter) + + budget = _TRAJECTORY_BUDGET[niter] + for name, got, expected in ( + ("A", result["A"], snapshot["A"]), + ("v", result["v"], snapshot["v"]), + ("WM", result["WM"], snapshot["WM"]), + ("WA", result["WA"], snapshot["WA"]), + ("WB", result["WB"], snapshot["WB"]), + ): + error = rel(got, expected) + record_property(f"rel_error_{name}", error) + assert error < budget, f"n={niter} {name}: {error:.3e} exceeds {budget:.0e}" + + +#: Matches ``IMAGE_PARAMS`` / ``IMAGE_ITERS`` in the generator. The image path works in +#: pixel units, so the kernel width and velocity step are far below the micron-scale +#: point-cloud defaults. +_IMAGE_PARAMS = {"a": 8.0, "p": 2.0, "expand": 2.0, "nt": 2, "diffeo_start": 4, "epV": 1.0} +_IMAGE_ITERS = 12 + + +@pytest.fixture(scope="module") +def image_reference(): + return _load("image_trajectory") + + +def test_image_path_axes_match_the_reference(image_reference): + """The image entry point's coordinate convention, checked before its results. + + ``fit_stalign_image`` centres pixel coordinates rather than using the point path's + physical microns. The solver agreeing on rasters says nothing about that convention, + so it is pinned separately -- if it drifted, the comparisons below would still pass + while aligning in a different frame. + """ + fit = fit_stalign_image(image_reference["ref"], image_reference["query"], niter=0, **_IMAGE_PARAMS) + + for got, expected in ( + (fit.query_axes[0], image_reference["source_axis_0"]), + (fit.query_axes[1], image_reference["source_axis_1"]), + (fit.ref_axes[0], image_reference["target_axis_0"]), + (fit.ref_axes[1], image_reference["target_axis_1"]), + ): + np.testing.assert_allclose(np.asarray(got), expected, rtol=0, atol=1e-12) + + +def _run_image(fixture, axes_ref: bool = True): + """Drive the solver on the image fixture, from the reference's own starting affine.""" + source = (jnp.asarray(fixture["source_axis_0"]), jnp.asarray(fixture["source_axis_1"])) + target = (jnp.asarray(fixture["target_axis_0"]), jnp.asarray(fixture["target_axis_1"])) + return _core.lddmm( + source, + jnp.asarray(fixture["query"]), + target, + jnp.asarray(fixture["ref"]), + L=fixture["start_L"], + T=fixture["start_T"], + niter=_IMAGE_ITERS, + **_IMAGE_PARAMS, + ) + + +@pytest.mark.parametrize("name", ["A", "v", "WM", "WA", "WB"]) +def test_image_path_matches_upstream(image_reference, name, record_property): + """The image entry point vs upstream's LDDMM on the same images, axes and start. + + Closes the last gap in the port's public surface: ``align(by="images")`` goes straight + through here and had no reference comparison at all. + + Note that ~11 % of the target grid samples the source through padding, because the two + rasters have different shapes and each is centred on its own centre. That agrees to + 1e-12 as well -- padding was the first suspect for an earlier disagreement and turned + out to be innocent on both values and gradients. + """ + result = _run_image(image_reference) + error = rel(result[name], image_reference[name]) + record_property("rel_error", error) + assert error < 1e-10, f"{name}: {error:.3e}" + + +def test_image_energy_trace_matches_upstream(image_reference, record_property): + """Every iteration of the objective, not just the endpoint. + + A trajectory can agree at the end while disagreeing throughout; comparing the whole + trace is what localises a disagreement to the step it starts at. + """ + result = _run_image(image_reference) + got = np.asarray(result["energies"]) + expected = np.asarray(image_reference["energies"]) + worst = max(abs(got[i] - expected[i]) / abs(expected[i]) for i in range(expected.size)) + record_property("worst_iteration_rel_error", worst) + assert worst < 1e-10 + + +def test_image_warp_matches_upstream(image_reference, record_property): + """``warp_image`` vs ``STalign.transform_image_source_to_target``. + + Exactly what ``align(by="images", out="images/...")`` writes, against upstream's own + image-warping composition rather than a reassembled one. + """ + result = _run_image(image_reference) + fit = StalignResult( + affine=result["A"], + velocity=result["v"], + velocity_grid=result["xv"], + aligned_points=jnp.zeros((0, 2)), + query_axes=(jnp.asarray(image_reference["source_axis_0"]), jnp.asarray(image_reference["source_axis_1"])), + ref_axes=(jnp.asarray(image_reference["target_axis_0"]), jnp.asarray(image_reference["target_axis_1"])), + ) + error = rel(fit.warp_image(jnp.asarray(image_reference["query"])), image_reference["warped"]) + record_property("rel_error", error) + assert error < 1e-9 + + +def test_on_grid_sampling_costs_six_orders_of_magnitude(record_property): + """Ledger row D10, measured rather than asserted. + + The control fixture crops both rasters to a common extent, so their axes are the + *same* integers and interpolation samples keep landing exactly on grid lines -- where + upstream and squidpy can floor() to different neighbours. Same solver, same inputs + otherwise, yet accuracy drops from 1e-12 to ~1e-3 on the velocity field. This is why + both fixtures start from a deliberately off-grid affine, and why the point-cloud + fixture asserts `test_fixture_samples_are_off_grid`. + """ + matched = _load("image_trajectory_matched") + axes = (jnp.asarray(matched["axis_0"]), jnp.asarray(matched["axis_1"])) + result = _core.lddmm( + axes, + jnp.asarray(matched["query"]), + axes, + jnp.asarray(matched["ref"]), + L=matched["start_L"], + T=matched["start_T"], + niter=_IMAGE_ITERS, + **_IMAGE_PARAMS, + ) + degraded = rel(result["v"], matched["v"]) + record_property("rel_error_on_grid", degraded) + assert 1e-6 < degraded < 1e-1, ( + f"on-grid sampling now costs {degraded:.3e}; if this has become exact, D10 is gone " + f"and the off-grid fixture design can be simplified" + ) + + +def test_velocity_grid_is_the_one_the_reference_used(primitives, source_grid): + """The trajectory comparison is only meaningful on a shared velocity grid. + + Upstream was driven onto this grid explicitly via its ``xv=``/``v=`` parameters, so + if squidpy's construction ever diverges again the trajectory numbers would be + comparing two different problems rather than two implementations. + """ + axes, _ = source_grid + built = _core._build_velocity_grid(axes, a=F.LDDMM_PARAMS["a"], expand=F.LDDMM_PARAMS["expand"]) + snapshot = _load("trajectory_n1") + np.testing.assert_allclose(np.asarray(built[0]), snapshot["xv_0"], rtol=0, atol=1e-9) + np.testing.assert_allclose(np.asarray(built[1]), snapshot["xv_1"], rtol=0, atol=1e-9) + + +def test_converged_solution_matches_upstream(primitives, record_property): + """500 iterations -- "enough to actually converge", per #1243. + + Elementwise equality is the wrong instrument this far in: 500 steps of descent + amplify last-ulp backend differences without either answer being wrong. What must + hold is that both converge to the same registration. + """ + snapshot = _load("converged_n500") + result = _run_lddmm(primitives, snapshot, 500) + + energy = rel(result["E"], snapshot["E_last"]) + record_property("rel_error_E", energy) + assert energy < 0.01, f"final energy differs by {energy:.3%}" + + # Convergence is a statement about the registration, not the objective: the affine + # initialised from landmarks is already near-optimal, so `E` barely moves over the + # run even though the fit is good. Target registration error is the honest measure. + clouds = F.make_clouds() + landmarks_ref = clouds.landmarks_ref_rc + before = float(np.mean(np.linalg.norm(clouds.landmarks_query_rc - landmarks_ref, axis=1))) + after = float(snapshot["tre_mean"]) + record_property("tre_before", before) + record_property("tre_after", after) + assert after < 0.4 * F.RASTER_PARAMS["dx"], f"reference did not converge: TRE {after:.2f}" + assert after < before / 5.0, f"reference barely moved: TRE {before:.2f} -> {after:.2f}" + + converged = StalignResult( + affine=result["A"], + velocity=result["v"], + velocity_grid=result["xv"], + aligned_points=jnp.zeros((0, 2)), + ) + aligned = _transform_rc(converged, np.asarray(primitives["query"])[:, ::-1]) + displacement = np.linalg.norm(aligned - np.asarray(snapshot["aligned_points_rc"]), axis=1) + percentile = float(np.percentile(displacement, 95)) + record_property("p95_displacement", percentile) + assert percentile < 0.1 * F.RASTER_PARAMS["dx"], ( + f"95th-percentile point disagreement {percentile:.4f} exceeds a tenth of a grid cell" + ) + + +# -------------------------------------------------------------------------------------- +# The ledger has to stay in sync +# -------------------------------------------------------------------------------------- + + +def test_divergences_doc_covers_all_xfails(): + """Every strict xfail cites a ledger row, and every cited row exists.""" + ledger = LEDGER.read_text() + documented = set(re.findall(r"\*\*(D\d+)\*\*", ledger)) + assert documented, "no ledger rows found; STALIGN_DIVERGENCES.md changed shape" + + source = Path(__file__).read_text() + cited = set(re.findall(r"ledger row (D\d+)", source)) + assert cited, "no test cites a ledger row" + + missing = cited - documented + assert not missing, f"tests cite ledger rows that do not exist: {sorted(missing)}" diff --git a/tests/experimental/tl/__init__.py b/tests/experimental/tl/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/experimental/tl/test_align.py b/tests/experimental/tl/test_align.py new file mode 100644 index 000000000..7885a6f98 --- /dev/null +++ b/tests/experimental/tl/test_align.py @@ -0,0 +1,297 @@ +"""Integration tests for the public ``align`` (STalign) function. + +Tiny synthetic fixtures with ``niter=1`` keep these fast; they verify wiring, +path resolution, and write-back -- not solver quality. +""" + +from __future__ import annotations + +import numpy as np +import pytest +from anndata import AnnData + +pytest.importorskip("jax") + +from squidpy.experimental.methods.align_samples import StalignResult +from squidpy.experimental.tl import align + +# Flat solver kwargs (assembled into the config internally) -- smallest possible solve. +_TINY = {"dx": 0.5, "blur": 1.0, "a": 1.0, "expand": 1.0, "nt": 1, "niter": 1, "epV": 1.0} +_TINY_IMAGE = {"a": 4.0, "nt": 1, "niter": 1, "epV": 1.0} + + +def _adata(*, key: str = "spatial") -> AnnData: + pts = np.array([[10.0, 1.0], [12.0, 1.0], [11.0, 2.0], [10.0, 3.0], [12.0, 3.0]]) + adata = AnnData(np.zeros((pts.shape[0], 1))) + adata.obsm[key] = pts + return adata + + +def _sdata_tables(**tables: AnnData): + sd = pytest.importorskip("spatialdata") + from spatialdata.models import TableModel + + return sd.SpatialData(tables={name: TableModel.parse(adata) for name, adata in tables.items()}) + + +def _sdata_images(): + sd = pytest.importorskip("spatialdata") + from spatialdata.models import Image2DModel + + ref = np.zeros((1, 16, 16), dtype=float) + ref[0, 5:11, 5:11] = 1.0 + query = np.zeros((1, 16, 16), dtype=float) + query[0, 6:12, 4:10] = 1.0 + return sd.SpatialData( + images={ + "ref": Image2DModel.parse(ref, dims=("c", "y", "x")), + "query": Image2DModel.parse(query, dims=("c", "y", "x")), + } + ) + + +# --- out=None: fit and return, touch nothing ------------------------------------------ + + +def test_out_none_returns_result_and_writes_nothing() -> None: + ref, query = _adata(), _adata() + result = align(ref, query, in_="obsm/spatial", method="stalign", **_TINY) + assert isinstance(result, StalignResult) + assert result.aligned_points.shape == query.obsm["spatial"].shape + assert list(query.obsm) == ["spatial"] + + +def test_result_satisfies_align_result_protocol() -> None: + from squidpy.experimental.tl import AlignResult + + assert isinstance(align(_adata(), _adata(), in_="obsm/spatial", **_TINY), AlignResult) + + +def test_public_surface_is_align_result_only() -> None: + import squidpy.experimental.tl as tl + + # `AlignResult` is the only result type exposed; concretes stay in their home modules. + assert "AlignResult" in tl.__all__ + assert not hasattr(tl, "StalignResult") + assert not hasattr(tl, "AffineFitResult") + + +# --- writing --------------------------------------------------------------------------- + + +def test_out_writes_in_place() -> None: + ref, query = _adata(), _adata() + assert align(ref, query, in_="obsm/spatial", out="obsm/aligned", **_TINY) is None + assert query.obsm["aligned"].shape == query.obsm["spatial"].shape + + +def test_copy_leaves_original_untouched() -> None: + ref, query = _adata(), _adata() + out = align(ref, query, in_="obsm/spatial", out="obsm/aligned", copy=True, **_TINY) + assert isinstance(out, AnnData) and out is not query + assert "aligned" in out.obsm + assert "aligned" not in query.obsm + + +def test_out_may_overwrite_the_input_path() -> None: + """``out`` equal to ``in_`` is allowed -- destructive, but explicitly asked for.""" + ref, query = _adata(), _adata() + original = query.obsm["spatial"].copy() + align(ref, query, in_="obsm/spatial", out="obsm/spatial", **_TINY) + assert list(query.obsm) == ["spatial"] + assert not np.array_equal(query.obsm["spatial"], original) + + +# --- SpatialData tables ---------------------------------------------------------------- + + +def test_sdata_pair_of_tables() -> None: + sdata = _sdata_tables(ref=_adata(), query=_adata()) + result = align( + sdata, + in_=("tables/ref/obsm/spatial", "tables/query/obsm/spatial"), + **_TINY, + ) + assert isinstance(result, StalignResult) + assert "aligned" not in sdata.tables["query"].obsm + + +def test_sdata_writes_into_the_named_table_only() -> None: + sdata = _sdata_tables(ref=_adata(), query=_adata()) + align( + sdata, + in_=("tables/ref/obsm/spatial", "tables/query/obsm/spatial"), + out="tables/query/obsm/aligned", + **_TINY, + ) + assert "aligned" in sdata.tables["query"].obsm + assert "aligned" not in sdata.tables["ref"].obsm + + +def test_sdata_copy_leaves_original_untouched() -> None: + sd = pytest.importorskip("spatialdata") + + sdata = _sdata_tables(ref=_adata(), query=_adata()) + out = align( + sdata, + in_=("tables/ref/obsm/spatial", "tables/query/obsm/spatial"), + out="tables/query/obsm/aligned", + copy=True, + **_TINY, + ) + assert isinstance(out, sd.SpatialData) and out is not sdata + assert "aligned" in out.tables["query"].obsm + assert "aligned" not in sdata.tables["query"].obsm + + +def test_two_sdata_objects_share_one_path() -> None: + ref_sdata = _sdata_tables(slice=_adata()) + query_sdata = _sdata_tables(slice=_adata()) + align(ref_sdata, query_sdata, in_="tables/slice/obsm/spatial", out="tables/slice/obsm/aligned", **_TINY) + assert "aligned" in query_sdata.tables["slice"].obsm + assert "aligned" not in ref_sdata.tables["slice"].obsm + + +# --- images ---------------------------------------------------------------------------- + + +def test_images_fit_returns_result() -> None: + sdata = _sdata_images() + result = align(sdata, in_=("images/ref", "images/query"), by="images", **_TINY_IMAGE) + assert isinstance(result, StalignResult) + assert sorted(sdata.images) == ["query", "ref"] + + +def test_images_out_materialises_a_warped_image() -> None: + """A diffeomorphism has no SpatialData transformation to be registered as.""" + sdata = _sdata_images() + expected = np.asarray(sdata.images["query"].data).shape + + align(sdata, in_=("images/ref", "images/query"), by="images", out="images/query_aligned", **_TINY_IMAGE) + + assert "query_aligned" in sdata.images + assert np.asarray(sdata.images["query_aligned"].data).shape == expected + + +def test_image_alignment_recovers_a_known_shift() -> None: + """The query square sits one row down and one column left of the reference.""" + sdata = _sdata_images() + ref = np.asarray(sdata.images["ref"].data) + query = np.asarray(sdata.images["query"].data) + + align(sdata, in_=("images/ref", "images/query"), by="images", out="images/query_aligned", a=4.0, nt=2) + aligned = np.asarray(sdata.images["query_aligned"].data) + + before = float(np.sum((query - ref) ** 2)) + after = float(np.sum((aligned - ref) ** 2)) + assert after < before / 2.0, f"overlap barely improved: {before:.1f} -> {after:.1f}" + + +# --- `by` ------------------------------------------------------------------------------ + + +def test_by_defaults_to_obs() -> None: + ref, query = _adata(), _adata() + explicit = align(ref, query, in_="obsm/spatial", by="obs", **_TINY) + assert isinstance(explicit, StalignResult) + + +def test_by_images_needs_an_image_path() -> None: + """Caught here rather than as a shape error from inside the solver.""" + ref, query = _adata(), _adata() + with pytest.raises(ValueError, match="`by='images'` needs an image path"): + align(ref, query, in_="obsm/spatial", by="images", **_TINY) + + +def test_by_obs_rejects_an_image_path() -> None: + sdata = _sdata_images() + with pytest.raises(ValueError, match="needs an \\(N, 2\\) coordinate path"): + align(sdata, in_=("images/ref", "images/query"), by="obs", **_TINY) + + +def test_unknown_by_lists_the_modalities() -> None: + ref, query = _adata(), _adata() + with pytest.raises(ValueError, match="Unknown `by='nope'`.*obs, images, landmarks"): + align(ref, query, in_="obsm/spatial", by="nope", **_TINY) + + +def test_warp_image_rejects_a_point_cloud_fit() -> None: + result = align(_adata(), _adata(), in_="obsm/spatial", **_TINY) + with pytest.raises(ValueError, match="fitted on point clouds"): + result.warp_image(np.zeros((1, 4, 4))) + + +# --- path errors ----------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("path", "match"), + [ + ("nope/x", "does not start with a known collection"), + ("obsm", "Expected `obsm/`"), + ("tables/t/spatial", "Expected `tables/
/obsm/`"), + ("images/a/b", "Expected `images/`"), + ("labels/x", "does not read or write yet"), + ("", "is empty"), + ], +) +def test_invalid_paths_explain_the_grammar(path: str, match: str) -> None: + with pytest.raises(ValueError, match=match): + align(_adata(), _adata(), in_=path, **_TINY) + + +def test_missing_obsm_key_lists_what_is_available() -> None: + with pytest.raises(KeyError, match="no `obsm\\['missing'\\]`.*spatial"): + align(_adata(), _adata(), in_="obsm/missing", **_TINY) + + +def test_missing_table_lists_what_is_available() -> None: + sdata = _sdata_tables(ref=_adata(), query=_adata()) + with pytest.raises(KeyError, match="no table 'nope'.*query.*ref"): + align(sdata, in_="tables/nope/obsm/spatial", **_TINY) + + +def test_table_path_on_anndata_is_rejected() -> None: + with pytest.raises(ValueError, match="but the container is an AnnData"): + align(_adata(), _adata(), in_="tables/t/obsm/spatial", **_TINY) + + +def test_bare_obsm_path_on_sdata_is_ambiguous() -> None: + sdata = _sdata_tables(ref=_adata(), query=_adata()) + with pytest.raises(ValueError, match="ambiguous for a SpatialData"): + align(sdata, in_="obsm/spatial", **_TINY) + + +def test_mixed_modalities_in_in_are_rejected() -> None: + sdata = _sdata_images() + with pytest.raises(ValueError, match="mixes modalities"): + align(sdata, in_=("images/ref", "tables/t/obsm/spatial"), by="images", **_TINY) + + +def test_out_modality_must_match_in() -> None: + sdata = _sdata_images() + with pytest.raises(ValueError, match="does not convert between the two"): + align(sdata, in_=("images/ref", "images/query"), by="images", out="tables/t/obsm/x", **_TINY_IMAGE) + + +def test_query_required_for_anndata() -> None: + with pytest.raises(ValueError, match="`data_query` is required"): + align(_adata(), in_="obsm/spatial", **_TINY) + + +def test_align_with_landmarks() -> None: + ref, query = _adata(), _adata() + landmarks = ref.obsm["spatial"][:3] + + result = align( + ref, + query, + in_="obsm/spatial", + method="stalign", + landmarks_source=landmarks, + landmarks_target=landmarks, + **_TINY, + ) + + assert isinstance(result, StalignResult) + assert result.aligned_points.shape == query.obsm["spatial"].shape diff --git a/tests/experimental/tl/test_align_io.py b/tests/experimental/tl/test_align_io.py new file mode 100644 index 000000000..e575aba5b --- /dev/null +++ b/tests/experimental/tl/test_align_io.py @@ -0,0 +1,103 @@ +"""Unit tests for the SpatialData transformation write-back. + +No estimator, no JAX -- these exercise registering a fitted affine on a coordinate +system directly, including the copy semantics, which are easy to get subtly wrong +because :func:`shallow_copy_sdata` shares element objects with the original. + +Array reads and writes at a path live in :mod:`._paths` and are covered by +``test_align_paths.py``. +""" + +from __future__ import annotations + +import numpy as np +import pytest +from anndata import AnnData + +from squidpy.experimental.methods.align_landmarks import AffineFitResult +from squidpy.experimental.tl._align._io import shallow_copy_sdata, writeback_affine_sdata + +_PTS = np.array([[10.0, 1.0], [12.0, 1.0], [11.0, 2.0], [10.0, 3.0], [12.0, 3.0]]) + + +def _adata(coords: np.ndarray = _PTS, *, key: str = "spatial") -> AnnData: + adata = AnnData(np.zeros((coords.shape[0], 1))) + adata.obsm[key] = coords.copy() + return adata + + +def _sdata_tables(**tables: AnnData): + sd = pytest.importorskip("spatialdata") + from spatialdata.models import TableModel + + return sd.SpatialData(tables={name: TableModel.parse(adata) for name, adata in tables.items()}) + + +def _sdata_points(cs: str = "qcs"): + sd = pytest.importorskip("spatialdata") + from spatialdata.models import PointsModel + from spatialdata.transformations import Identity + + pts = PointsModel.parse(_PTS, transformations={cs: Identity()}) + return sd.SpatialData(points={"pts": pts}) + + +# --------------------------------------------------------------------------- +# writeback_affine_sdata +# --------------------------------------------------------------------------- + + +def test_writeback_affine_inplace_registers_transform() -> None: + pytest.importorskip("spatialdata") + from spatialdata.transformations import get_transformation + + sdata = _sdata_points() + out = writeback_affine_sdata( + AffineFitResult(matrix=np.eye(3)), sdata, output_mode="inplace", moving_cs="qcs", target_cs="tcs" + ) + assert out is None + assert "tcs" in get_transformation(sdata.points["pts"], get_all=True) + + +def test_writeback_affine_copy_leaves_original_untouched() -> None: + pytest.importorskip("spatialdata") + from spatialdata.transformations import get_transformation + + sdata = _sdata_points() + out = writeback_affine_sdata( + AffineFitResult(matrix=np.eye(3)), sdata, output_mode="copy", moving_cs="qcs", target_cs="tcs" + ) + assert out is not sdata + assert "tcs" in get_transformation(out.points["pts"], get_all=True) + # regression: copy must not register the transform on the input container + assert "tcs" not in get_transformation(sdata.points["pts"], get_all=True) + + +def test_writeback_affine_requires_cs_names() -> None: + pytest.importorskip("spatialdata") + with pytest.raises(ValueError, match="`cs_query` and `cs_ref` are required"): + writeback_affine_sdata( + AffineFitResult(matrix=np.eye(3)), _sdata_points(), output_mode="inplace", moving_cs=None, target_cs="tcs" + ) + + +def test_writeback_affine_no_matching_cs() -> None: + pytest.importorskip("spatialdata") + sdata = _sdata_points(cs="qcs") + with pytest.raises(KeyError, match="No elements .* registered to coordinate system 'other'"): + writeback_affine_sdata( + AffineFitResult(matrix=np.eye(3)), sdata, output_mode="inplace", moving_cs="other", target_cs="tcs" + ) + + +# --------------------------------------------------------------------------- +# shallow_copy_sdata +# --------------------------------------------------------------------------- + + +def test_shallow_copy_sdata_preserves_elements() -> None: + pytest.importorskip("spatialdata") + sdata = _sdata_tables(ref=_adata(), query=_adata(_PTS + 5)) + copy = shallow_copy_sdata(sdata) + assert copy is not sdata + assert set(copy.tables) == {"ref", "query"} diff --git a/tests/experimental/tl/test_align_landmarks.py b/tests/experimental/tl/test_align_landmarks.py new file mode 100644 index 000000000..d8ef944b4 --- /dev/null +++ b/tests/experimental/tl/test_align_landmarks.py @@ -0,0 +1,300 @@ +"""Integration tests for the landmark path of the public ``align``. + +Closed-form and JAX-free, so these run everywhere. They cover both write-backs a landmark +fit can have: transforming coordinates into an ``obsm`` key, and registering the affine on +a whole SpatialData coordinate system. +""" + +from __future__ import annotations + +import numpy as np +import pytest +from anndata import AnnData + +from squidpy.experimental.methods.align_landmarks import AffineFitResult +from squidpy.experimental.tl import align + +# square corners; query = ref shifted by (5, 7) -> a pure translation both models recover +_REF = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]) +_SHIFT = np.array([5.0, 7.0]) +_QUERY = _REF + _SHIFT + + +def _adata(landmarks: np.ndarray, coords: np.ndarray, *, key: str = "spatial") -> AnnData: + adata = AnnData(np.zeros((coords.shape[0], 1))) + adata.obsm["landmarks"] = landmarks.copy() + adata.obsm[key] = coords.copy() + return adata + + +def _shapes(points: np.ndarray, cs: str = "global"): + pytest.importorskip("spatialdata") + import geopandas + import shapely + from spatialdata.models import ShapesModel + from spatialdata.transformations import Identity + + frame = geopandas.GeoDataFrame(geometry=[shapely.Point(*p) for p in points]) + frame["radius"] = 1.0 + return ShapesModel.parse(frame, transformations={cs: Identity()}) + + +# --- fitting --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("method", ["similarity", "affine"]) +def test_returns_an_affine_result(method: str) -> None: + ref, query = _adata(_REF, _REF), _adata(_QUERY, _QUERY) + result = align(ref, query, in_="obsm/landmarks", by="landmarks", method=method) + + assert isinstance(result, AffineFitResult) + assert result.matrix.shape == (3, 3) + np.testing.assert_allclose(result.transform(_QUERY), _REF, atol=1e-6) + assert result.metadata["method"] == method + + +def test_landmarks_default_to_similarity() -> None: + sdata = pytest.importorskip("spatialdata").SpatialData( + shapes={"lm_ref": _shapes(_REF), "lm_query": _shapes(_QUERY)} + ) + result = align(sdata, in_=("shapes/lm_ref", "shapes/lm_query"), by="landmarks") + assert result.metadata["method"] == "similarity" + + +def test_landmarks_need_not_live_in_a_shapes_element() -> None: + """``in_`` says where to read; ``by`` says what the arrays mean. + + Requiring a SpatialData just to hold four points would tax AnnData users for nothing. + """ + ref, query = _adata(_REF, _REF), _adata(_QUERY, _QUERY) + result = align(ref, query, in_="obsm/landmarks", by="landmarks", method="affine") + np.testing.assert_allclose(result.transform(_QUERY), _REF, atol=1e-6) + + +def test_stalign_rejects_a_landmark_path() -> None: + sdata = pytest.importorskip("spatialdata").SpatialData( + shapes={"lm_ref": _shapes(_REF), "lm_query": _shapes(_QUERY)} + ) + with pytest.raises(ValueError, match="does not support landmarks alignment.*obs, images"): + align(sdata, in_=("shapes/lm_ref", "shapes/lm_query"), by="landmarks", method="stalign") + + +# --- writing coordinates ---------------------------------------------------------------- + + +def test_apply_to_selects_what_moves() -> None: + ref, query = _adata(_REF, _REF), _adata(_QUERY, _QUERY) + out = align( + ref, query, in_="obsm/landmarks", by="landmarks", method="affine", apply_to="obsm/spatial", out="obsm/aligned" + ) + + assert out is None + np.testing.assert_allclose(query.obsm["aligned"], _REF, atol=1e-6) + np.testing.assert_allclose(query.obsm["spatial"], _QUERY) + + +def test_apply_to_is_required_for_landmarks() -> None: + """``in_`` holds correspondences, so it cannot also say which array to transform.""" + ref, query = _adata(_REF, _REF), _adata(_QUERY, _QUERY) + with pytest.raises(ValueError, match="needs `apply_to` when aligning by landmarks"): + align(ref, query, in_="obsm/landmarks", by="landmarks", method="affine", out="obsm/aligned") + + +def test_copy_leaves_original_untouched() -> None: + ref, query = _adata(_REF, _REF), _adata(_QUERY, _QUERY) + out = align( + ref, + query, + in_="obsm/landmarks", + by="landmarks", + method="affine", + apply_to="obsm/spatial", + out="obsm/aligned", + copy=True, + ) + assert isinstance(out, AnnData) and out is not query + assert "aligned" in out.obsm + assert "aligned" not in query.obsm + + +# --- registering a transformation -------------------------------------------------------- + + +def test_registers_a_transformation_on_the_coordinate_system() -> None: + sd = pytest.importorskip("spatialdata") + from spatialdata.models import PointsModel + from spatialdata.transformations import Identity, get_transformation + + sdata = sd.SpatialData( + shapes={"lm_ref": _shapes(_REF, "ref_cs"), "lm_query": _shapes(_QUERY, "query_cs")}, + points={"pts": PointsModel.parse(_QUERY, transformations={"query_cs": Identity()})}, + ) + out = align(sdata, in_=("shapes/lm_ref", "shapes/lm_query"), by="landmarks", method="affine", out="cs/ref_cs") + + assert out is None + assert "ref_cs" in get_transformation(sdata.points["pts"], get_all=True) + + +def test_registering_copy_leaves_original_untouched() -> None: + sd = pytest.importorskip("spatialdata") + from spatialdata.models import PointsModel + from spatialdata.transformations import Identity, get_transformation + + sdata = sd.SpatialData( + shapes={"lm_ref": _shapes(_REF, "ref_cs"), "lm_query": _shapes(_QUERY, "query_cs")}, + points={"pts": PointsModel.parse(_QUERY, transformations={"query_cs": Identity()})}, + ) + out = align( + sdata, in_=("shapes/lm_ref", "shapes/lm_query"), by="landmarks", method="affine", out="cs/ref_cs", copy=True + ) + + assert out is not sdata + assert "ref_cs" in get_transformation(out.points["pts"], get_all=True) + assert "ref_cs" not in get_transformation(sdata.points["pts"], get_all=True) + + +def test_registration_composes_with_an_existing_transform() -> None: + """The fit maps ``query_cs`` coords into the target, not the element's intrinsic frame. + + An element placed into ``query_cs`` by a non-identity transform must keep that + placement, so the registered transform has to compose the two. + """ + sd = pytest.importorskip("spatialdata") + from spatialdata.models import PointsModel + from spatialdata.transformations import Translation, get_transformation + + offset = np.array([100.0, 200.0]) + sdata = sd.SpatialData( + # Landmarks are expressed in `query_cs`: intrinsic coords shifted by `offset`. + shapes={"lm_ref": _shapes(_REF, "ref_cs"), "lm_query": _shapes(_QUERY + offset, "query_cs")}, + points={"pts": PointsModel.parse(_QUERY, transformations={"query_cs": Translation(offset, axes=("x", "y"))})}, + ) + align(sdata, in_=("shapes/lm_ref", "shapes/lm_query"), by="landmarks", method="affine", out="cs/ref_cs") + + matrix = get_transformation(sdata.points["pts"], to_coordinate_system="ref_cs").to_affine_matrix( + input_axes=("x", "y"), output_axes=("x", "y") + ) + mapped = _QUERY @ matrix[:2, :2].T + matrix[:2, 2] + np.testing.assert_allclose(mapped, _REF, atol=1e-6) + + +def test_shared_coordinate_system_is_refused() -> None: + """Registering moves everything in the coordinate system -- including the reference. + + With both samples in one coordinate system the write-back would drag the reference + along with the query and silently produce a wrong answer, so it has to refuse. + """ + sdata = pytest.importorskip("spatialdata").SpatialData( + shapes={"lm_ref": _shapes(_REF, "global"), "lm_query": _shapes(_QUERY, "global")} + ) + with pytest.raises(ValueError, match="both in coordinate system 'global'.*move the reference too"): + align(sdata, in_=("shapes/lm_ref", "shapes/lm_query"), by="landmarks", out="cs/aligned") + + +def test_reference_in_another_object_is_fine() -> None: + """Only a *shared* coordinate system is a problem; two objects cannot collide.""" + sd = pytest.importorskip("spatialdata") + from spatialdata.transformations import get_transformation + + ref_sdata = sd.SpatialData(shapes={"lm": _shapes(_REF, "global")}) + query_sdata = sd.SpatialData(shapes={"lm": _shapes(_QUERY, "global")}) + + align(ref_sdata, query_sdata, in_="shapes/lm", by="landmarks", out="cs/aligned") + + assert "aligned" in get_transformation(query_sdata["lm"], get_all=True) + assert "aligned" not in get_transformation(ref_sdata["lm"], get_all=True) + + +def test_ambiguous_coordinate_system_is_refused() -> None: + """Which system moves has to be unambiguous, so exactly one is required.""" + sd = pytest.importorskip("spatialdata") + import geopandas + import shapely + from spatialdata.models import ShapesModel + from spatialdata.transformations import Identity + + frame = geopandas.GeoDataFrame(geometry=[shapely.Point(*p) for p in _QUERY]) + frame["radius"] = 1.0 + both = ShapesModel.parse(frame, transformations={"a": Identity(), "b": Identity()}) + sdata = sd.SpatialData(shapes={"lm_ref": _shapes(_REF, "ref_cs"), "lm_query": both}) + + with pytest.raises(ValueError, match="registered to 2 coordinate systems"): + align(sdata, in_=("shapes/lm_ref", "shapes/lm_query"), by="landmarks", out="cs/aligned") + + +def test_table_landmarks_cannot_target_a_coordinate_system() -> None: + """A table annotates elements; it has no coordinate system of its own to move.""" + sd = pytest.importorskip("spatialdata") + from spatialdata.models import TableModel + + def table(points: np.ndarray) -> AnnData: + adata = AnnData(np.zeros((points.shape[0], 1))) + adata.obsm["landmarks"] = points + return TableModel.parse(adata) + + sdata = sd.SpatialData(tables={"r": table(_REF), "q": table(_QUERY)}) + with pytest.raises(ValueError, match="a table has no coordinate system of its own"): + align( + sdata, + in_=("tables/r/obsm/landmarks", "tables/q/obsm/landmarks"), + by="landmarks", + out="cs/aligned", + ) + + +def test_diffeomorphism_cannot_be_registered() -> None: + """SpatialData's transformations are affine at most, so stalign has to materialise.""" + sd = pytest.importorskip("spatialdata") + pytest.importorskip("jax") + from spatialdata.models import TableModel + + def table(points: np.ndarray) -> AnnData: + adata = AnnData(np.zeros((points.shape[0], 1))) + adata.obsm["spatial"] = points + return TableModel.parse(adata) + + pts = np.array([[10.0, 1.0], [12.0, 1.0], [11.0, 2.0], [10.0, 3.0], [12.0, 3.0]]) + sdata = sd.SpatialData(tables={"ref": table(pts), "query": table(pts)}) + + with pytest.raises(ValueError, match="no transformation type for"): + align( + sdata, + in_=("tables/ref/obsm/spatial", "tables/query/obsm/spatial"), + out="cs/aligned", + dx=0.5, + blur=1.0, + a=1.0, + expand=1.0, + nt=1, + niter=1, + epV=1.0, + ) + + +# --- input validation -------------------------------------------------------------------- + + +def test_too_few_landmarks() -> None: + ref, query = _adata(_REF[:2], _REF[:2]), _adata(_QUERY[:2], _QUERY[:2]) + with pytest.raises(ValueError, match="at least 3 landmark pairs"): + align(ref, query, in_="obsm/landmarks", by="landmarks", method="affine") + + +def test_length_mismatch() -> None: + ref, query = _adata(_REF, _REF), _adata(_QUERY[:3], _QUERY[:3]) + with pytest.raises(ValueError, match="same shape"): + align(ref, query, in_="obsm/landmarks", by="landmarks", method="affine") + + +def test_unknown_method_lists_available() -> None: + ref, query = _adata(_REF, _REF), _adata(_QUERY, _QUERY) + with pytest.raises(ValueError, match="Unknown align method 'nope'.*affine.*similarity.*stalign"): + align(ref, query, in_="obsm/landmarks", method="nope") + + +def test_non_finite_landmarks_rejected() -> None: + bad = _QUERY.copy() + bad[0, 0] = np.nan + ref, query = _adata(_REF, _REF), _adata(bad, bad) + with pytest.raises(ValueError, match="finite"): + align(ref, query, in_="obsm/landmarks", by="landmarks", method="affine") diff --git a/tests/experimental/tl/test_align_paths.py b/tests/experimental/tl/test_align_paths.py new file mode 100644 index 000000000..9e5c74824 --- /dev/null +++ b/tests/experimental/tl/test_align_paths.py @@ -0,0 +1,176 @@ +"""Unit tests for the ``in_`` / ``out`` path layer. + +No estimator, no JAX -- these exercise parsing, reading and writing directly, so the +AnnData *and* SpatialData branches and every error guard are covered cheaply. +""" + +from __future__ import annotations + +import numpy as np +import pytest +from anndata import AnnData + +from squidpy.experimental.tl._align._paths import parse_path, read_path, write_path + +_PTS = np.array([[10.0, 1.0], [12.0, 1.0], [11.0, 2.0], [10.0, 3.0], [12.0, 3.0]]) + + +def _adata(coords: np.ndarray = _PTS, *, key: str = "spatial") -> AnnData: + adata = AnnData(np.zeros((coords.shape[0], 1))) + adata.obsm[key] = coords.copy() + return adata + + +def _sdata_tables(**tables: AnnData): + sd = pytest.importorskip("spatialdata") + from spatialdata.models import TableModel + + return sd.SpatialData(tables={name: TableModel.parse(adata) for name, adata in tables.items()}) + + +def _sdata_image(name: str = "he", shape: tuple[int, int, int] = (2, 8, 6)): + sd = pytest.importorskip("spatialdata") + from spatialdata.models import Image2DModel + + data = np.arange(int(np.prod(shape)), dtype=float).reshape(shape) + return sd.SpatialData(images={name: Image2DModel.parse(data, dims=("c", "y", "x"))}) + + +# --- parsing --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("path", "modality", "element", "key"), + [ + ("obsm/spatial", "obs", None, "spatial"), + ("/obsm/spatial/", "obs", None, "spatial"), + ("tables/s1/obsm/xy", "obs", "s1", "xy"), + ("images/he", "images", "he", None), + ("shapes/lm", "landmarks", "lm", None), + ("cs/aligned", "landmarks", "aligned", None), + ], +) +def test_parse_accepts_the_documented_forms(path, modality, element, key) -> None: + parsed = parse_path(path, name="in_") + assert (parsed.modality, parsed.element, parsed.key) == (modality, element, key) + assert parsed.raw == path + + +@pytest.mark.parametrize( + ("path", "match"), + [ + ("", "is empty"), + ("///", "is empty"), + ("obsm", "Expected `obsm/`"), + ("obsm/a/b", "Expected `obsm/`"), + ("tables/s1/obs/x", "Expected `tables/
/obsm/`"), + ("tables/s1", "Expected `tables/
/obsm/`"), + ("images", "Expected `images/`"), + ("shapes/a/b", "Expected `shapes/`"), + ("cs/a/b", "Expected `cs/`"), + ("labels/x", "does not read or write yet"), + ("varm/x", "does not start with a known collection"), + ], +) +def test_parse_rejects_malformed_paths(path, match) -> None: + with pytest.raises(ValueError, match=match): + parse_path(path, name="in_") + + +def test_parse_rejects_non_strings() -> None: + with pytest.raises(TypeError, match="must be a string path"): + parse_path(("obsm/spatial",), name="in_") # type: ignore[arg-type] + + +# --- reading --------------------------------------------------------------------------- + + +def test_read_anndata_obsm() -> None: + got = read_path(_adata(), parse_path("obsm/spatial", name="in_"), name="in_") + np.testing.assert_array_equal(got, _PTS) + + +def test_read_sdata_table_obsm() -> None: + sdata = _sdata_tables(s1=_adata(), s2=_adata(_PTS + 5)) + got = read_path(sdata, parse_path("tables/s2/obsm/spatial", name="in_"), name="in_") + np.testing.assert_array_equal(got, _PTS + 5) + + +def test_read_image_is_channels_first() -> None: + got = read_path(_sdata_image(), parse_path("images/he", name="in_"), name="in_") + assert got.shape == (2, 8, 6) + + +def test_read_rejects_non_2d_coordinates() -> None: + adata = _adata() + adata.obsm["bad"] = np.zeros((5, 3)) + with pytest.raises(ValueError, match=r"must be an \(N, 2\) array"): + read_path(adata, parse_path("obsm/bad", name="in_"), name="in_") + + +def test_missing_obsm_key_lists_alternatives() -> None: + with pytest.raises(KeyError, match="no `obsm\\['nope'\\]`.*spatial"): + read_path(_adata(), parse_path("obsm/nope", name="in_"), name="in_") + + +def test_missing_table_lists_alternatives() -> None: + sdata = _sdata_tables(s1=_adata(), s2=_adata()) + with pytest.raises(KeyError, match="no table 'nope'.*s1.*s2"): + read_path(sdata, parse_path("tables/nope/obsm/spatial", name="in_"), name="in_") + + +def test_missing_image_lists_alternatives() -> None: + with pytest.raises(KeyError, match="no image 'nope'.*he"): + read_path(_sdata_image(), parse_path("images/nope", name="in_"), name="in_") + + +def test_table_path_against_anndata_is_rejected() -> None: + with pytest.raises(ValueError, match="but the container is an AnnData"): + read_path(_adata(), parse_path("tables/s1/obsm/spatial", name="in_"), name="in_") + + +def test_bare_obsm_path_against_sdata_is_ambiguous() -> None: + sdata = _sdata_tables(s1=_adata(), s2=_adata()) + with pytest.raises(ValueError, match="ambiguous for a SpatialData"): + read_path(sdata, parse_path("obsm/spatial", name="in_"), name="in_") + + +def test_image_path_against_anndata_is_rejected() -> None: + with pytest.raises(TypeError, match="only a SpatialData holds"): + read_path(_adata(), parse_path("images/he", name="in_"), name="in_") + + +def test_coordinate_system_is_write_only() -> None: + """A coordinate system holds transformations, not data to align.""" + with pytest.raises(ValueError, match="only valid as an `out`"): + read_path(_adata(), parse_path("cs/aligned", name="in_"), name="in_") + + +# --- writing --------------------------------------------------------------------------- + + +def test_write_anndata_obsm() -> None: + adata = _adata() + write_path(adata, parse_path("obsm/aligned", name="out"), _PTS + 1) + np.testing.assert_array_equal(adata.obsm["aligned"], _PTS + 1) + np.testing.assert_array_equal(adata.obsm["spatial"], _PTS) + + +def test_write_sdata_table_obsm() -> None: + sdata = _sdata_tables(s1=_adata(), s2=_adata()) + write_path(sdata, parse_path("tables/s2/obsm/aligned", name="out"), _PTS + 1) + assert "aligned" in sdata.tables["s2"].obsm + assert "aligned" not in sdata.tables["s1"].obsm + + +def test_write_image_creates_a_new_element() -> None: + sdata = _sdata_image() + write_path(sdata, parse_path("images/warped", name="out"), np.zeros((2, 8, 6))) + assert sorted(sdata.images) == ["he", "warped"] + assert np.asarray(sdata.images["warped"].data).shape == (2, 8, 6) + + +def test_write_promotes_a_2d_image() -> None: + sdata = _sdata_image() + write_path(sdata, parse_path("images/warped", name="out"), np.zeros((8, 6))) + assert np.asarray(sdata.images["warped"].data).shape == (1, 8, 6)