API reference
Transport
geodesic
geodesic(G, rhoA, rhoB, *, method=DEFAULT_METHOD, **kwargs)
The discrete transport geodesic between densities rhoA and rhoB on G.
method="shooting" (default): Newton shooting on the Hamiltonian flow
(shooting.geodesic_shooting), then the flow integrated for the path.
Fourth-order in time (RK4 truncation); both endpoints must be strictly
positive. If one is not, or shooting fails on them (densities close to
zero, or a transport that needs more than maxiters Newton steps), this warns
(ShootingFallbackWarning) and returns method="socp"'s answer at its
default N=10 instead -- or raises, with fallback=False or
without cvxpy installed. Keywords: fallback, nsteps (integrator
steps, default 150; rho then has nsteps + 1 columns), tol,
maxiters, phi0_init, segments (default "auto": single
shooting, switching to multiple shooting on a long transport; an integer
fixes it -- see shooting.log_map), verbose. Honors every
AdmissibleMean, including the exact LogarithmicMean. phi0 and phi1 use
the SOCP's W2-gradient convention, so the two methods' potentials are
directly comparable; status is "converged".
method="socp": a single second-order-cone program
(socp.geodesic_socp). Handles any densities, including boundary-supported
ones; time-discretization error O(1/N). Keywords: N, solver,
check, verbose, plus solver options. Honors every conic
AdmissibleMean in G.mean (QuadLogMean for the logarithmic mean).
method="sinkhorn": the entropic displacement interpolation for a ground
cost -- the path is the entropic barycenter of the two endpoints at
weights (1-t, t) for t = 0, 1/N, ..., 1, a different object from the
discrete transport geodesic of the other methods. Requires cost (see
ground_cost) and epsilon; keywords N (default 10) and iters
(Sinkhorn budget, default 256). W2 is the entropic transport cost
tol
(l1, default 1e-6), otherwise "iteration_limit". The
path's end columns are the blurred endpoints the Sinkhorn barycenter
returns, not rhoA/rhoB exactly; m, phi0, phi1 are NaN-filled.
Torch tensors: rhoA and rhoB may be torch tensors (computed in float64 on the CPU), and the solution's fields are then tensors. With method="shooting" all of them -- W2, the path, the momenta, the potentials -- are differentiable with respect to both endpoints, exactly for the discrete solve (shooting.differentiable). The other methods are not differentiable: they raise if an input requires grad, and otherwise return tensors with no graph. Shooting's fallback to the SOCP is disabled when an input requires grad; shooting raises instead. Gradients are first order only: differentiating a gradient again (create_graph=True, for a gradient penalty or a Hessian-vector product) raises.
transport_cost
transport_cost(G, rhoA, rhoB, *, method=DEFAULT_METHOD, **kwargs)
The discrete transport distance W(rhoA, rhoB) (not squared): sqrt(geodesic(...).W2). See geodesic for the methods and keywords.
method="shooting" (default) needs only the log map, not the path, and
takes geodesic's keywords, segments included.
method="sinkhorn" solves only the endpoint plan rather than the whole
path, and warns if that plan has not converged (there is no status to
return).
Takes torch tensors as geodesic does, returning a 0-d tensor. W is not differentiable where W = 0 (rhoA == rhoB); its gradient there is 0, the subgradient at W's minimum and torch.linalg.norm's convention, rather than the nan the square root's infinite derivative would give. Gradients are first order only (see geodesic).
barycenter
barycenter(G, refs, lam, *, method=DEFAULT_METHOD, **kwargs)
The discrete transport barycenter of the reference densities refs
with weights lam: the minimizer of J(nu) = sum_i lam_i W^2(refs_i, nu).
method="shooting" (default): intrinsic gradient descent with shooting
geodesics (shooting.barycenter_shooting). Every reference with
lam_i > 0 must be strictly positive. info = {"iters", "status",
"J_hist", "grad_hist", "h"}; status is "converged", "stalled" (at the
precision of the log maps) or "maxiters" (which warns). Keywords: h,
maxiters (descent iterations), tol (gradient norm, default
1e-5), ftol, log_tol and log_maxiters (each log map's Newton
tolerance and budget, default 50), segments (multiple shooting for
each log map, see shooting.log_map), nsteps, init, verbose.
First order, so it
converges linearly; method="socp" solves the same problem to its global
optimum of its discretization and is the reference.
method="socp": one joint second-order-cone program
(socp.barycenter_socp), solved to its global optimum.
info = {"geodesics": [...]} holds one GeodesicSolution per reference
with lam_i > 0, each carrying the index of its reference as
ref_index: references at weight zero are not solved, so the list can
be shorter than refs and geodesics[k] need not be the geodesic of
refs[k]. Keywords: N, solver, check, verbose.
method="sinkhorn": the entropically regularized Wasserstein barycenter
for a ground cost (Benamou et al. 2015; Bonneel, Peyré & Cuturi 2016).
This is a different object from the discrete transport barycenter of
the other methods: it depends on cost (see ground_cost) and
epsilon, both required, and even for a single reference returns a
blurred copy of it. iters (default 256) is the Sinkhorn budget.
J = sum_i lam_i
Returns (nu, J, info).
analysis
analysis(G, target, refs, *, method=DEFAULT_METHOD, **kwargs)
Recover the barycentric coordinates of target with respect to the
reference densities refs.
method="shooting" (default): shooting.analyze_shooting -- potentials from
log_map at target, then the same Gram matrix and simplex QP as the
SOCP. target and every reference must be strictly positive.
Keywords: nsteps, tol and maxiters (each log map's Newton
tolerance and budget, default 50), segments (multiple shooting for
each log map, see shooting.log_map), phi0_inits, compute_condition,
return_system, qp_method, qp_solver.
A barycenter is recovered to solver tolerance only by the method that synthesized it; the others recover it to their discretization error, since each checks stationarity in its own discrete convention.
method="socp": socp.analyze_socp -- geodesic SOCPs from the
target to each reference, the Gram matrix of their endpoint potentials
(or initial momenta, convention="momentum"), and the simplex QP.
Keywords: N, solver, convention, compute_condition,
return_system (also return the Gram matrix A).
method="sinkhorn": Wasserstein barycentric coordinates for a ground cost
(Bonneel, Peyré & Cuturi 2016; simplex_regression): L-BFGS over the
simplex on 1/2 ||P(lam) - target||^2, differentiated through the
Sinkhorn iterations. Not a Gram-matrix method, so return_system and
compute_condition raise. Requires cost and epsilon; keywords
iters (default 256) and alpha0 (initial pre-softmax point,
default 0 = uniform weights). Recovers exactly the weights of a
barycenter synthesized by barycenter(method="sinkhorn") with the same
cost, epsilon and iters.
Returns lam_hat on the simplex.
GeodesicSolution
dataclass
A discrete transport geodesic.
W2: the squared transport distance. rho: the (n, steps+1) density path.
m: the (|E|, steps) momentum path and m0 its first column. phi0, phi1:
endpoint potentials (gradients of W2 with respect to each endpoint).
status: solver status; for method="sinkhorn", "converged" if the endpoint
plan's marginals are met to tol, else "iteration_limit".
solvetime: wall-clock seconds.
ref_index: for a geodesic that is part of a barycenter, the index of its
reference in the refs passed to barycenter; None for a standalone
geodesic. References with lam_i == 0 are not solved, so the list of
geodesics can be shorter than refs and this is what ties each one back.
Methods that do not produce a quantity fill it with NaN (the Sinkhorn method has no momenta or potentials). For numpy inputs the fields are a float and numpy arrays; for torch inputs (geodesic on tensors) they are tensors.
ShootingFallbackWarning
Bases: UserWarning
method="shooting" could not take the data, and the SOCP was used instead.
Raised by the default method when a density is not strictly positive, or when shooting fails on the data -- because a density is close to zero, or because Newton ran out of iterations (maxiters) on a hard transport. Filter it to silence the fallback, or escalate it with warnings.simplefilter("error", ShootingFallbackWarning); pass fallback=False to get the shooting error itself instead.
Graphs
MarkovGraph
Graph primitive: an undirected graph plus a reversible Markov chain on it.
Attributes:
| Name | Type | Description |
|---|---|---|
n |
number of vertices. |
|
E |
(|E|, 2) int array of oriented edges; E[e] = (x, y) means the edge is oriented from x to y (the orientation is the one with x < y). |
|
pi |
(n,) stationary distribution. |
|
Q |
(n, n) sparse transition rate matrix. |
|
kappa |
(|E|,) edge weights kappa[e] = Q[x, y] * pi[x], which by reversibility equals Q[y, x] * pi[y]. |
|
D |
(n, |E|) sparse incidence matrix with D @ m == graph_divergence(G, m) for any edge field m. D[x, e] = -Q[x, y], D[y, e] = Q[y, x] for edge e = (x, y). |
|
mean |
AdmissibleMean
|
the mobility theta(s, t) of the transport metric on this graph (GeometricMean() by default). The mean is part of the geometry, like the graph itself, so it lives here rather than as a per-call option: every geodesic, barycenter and analysis computed from the same MarkovGraph uses the same metric, and a target synthesized on one graph is analyzed with the mean it was made with. |
with_mean
with_mean(mean)
The same graph with a different mean; shares the cached matrices.
grid_markov_chain
grid_markov_chain(n)
The n x n grid graph (n^2 nodes, row-major, nearest-neighbor edges), n >= 2.
markov_chain_from_edge_list
markov_chain_from_edge_list(E)
Uniform random-walk Markov chain on the graph defined by an edge list.
E is a sequence of 0-indexed (i, j) pairs; all edges are unweighted. Returns (Q, pi).
markov_chain_from_adjacency_matrix
markov_chain_from_adjacency_matrix(A)
Uniform random-walk Markov chain from a binary adjacency matrix.
Equivalent to markov_chain_from_weight_matrix(A): for a symmetric 0/1 matrix the two formulas coincide (both normalize by node degree). Returns (Q, pi).
markov_chain_from_weight_matrix
markov_chain_from_weight_matrix(W)
Weighted random-walk Markov chain from a non-negative symmetric weight matrix.
Q[x, y] = W[x, y] / sum(W[x, :]), pi[x] = sum(W[x, :]) / sum(W). pi is stationary for Q (in fact Q is reversible) because W is symmetric, so W must be square, finite, nonnegative and symmetric, and every node needs positive total weight. Returns (Q, pi).
Means
GeometricMean
Bases: AdmissibleMean
theta(s, t) = sqrt(s t). The package default (Erbar et al. 2020).
ArithmeticMean
Bases: AdmissibleMean
theta(s, t) = (s + t) / 2. Does not vanish at an empty node.
HarmonicMean
Bases: AdmissibleMean
theta(s, t) = 2 s t / (s + t). The smallest of the four; bounded partial_s near zero.
LogarithmicMean
Bases: AdmissibleMean
theta(s, t) = (s - t) / (ln s - ln t). The mean for which the heat flow is the gradient flow of the relative entropy (Maas 2011). Not finitely conic-representable; the SOCP uses QuadLogMean instead.
QuadLogMean
Bases: AdmissibleMean
Gauss-Legendre approximation of the logarithmic mean with K nodes, Lambda_K(s, t) = sum_k w_k s^a_k t^(1 - a_k) ~ int_0^1 s^a t^(1-a) da.
Lambda_K is itself an admissible mean, so a metric built on it is exact, not approximate; each term is a power-cone constraint, which is how the SOCP represents the logarithmic mean. Worst relative error against LogarithmicMean over density ratios up to 1000: 8.5e-4 (K=4), 6e-7 (K=6), 1.2e-10 (K=8), 2e-15 (K=12).
Shooting
log_map
log_map(G, nu, target, *, phi0_init=None, tol=1e-09, maxiters=50, nsteps=150, floor_rtol=1e-06, segments='auto', verbose=False)
The Riemannian logarithm of target at nu, by single or multiple
shooting.
Solves F(phi0) = rho(1; nu, phi0) - target = 0 by damped Newton over the
mean-zero potentials: n - 1 unknowns, since the gauge
Initialization is the linearized geodesic L(nu) phi0 = pi * (target - nu),
exact to first order in target - nu, or phi0_init if given -- the
warm-start mechanism: callers keep their own phi0 from a nearby solve. If
the first-order guess overshoots through the positivity floor, as it does
for far-apart concentrated endpoints, it is halved until the first shot
survives.
segments="auto" (the default) starts with single shooting and, if
the line search cut the first Newton step to a quarter or less -- the
signature of a long transport -- switches to multiple shooting with 8
segments from its own default start; if multiple shooting fails, or is
evidently stalling (a first step cut to 1/8, or a residual not halved in
three steps), single shooting carries on from its first step. Short transports stay single
shooting at no extra cost; a long one pays for one single-shooting step.
segments=1 is single shooting throughout.
An integer segments > 1 solves by multiple shooting (shooting.multiple): Newton
for the state at the start of each of that many segments of the
integrator's step grid. It solves the same discrete problem, so the
answer agrees with single shooting's to Newton's tolerance; what changes
is the cost. A long transport needs fewer Newton steps -- 8 at K=4 against
19 for 10x10 corner bumps -- but each step carries about twice the
tangents, so a short transport costs more (the documentation's "Choosing a
method" has the timings). A
phi0_init warm-starts it by one sweep of the flow, cut at the segment
starts; if that sweep hits the positivity floor, the default start is
used. The result's starts holds the solved segment start states.
Returns a LogMapResult. Raises ShootingError if Newton has not reached
tol after maxiters steps, the line search fails, or no admissible
initial potential exists; the SOCP is the fallback in every case, or
log_map_mollified for data near the boundary.
exp_map
exp_map(G, nu, tangent, *, t=1.0, nsteps=150, kind='auto', floor_rtol=1e-06)
The Riemannian exponential map at nu: integrate the Hamiltonian flow
from (nu, phi0) for time t and return the endpoint density.
tangent is a potential phi0 (length n, kind="potential") or a
momentum m0 (length |E|, kind="momentum"), the latter converted with
momentum_to_potential. kind="auto" infers it from the length, which is
ambiguous when n == |E| (a cycle, for instance) and then raises.
Raises PositivityFloorError rather than returning garbage if the flow hits
the positivity floor before time t.
TorchThreadsWarning
Bases: UserWarning
torch is using more than one thread for the shooting solver's small operations, which costs CPU time without speeding them up. Issued once per process; filter it to silence it.
Sinkhorn
ground_cost
ground_cost(G, rule='shortest_path', *, t=None, laziness=0.5, normalize=True)
A ground cost matrix on the nodes of G for the Sinkhorn barycenter.
- "shortest_path" (default): squared BFS hop count (edges of G unweighted).
- "diffusion": squared diffusion distance at time t, D_t(x, y)^2 = sum_z (P^t[x, z] - P^t[y, z])^2 / pi(z), under the lazy chain P = laziness * I + (1 - laziness) * G.Q (same stationary pi). t defaults to the graph diameter. Uses the chain G.Q itself, which coincides with the random walk on the adjacency only for unweighted chains. laziness=0 is the plain walk, as in GraphTransportation.jl.
With normalize=True (default) the matrix is divided by its maximum so entries lie in [0, 1], which keeps the kernel exp(-cost/epsilon) from underflowing at the usual epsilon. Raises if the graph is disconnected.
Degeneracy of the diffusion distance. Spectrally D_t^2 = sum_j lambda_j^(2t) (psi_j(x) - psi_j(y))^2 over the eigenpairs of P, so D_t is a metric whenever P has no zero eigenvalue. A zero eigenvalue removes its coordinates, and D_t(x, y) = 0 exactly when (e_x - e_y) P^t = 0, i.e. e_x - e_y lies in the left kernel of P^t; that need not happen (the plain walk on the 3x3 grid has a 3-dimensional kernel and still separates every pair), but when it does D_t is only a pseudo-metric. For the plain walk this is common: two nodes with the same neighborhood (the ends of a 3-path, opposite corners of a 4-cycle, the leaves of a star) have identical rows of Q, hence zero cost for every t >= 1, and a Sinkhorn barycenter moves mass between them for free. The default laziness 1/2 maps each eigenvalue lambda of Q to (1 + lambda) / 2 >= 0, which sends those pairs to 1/2 and also removes the period-2 oscillation on bipartite graphs; the only zero left is lambda = -1, and it collapses a pair only on the two-node graph (any laziness > 1/2 avoids even that). Because a custom chain or laziness can reintroduce the problem, a cost with a zero between distinct nodes warns, and an identically zero cost raises.