Quick Start

Installation

git clone https://github.com/chansigit/torchgw.git
cd torchgw
pip install -e .

Or from PyPI:

pip install torchgw

Requirements: numpy, scipy, scikit-learn, torch>=2.0, joblib. Source code: github.com/chansigit/torchgw.

Basic Usage

import numpy as np
from torchgw import sampled_gw, build_knn_graph

# Two point clouds (dimensions may differ)
X = np.random.randn(500, 3).astype(np.float32)
Y = np.random.randn(600, 5).astype(np.float32)

# Compute transport plan
T = sampled_gw(X, Y, epsilon=0.005, M=80, max_iter=200)

# T[i,j] is the coupling weight between X[i] and Y[j]
print(T.shape)  # (500, 600)

With Precomputed Graphs

Building the kNN graph once and reusing it avoids redundant computation:

g_src = build_knn_graph(X, k=10)
g_tgt = build_knn_graph(Y, k=10)

T = sampled_gw(X, Y, graph_source=g_src, graph_target=g_tgt,
               epsilon=0.005, M=80, max_iter=200)

Semi-Relaxed Mode

When source and target have different compositions (e.g., a cell type present in source but absent in target), balanced GW forces mass onto wrong matches. Semi-relaxed GW fixes the source marginal but lets the target marginal adapt:

# Balanced (default): T @ 1 = p,  T.T @ 1 = q  (both enforced)
T = sampled_gw(X, Y, epsilon=0.005)

# Semi-relaxed: T @ 1 = p (enforced),  T.T @ 1 ≈ q (soft KL penalty)
T = sampled_gw(X, Y, epsilon=0.005, semi_relaxed=True, rho=1.0)

# rho controls how strictly q is enforced:
#   rho → ∞  : recovers balanced GW
#   rho → 0  : target marginal is completely free

Convergence Logging

T, info = sampled_gw(X, Y, epsilon=0.005, max_iter=200, log=True)
print(info["n_iter"])    # actual iterations run
print(info["err_list"])  # per-iteration convergence errors

Differentiable Mode

For end-to-end training, keep the computation graph:

import torch
from torchgw import sampled_gw

C_feat = torch.cdist(encoder(X), encoder(Y))

# Default: exact gradients via implicit differentiation
T = sampled_gw(fgw_alpha=1.0, C_linear=C_feat, differentiable=True, epsilon=0.01)
loss = (C_feat.detach() * T).sum()
loss.backward()  # exact gradients flow to encoder parameters

# Alternative: unrolled autograd (higher memory, useful as fallback)
T = sampled_gw(fgw_alpha=1.0, C_linear=C_feat, epsilon=0.01,
               differentiable=True, grad_mode="unrolled")

The default grad_mode="implicit" solves the adjoint system at the Sinkhorn fixed point for exact gradients with O(NK) memory. See Algorithm for the mathematical derivation.

Note

Gradients are only as exact as the fixed point. Below epsilon=0.01 entropic Sinkhorn converges too slowly for the plan to reach one, and sampled_gw emits a warning when differentiable=True.

Note

differentiable=True requires fgw_alpha > 0 with a differentiable C_linear. Pure GW (fgw_alpha=0) uses precomputed graph distances that are not part of the computation graph; a warning is emitted in this case.

Joint Embedding

After computing a transport plan, embed both datasets into a shared space:

from torchgw import joint_embedding

emb = joint_embedding(
    anchor_name="tgt",
    data_by_name={"src": X, "tgt": Y},
    graphs_by_name={"src": g_src, "tgt": g_tgt},
    transport_plans={("src", "tgt"): T},
    out_dim=10,
)
print(emb["src"].shape)  # (500, 10)
print(emb["tgt"].shape)  # (600, 10)