Riemannian Metric Learning for Alignment of Spatial Multiomics (Manifold Gromov-Wasserstein or MGW)

This the repository for "Riemannian Metric Learning for Alignment of Spatial Multiomics." a technique which:

1. Performs Riemannian metric learning across spatial modalities (multimoics, transcriptomics, and so on) using the Riemannian pull-back metric.
2. Infers Riemannian (geodesic) distances.
3. Aligns Riemannian distances with Gromov-Wasserstein Optimal Transport.

In the section below, we detail the usage of MGW which complements the simple demo notebooks:

- [demo_mgw_y7.ipynb](demo_mgw_y7.ipynb)
- [mouse_align.ipynb](mouse_align.ipynb)
- [riemannian_mouse_tensor.ipynb](riemannian_mouse_tensor.ipynb)
- [riemannian_mouse_geodesics.ipynb](riemannian_mouse_geodesics.ipynb)

Contents

mgw/

• mgw.py — main solver: mgw_preprocess, mgw_align_core, mgw_align, transfer
• geometry.py — metric-tensor, geodesic distance, k-NN graph, APSP utilities
• models.py — neural field models (MLP)
• metric.py — evaluation metrics (e.g. migration, AMI, cosine similarity)
• plotting.py — visualization utilities
• tensor_vis.py — Riemannian deformation-grid visualization
• util.py — miscellaneous helpers, barycentric projection

validation/

• dopamine.py — validation utilities for dopamine experiments (AUROC, AUPRC)
• run_methods.py — code for running other methods (moscot Translation, SCOT, SCOTv2, PASTE2 FGW spatial, POT FGW spatial only)

demos/

• demo_mgw_y7.ipynb — demo notebook for running MGW on the Y_7 ccRCC slice (Hu '24)
• mouse_align.ipynb - demo notebook for aligning Spatiotemporal Transcriptomics with MGW on E9.5-10.5 mouse embryo timepoint pair (Chen '22)
• riemannian_mouse_geodesics.ipynb — demo visualization of the geodesics in the Riemannian pull-back metric of E9.5-10.5 mouse embryo (Chen '22)

experiments/

• Reproducible experimental notebooks on Stereo-Seq Mouse Embryo, Visium-Xenium alignment of colorectal cancer, MALDI-MSI metabolomics and Visium transcriptomics alignment of human striatum, AFADESI-MSI and Visium alignment of renal cancer.

Getting Started

1. Load the two multiomic datasets

Load two AnnData objects such as spatial transcriptomics (st) and spatial metabolomics (msi) after appropriate filtering.

import anndata as ad
st = ad.read_h5ad(ST_PATH)
msi = ad.read_h5ad(MSI_PATH)

2. Run MGW preprocessing

Call mgw.mgw_preprocess on two AnnDatas. This is unlikely to be an all encompassing across arbitrary modalities which each have their own preprocessing and QC protocols. In our demos, we offered a few simple examples of such steps for transcriptomics and metabolomics which precede alignment. If you want your raw st.X and msi.X to be run as-is (e.g. with some feature processing of your own) without further processing, set use_cca_feeler=False, use_pca_X/Z=False, and log1p_X/Z=False.

You can run PCA (will default to pre-computed PCA if already done) with PCA_comp components, and an additional CCA step for multimodal data. Set use_cca_feeler=True for this CCA step, which involves basic/coarse feeler alignment (spatial_only: bool = True to do a spatial-only feeler, feature_only = True to do a feature-only feeler, or if both False a basic spatial-feature feeler). This subsets feature dimensions which are correlated across modalities, and you can specify the number of final CCA dimensions with CCA_comp.

We do not assume common/joint features in multimodal data generally and do independent internal PCA steps. For unimodal (e.g. transcriptomics-transcriptomics) alignment we strongly recommend an external joint PCA across the two datasets: see, e.g. experiments/mgw_mouse_embryo.ipynb for an example of this pre-processing. Lastly, for unimodal alignment we recommend using the MGW cost matrix as the geometric term in an FGW as these formulations can use rich cross-data feature similarity in addition to intrinsic geometry.

import mgw.mgw as mgw
pre = mgw.mgw_preprocess(
    st, msi,
    PCA_comp=PCA_componet,
    CCA_comp=CCA_componet,
    use_cca_feeler=True, 
    use_pca_X=True,
    use_pca_Z=False, #False if the features from second modality are intensities which doesn't make sense to run pca on
    log1p_X=True,
    log1p_Z=False, #False if the features from second modality are not counts which doesn't make sense to run log1p on
    verbose=True
)

3. Run MGW

Next, we run mgw.mgw_align_core on the data pre to both infer the neural fields, learn metric tensors, and align the result with Gromov-Wasserstein.

We recommend trying larger DEFAULT_EPS (e.g. 1e0) if the setting below yields a weak alignment. The Jacobian may become singular in degenerate cases such as fully constant expression regions, so larger DEFAULT_EPS can significantly improve Jacobian conditioning and alignment quality.

PHI_ARC = (128,256,256,128)
KNN_K= 12
DEFAULT_GW_PARAMS = dict(verbose=True, inner_maxit=3000, outer_maxit=3000, inner_tol=1e-7,   outer_tol=1e-7,   epsilon=1e-4)
DEFAULT_LR = 1e-3
DEFAULT_EPS = 1e-2
DEFAULT_ITER = 20_000
EXP_PATH = "your_path"
EXP_TAG = "your_tag"

out = mgw.mgw_align_core(
        pre,
        widths=PHI_ARC,
        lr=DEFAULT_LR,
        niter=DEFAULT_ITER,
        knn_k=KNN_K,
        geodesic_eps=DEFAULT_EPS,
        cost_p=2,             # 1 or 2 — see note below
        save_dir=EXP_PATH, 
        tag=EXP_TAG, 
        verbose=True,
        plot_net=True,
        gw_params=DEFAULT_GW_PARAMS
    )

Here, the key parameters are

• PHI_ARC: Layers of the MLP
• KNN_K: Resolution of the K nearest neighbor graph used for Riemannian geodesics
• DEFAULT_EPS: Epsilon for stability of Jacobian. Smaller yields more faithful Riemannian geodesics (e.g. 1e-5 for mouse embryo), but in case of degenerate Jacobian larger is better.
• DEFAULT_GW_PARAMS: Default parameters for the optimal transport solver of ott jax
• DEFAULT_LR: Learning-rate for the network.
• DEFAULT_ITER: Number of training iterations for the network.
• save_dir: Where to save outputs
• tag: Tag for generated files.
• cost_p: Power applied to Riemannian geodesic distances before GW alignment — see note below.

Choosing cost_p. The GW objective minimised by MGW is

$$\min_{\pi} \sum_{i,i',j,j'} \bigl(C_M(i,i') - C_N(j,j')\bigr)^2 \pi_{ij},\pi_{i'j'}$$

where $C_M(i,i') \propto d_M(i,i')^p$ and $d_M$ is the Riemannian geodesic distance. The choice of $p$ determines which aspects of the geometry the alignment prioritises:

cost_p	Cost matrix	Behaviour
2 (default)	$C \propto d^2$	Emphasises large inter-point separations quadratically; equivalent to comparing squared geodesic distances and more sensitive to the global spread of the tissue.
1	$C \propto d$	Linear in distance; de-emphasises very large pairs and is more robust to outlier spots or disconnected tissue regions. Can be preferable when the two datasets differ substantially in spatial scale or cell density.

We recommend starting with cost_p=2 and trying cost_p=1 if the alignment appears dominated by a small number of outlier long-range pairs.

4. Transfer features across modalities

We have a number of variables which can be accessed from out.

• P: MGW coupling/alignment
• xs: Spatial coordinates 1 (normalized)
• xs2: Spatial coordinates 2 (normalized)
• phi: Neural field mapping into modality 1
• psi: Neural field mapping into modality 2
• G_M/G_N: Pull-back metric tensor field evaluated at the coordinates
• C_M/C_N: MGW Riemannian distance matrices

As an example, let us return the alignment and barycentrically project across modalities.

P = out["P"]
adata_sm2st = mgw.transfer(st, msi, P)
adata_st2sm = mgw.transfer(msi, st, P.T)

P represents the MGW alignment, adata_sm2st is the metabolomics to transcriptomics projection (added to st as metabolite annotation), and adata_st2sm is the transcriptomics to metabolomics projection (added to msi as transcriptomics annotation).

5. Visualize the Riemannian deformation grid

The deformation grid shows how the learned metric $g_{\mu\nu}(x) = J_\varphi^\top J_\varphi$ deforms the tissue. Grid lines are level sets of harmonic coordinate fields on the Riemannian kNN graph; the background heatmap encodes local anisotropy $\log(\lambda_{\max}/\lambda_{\min})$.

from mgw.tensor_vis import plot_deformation_grids
plot_deformation_grids(out, labels=['Dataset 1', 'Dataset 2'])