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.

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.

To run on the raw st.X and msi.X as-is without processing, set use_cca_feeler=False, use_pca_X/Z=False, and log1p_X/Z=False. We do not assume common/joint features in multimodal data generally and do independent internal PCA steps. For unimodal (e.g. transcriptomics-transcriptomics) we recommend an external joint PCA: see, e.g. experiments/mgw_mouse_embryo.ipynb for an example of this pre-processing.

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.

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. Generally not an issue, and smaller yields more faithful Riemannian geodesics (e.g. 1e-5 for mouse embryo).
• 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'])