"""Public solvers: input handling, multiscale warm start, and the two entry points.
Layering: ``_solver`` (this file) → ``_loop`` (GW iteration) → ``_sinkhorn`` /
``_lowrank`` (inner OT step), with ``_distances`` / ``_graph`` / ``_sampling`` /
``_multiscale`` as leaf helpers.
"""
import warnings
import numpy as np
import torch
from torchgw._distances import DijkstraProvider, LandmarkProvider, PrecomputedProvider
from torchgw._graph import build_knn_graph
from torchgw._loop import gw_loop, make_augmented_step, make_lowrank_step
from torchgw._multiscale import fps_downsample, upsample_plan
from torchgw._sinkhorn import GRAD_MODES
from torchgw._utils import get_device
DISTANCE_MODES = ("precomputed", "dijkstra", "landmark")
# ── Input handling ──────────────────────────────────────────────────────
def _to_tensor(x):
"""numpy array or tensor → tensor; None passes through."""
if x is None or isinstance(x, torch.Tensor):
return x
return torch.as_tensor(x)
def _resolve_rho(rho, rho_a, rho_b):
"""Map the legacy single-sided ``rho`` (target side only) onto (rho_a, rho_b)."""
if rho is None:
return rho_a, rho_b
if rho_a != 1.0 or rho_b != 1.0:
raise ValueError("pass either rho (legacy) OR rho_a/rho_b, not both")
return 1.0, rho
def _prepare_inputs(
X_source, X_target, p, q, dist_source, dist_target, C_linear,
distance_mode, fgw_alpha, k, n_landmarks, device,
):
"""Coerce inputs, infer N/K, validate, build the distance provider.
Returns (X_source, X_target, p, q, dist_source, dist_target, C_linear,
N, K, provider, device). provider is None for pure Wasserstein.
"""
X_source, X_target = _to_tensor(X_source), _to_tensor(X_target)
p, q = _to_tensor(p), _to_tensor(q)
dist_source, dist_target = _to_tensor(dist_source), _to_tensor(dist_target)
C_linear = _to_tensor(C_linear)
if dist_source is not None and dist_target is not None:
N, K = dist_source.shape[0], dist_target.shape[0]
elif X_source is not None and X_target is not None:
N, K = X_source.shape[0], X_target.shape[0]
elif C_linear is not None:
N, K = C_linear.shape
else:
raise ValueError(
"Cannot infer dataset sizes. Provide (X_source, X_target), "
"(dist_source, dist_target), or C_linear."
)
if distance_mode not in DISTANCE_MODES:
raise ValueError(f"distance_mode must be one of {DISTANCE_MODES}, got {distance_mode!r}")
if fgw_alpha > 0 and C_linear is None:
raise ValueError("fgw_alpha > 0 requires C_linear to be provided")
if device is None:
device = get_device()
have_X = X_source is not None and X_target is not None
if fgw_alpha >= 1.0:
provider = None
elif distance_mode == "precomputed":
if dist_source is not None and dist_target is not None:
provider = PrecomputedProvider(dist_source=dist_source, dist_target=dist_target)
elif have_X:
provider = PrecomputedProvider(
graph_source=build_knn_graph(X_source.cpu().numpy(), k=k),
graph_target=build_knn_graph(X_target.cpu().numpy(), k=k),
)
else:
raise ValueError(
"distance_mode='precomputed' requires (dist_source, dist_target) "
"or (X_source, X_target)"
)
else:
if not have_X:
raise ValueError(f"distance_mode={distance_mode!r} requires X_source and X_target")
g_src = build_knn_graph(X_source.cpu().numpy(), k=k)
g_tgt = build_knn_graph(X_target.cpu().numpy(), k=k)
if distance_mode == "dijkstra":
provider = DijkstraProvider(g_src, g_tgt)
else:
provider = LandmarkProvider(g_src, g_tgt, n_landmarks=n_landmarks)
return X_source, X_target, p, q, dist_source, dist_target, C_linear, N, K, provider, device
def _marginals(p, q, N, K, device):
"""Float64 marginals on device; uniform when not given."""
p = p.to(dtype=torch.float64, device=device) if p is not None \
else torch.ones(N, device=device, dtype=torch.float64) / N
q = q.to(dtype=torch.float64, device=device) if q is not None \
else torch.ones(K, device=device, dtype=torch.float64) / K
return p, q
def _maybe_multiscale(
multiscale, n_coarse, X_source, X_target, N, K,
dist_source, dist_target, C_linear, fgw_alpha,
distance_mode, n_landmarks, device, p, q,
solver_fn, solver_kwargs,
):
"""Solve a farthest-point-sampled coarse problem and upsample its plan as T_init.
Returns None when multiscale is off or the problem is too small.
"""
if not multiscale or X_source is None or X_target is None:
return None
n_c = n_coarse if n_coarse is not None else min(500, N // 4, K // 4)
n_c = max(n_c, 10)
if n_c >= N or n_c >= K:
return None
idx_src, assign_src = fps_downsample(X_source, n_c)
idx_tgt, assign_tgt = fps_downsample(X_target, n_c)
C_lin_coarse = C_linear[idx_src][:, idx_tgt] if C_linear is not None and fgw_alpha > 0 else None
dist_src_coarse = dist_tgt_coarse = None
if dist_source is not None and dist_target is not None:
dist_src_coarse = dist_source[idx_src][:, idx_src]
dist_tgt_coarse = dist_target[idx_tgt][:, idx_tgt]
kw = {**solver_kwargs, "multiscale": False} # no recursion
kw["M"] = min(solver_kwargs.get("M", 50), max(n_c // 2, 10))
kw["k"] = min(solver_kwargs.get("k", 30), n_c - 1)
if solver_kwargs.get("s_shared") is not None:
kw["s_shared"] = min(solver_kwargs["s_shared"], n_c)
T_coarse = solver_fn(
X_source[idx_src], X_target[idx_tgt],
distance_mode=distance_mode,
dist_source=dist_src_coarse, dist_target=dist_tgt_coarse,
n_landmarks=n_landmarks,
fgw_alpha=fgw_alpha, C_linear=C_lin_coarse,
device=device,
**kw,
)
return upsample_plan(T_coarse, assign_src, assign_tgt, p, q)
# ── Public API: standard solver ─────────────────────────────────────────
[docs]
def sampled_gw(
X_source: np.ndarray | torch.Tensor | None = None,
X_target: np.ndarray | torch.Tensor | None = None,
p: np.ndarray | torch.Tensor | None = None,
q: np.ndarray | torch.Tensor | None = None,
*,
# distances
distance_mode: str = "dijkstra",
dist_source: np.ndarray | torch.Tensor | None = None,
dist_target: np.ndarray | torch.Tensor | None = None,
k: int = 30,
n_landmarks: int = 50,
# fused GW
fgw_alpha: float = 0.0,
C_linear: np.ndarray | torch.Tensor | None = None,
# sampled GW algorithm
M: int = 50,
alpha: float = 0.9,
epsilon: float = 0.001,
s_shared: int | None = None,
lambda_ema_beta: float | None = None,
# solver control
max_iter: int = 500,
tol: float = 1e-5,
min_iter_before_converge: int = 50,
mixed_precision: bool = False,
T_init: np.ndarray | torch.Tensor | None = None,
device: torch.device | None = None,
verbose: bool = False,
verbose_every: int = 20,
log: bool = False,
# unbalanced transport
semi_relaxed: bool = False,
rho_a: float = 1.0,
rho_b: float = 1.0,
rho: float | None = None,
# multiscale warm start
multiscale: bool = False,
n_coarse: int | None = None,
# differentiable mode
differentiable: bool = False,
grad_mode: str = "implicit",
) -> torch.Tensor | tuple[torch.Tensor, dict]:
"""Sampled Gromov-Wasserstein alignment between two datasets.
Standard log-domain Sinkhorn with slack augmentation for partial transport.
Parameters
----------
X_source, X_target : ndarray or Tensor, optional
Feature matrices (N, D) and (K, D'). Required unless distance
matrices or ``C_linear`` are given.
p, q : ndarray or Tensor, optional
Marginal distributions (uniform if None).
**Distances**
distance_mode : str
``"dijkstra"`` (default), ``"precomputed"``, or ``"landmark"``.
dist_source, dist_target : ndarray or Tensor, optional
Precomputed (N, N) and (K, K) distance matrices.
k : int
k-NN graph connectivity.
n_landmarks : int
Landmark count for ``distance_mode="landmark"``.
**Fused GW**
fgw_alpha : float
0 = pure GW, 1 = pure Wasserstein, in between = fused.
C_linear : ndarray or Tensor, optional
(N, K) feature cost matrix. Required when ``fgw_alpha > 0``.
**Sampled GW algorithm**
M : int
Anchor pairs sampled per iteration.
alpha : float
Momentum: ``T ← (1-alpha) T + alpha T_new``.
epsilon : float
Entropic regularization (decays 10x over the run).
s_shared : int, optional
Expected number of matched points (partial transport).
lambda_ema_beta : float, optional
EMA smoothing of the sampled cost matrix
(``Lambda ← (1-beta) Lambda + beta Lambda_sample``). Reduces sampling
variance at the cost of a small bias that vanishes at convergence.
Typical values 0.3–0.7. None disables (default).
**Solver control**
max_iter, tol : int, float
Iteration cap and plan-change tolerance.
min_iter_before_converge : int
Iterations before convergence checks start.
mixed_precision : bool
Run Sinkhorn in float32 (log domain, numerically safe), output float64.
T_init : ndarray or Tensor, optional
Initial transport plan (overrides multiscale).
device : torch.device, optional
verbose, verbose_every : bool, int
log : bool
Also return a dict with ``err_list``, ``n_iter``, ``gw_cost``.
**Unbalanced transport**
semi_relaxed : bool
Relax marginals via KL penalties.
rho_a, rho_b : float
KL penalty weights for the source and target marginals; smaller
relaxes more. ``rho`` is the legacy alias for ``rho_b``.
**Multiscale warm start**
multiscale : bool
Solve a coarse problem first and upsample its plan as ``T_init``.
n_coarse : int, optional
Coarse problem size (auto if None).
**Differentiable mode**
differentiable : bool
Keep the computation graph so gradients flow to ``C_linear``.
Requires ``fgw_alpha > 0``; use ``epsilon >= 0.01``.
grad_mode : str
``"implicit"`` (default): exact gradient via the adjoint system at the
Sinkhorn fixed point, memory O(NK + K^2). ``"unrolled"``: backprop
through the Sinkhorn iterations, memory O(NK · iterations).
Returns
-------
T : Tensor (N, K)
log_dict : dict (only if ``log=True``)
"""
rho_a, rho_b = _resolve_rho(rho, rho_a, rho_b)
if grad_mode not in GRAD_MODES:
raise ValueError(f"grad_mode must be one of {GRAD_MODES}, got {grad_mode!r}")
(X_source, X_target, p, q, dist_source, dist_target, C_linear,
N, K, provider, device) = _prepare_inputs(
X_source, X_target, p, q, dist_source, dist_target, C_linear,
distance_mode, fgw_alpha, k, n_landmarks, device,
)
p_real, q_real = _marginals(p, q, N, K, device)
# Warm start: user T_init > multiscale > uniform
if T_init is not None:
T_init = _to_tensor(T_init).to(device=device, dtype=torch.float64)
else:
T_init = _maybe_multiscale(
multiscale, n_coarse, X_source, X_target, N, K,
dist_source, dist_target, C_linear, fgw_alpha,
distance_mode, n_landmarks, device, p_real, q_real,
solver_fn=sampled_gw,
solver_kwargs=dict(
s_shared=s_shared, M=M, alpha=alpha, max_iter=max_iter, tol=tol,
epsilon=epsilon, k=k, min_iter_before_converge=min_iter_before_converge,
verbose=False, log=False, differentiable=False,
semi_relaxed=semi_relaxed, rho_a=rho_a, rho_b=rho_b,
),
)
if T_init is None:
T_init = torch.outer(p_real, q_real)
C_lin = C_linear.to(dtype=torch.float64, device=device) if fgw_alpha > 0 else None
if differentiable and fgw_alpha == 0.0:
warnings.warn(
"differentiable=True with fgw_alpha=0 (pure GW): gradients cannot "
"flow because the GW cost matrix is built from precomputed graph "
"distances that are not part of the computation graph. Set "
"fgw_alpha > 0 with a differentiable C_linear to get useful gradients.",
stacklevel=2,
)
if differentiable and epsilon < 0.01:
warnings.warn(
f"differentiable=True with epsilon={epsilon}: entropic Sinkhorn converges "
"too slowly at this regularization for the transport plan to reach a fixed "
"point, so gradients (implicit or unrolled) are unreliable. Use epsilon >= 0.01 "
"for differentiable training.",
stacklevel=2,
)
sink_dtype = torch.float32 if mixed_precision else torch.float64
step = make_augmented_step(
p_real, q_real, s_shared,
semi_relaxed=semi_relaxed, rho_a=rho_a, rho_b=rho_b,
differentiable=differentiable, grad_mode=grad_mode,
sink_dtype=sink_dtype, device=device,
)
with (torch.enable_grad() if differentiable else torch.no_grad()):
T_out, err_list, n_iter, gw_cost = gw_loop(
provider=provider, T_init=T_init, step=step,
fgw_alpha=fgw_alpha, C_lin=C_lin,
M=M, alpha=alpha, epsilon=epsilon, lambda_ema_beta=lambda_ema_beta,
max_iter=max_iter, tol=tol, min_iter_before_converge=min_iter_before_converge,
differentiable=differentiable, sink_dtype=sink_dtype,
verbose=verbose, verbose_every=verbose_every,
)
if log:
return T_out, {"err_list": err_list, "n_iter": n_iter, "gw_cost": gw_cost}
return T_out
# ── Public API: low-rank solver ─────────────────────────────────────────
[docs]
def sampled_lowrank_gw(
X_source: np.ndarray | torch.Tensor | None = None,
X_target: np.ndarray | torch.Tensor | None = None,
p: np.ndarray | torch.Tensor | None = None,
q: np.ndarray | torch.Tensor | None = None,
*,
# low-rank Sinkhorn
rank: int = 20,
lr_max_iter: int = 5,
lr_dykstra_max_iter: int = 50,
# distances
distance_mode: str = "dijkstra",
dist_source: np.ndarray | torch.Tensor | None = None,
dist_target: np.ndarray | torch.Tensor | None = None,
k: int = 30,
n_landmarks: int = 50,
# fused GW
fgw_alpha: float = 0.0,
C_linear: np.ndarray | torch.Tensor | None = None,
# sampled GW algorithm
M: int = 50,
alpha: float = 0.9,
epsilon: float = 0.001,
s_shared: int | None = None,
lambda_ema_beta: float | None = None,
# solver control
max_iter: int = 500,
tol: float = 1e-5,
min_iter_before_converge: int = 50,
mixed_precision: bool = False,
device: torch.device | None = None,
verbose: bool = False,
verbose_every: int = 20,
log: bool = False,
# unbalanced transport (not supported yet; kept for API symmetry)
semi_relaxed: bool = False,
rho_a: float = 1.0,
rho_b: float = 1.0,
rho: float | None = None,
# multiscale warm start
multiscale: bool = False,
n_coarse: int | None = None,
) -> torch.Tensor | tuple[torch.Tensor, dict]:
"""Sampled Gromov-Wasserstein with low-rank Sinkhorn.
Uses the low-rank factorization of Scetbon, Cuturi & Peyré (2021) to cut
memory from O(NK) to O((N+K)·rank) per Sinkhorn step. This is a memory
optimization for N, K > 50k; at smaller scales ``sampled_gw`` is faster.
Parameters
----------
rank : int
Nonnegative rank of the transport plan factorization.
lr_max_iter : int
Mirror-descent iterations per Sinkhorn call.
lr_dykstra_max_iter : int
Dykstra projection iterations per Sinkhorn call.
All other parameters are as in :func:`sampled_gw`. ``semi_relaxed`` and
``differentiable`` are not supported by the low-rank solver.
Returns
-------
T : Tensor (N, K)
log_dict : dict (only if ``log=True``)
"""
rho_a, rho_b = _resolve_rho(rho, rho_a, rho_b)
if semi_relaxed:
raise ValueError("semi_relaxed is not supported for low-rank Sinkhorn")
if rho_a != rho_b:
raise NotImplementedError(
"sampled_lowrank_gw does not yet support rho_a != rho_b "
"(low-rank Dykstra requires symmetric KL). Use sampled_gw for "
"fully-unbalanced (rho_a, rho_b) FGW."
)
(X_source, X_target, p, q, dist_source, dist_target, C_linear,
N, K, provider, device) = _prepare_inputs(
X_source, X_target, p, q, dist_source, dist_target, C_linear,
distance_mode, fgw_alpha, k, n_landmarks, device,
)
p_real, q_real = _marginals(p, q, N, K, device)
T_init = _maybe_multiscale(
multiscale, n_coarse, X_source, X_target, N, K,
dist_source, dist_target, C_linear, fgw_alpha,
distance_mode, n_landmarks, device, p_real, q_real,
solver_fn=sampled_lowrank_gw,
solver_kwargs=dict(
rank=rank, lr_max_iter=lr_max_iter, lr_dykstra_max_iter=lr_dykstra_max_iter,
s_shared=s_shared, M=M, alpha=alpha, max_iter=max_iter, tol=tol,
epsilon=epsilon, k=k, min_iter_before_converge=min_iter_before_converge,
verbose=False, log=False,
semi_relaxed=semi_relaxed, rho_a=rho_a, rho_b=rho_b,
),
)
if T_init is None:
T_init = torch.outer(p_real, q_real)
C_lin = C_linear.to(dtype=torch.float64, device=device) if fgw_alpha > 0 else None
sink_dtype = torch.float32 if mixed_precision else torch.float64
step = make_lowrank_step(
p_real, q_real, rank=rank, lr_max_iter=lr_max_iter,
lr_dykstra_max_iter=lr_dykstra_max_iter, sink_dtype=sink_dtype,
)
with torch.no_grad():
T_out, err_list, n_iter, gw_cost = gw_loop(
provider=provider, T_init=T_init, step=step,
fgw_alpha=fgw_alpha, C_lin=C_lin,
M=M, alpha=alpha, epsilon=epsilon, lambda_ema_beta=lambda_ema_beta,
max_iter=max_iter, tol=tol, min_iter_before_converge=min_iter_before_converge,
differentiable=False, sink_dtype=sink_dtype,
verbose=verbose, verbose_every=verbose_every,
)
if log:
return T_out, {"err_list": err_list, "n_iter": n_iter, "gw_cost": gw_cost}
return T_out