Quadratic problems (advanced)¶

This example shows an advanced quadratic problems usage, e.g., the LineageProblem, the SpatioTemporalProblem, the MappingProblem, the AlignmentProblem, the GWProblem, and the FGWProblem.

Imports and data loading¶

import warnings

warnings.simplefilter("ignore", FutureWarning)

from moscot import datasets
from moscot.problems.generic import GWProblem

import scanpy as sc

Simulate data using simulate_data().

adata = datasets.simulate_data(n_distributions=2, key="batch", quad_term="spatial")
sc.pp.pca(adata)
adata

AnnData object with n_obs × n_vars = 40 × 60
    obs: 'batch', 'celltype'
    uns: 'pca'
    obsm: 'spatial', 'X_pca'
    varm: 'PCs'

gwp = GWProblem(adata)
gwp = gwp.prepare(
    key="batch",
    x_attr={"attr": "obsm", "key": "spatial"},
    y_attr={"attr": "obsm", "key": "spatial"},
)
gwp

GWProblem[('0', '1')]

Threshold¶

The threshold parameter defines the convergence criterion. In the balanced setting the threshold denotes the deviation between prior and posterior marginals, while in the unbalanced setting the threshold corresponds to a Cauchy sequence stopping criterion.

Initializers¶

Different Initializers can help to improve convergence. For the full-rank case only the default initializer exists, hence the initializer argument must be set to None.

For low-rank problems the same initializers as for the linear low-rank solvers are available, and initializer_kwargs can be passed the same way, see Linear problems (advanced) for more information.

Number of iterations¶

To solve a quadratic optimal transport problem, a consecutively-updated linearized problem is solved. Here, min_iterations denotes a lower bound and max_iterations an upper bound on the number of outer iterations. If max_iterations is too low, the solver might not converge.

gwp = gwp.solve(epsilon=1e-1, min_iterations=0, max_iterations=1)

INFO     Solving `1` problems                                                                                      
INFO     Solving problem OTProblem[stage='prepared', shape=(20, 20)].                                              

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

WARNING  Solver did not converge

Linear solver keyword arguments¶

As mentioned above, each outer loop step of the Gromov-Wasserstein algorithm consists of solving a linear problem. Arguments for the linear solver can be specified via linear_solver_kwargs, keyword arguments for Sinkhorn in the full-rank case or keyword arguments for LRSinkhorn, respectively. This way, we can also set the minimum and maximum number of iterations for the linear solver:

ls_kwargs = {"min_iterations": 10, "max_iterations": 1000, "threshold": 0.01}
gwp = gwp.solve(
    epsilon=1e-1,
    threshold=0.1,
    min_iterations=2,
    max_iterations=20,
    linear_solver_kwargs=ls_kwargs,
)

INFO     Solving `1` problems                                                                                      
INFO     Solving problem OTProblem[stage='solved', shape=(20, 20)].                                                

Low-rank hyperparameters¶

The parameters gamma and gamma_rescale are the same as in the linear case, see example Linear problems (advanced).

Keyword arguments and implementation details¶

Whenever the solve method of a quadratic problem is called, a backend-specific quadratic solver is instantiated. Currently, ott is supported, its corresponding quadratic solvers is GromovWasserstein, handling both the full-rank and the low-rank case. moscot wraps this class in GWSolver, handling both the purely quadratic and the fused quadratic problem.