Wasserstein Network (wnet) is a Python/C++ library for working with Wasserstein distances. It uses the Min Cost Flow algorithm as implemented by the LEMON library, exposed to Python via the pylmcf module, enabling efficient computation and manipulation of Wasserstein distances between multidimensional distributions.
- Wasserstein and Truncated Wasserstein distance between multidimensional distributions (dimensions 1–20)
- Three distance metrics: L1, L2, L∞
- Order-p (Lp) Wasserstein for any real p ≥ 1 (per-pair cost =
ground_distance**p, result is the p-th root; fractional p via automatic cost scaling) - Derivatives with respect to peak intensities and spectrum mixture proportions
- Position gradients (∂cost/∂position) with warm-restart re-solving after peak position updates
- Support for distribution mixtures and efficient recalculation with changed mixture proportions
- Picklable
Distributionobjects
You can install the Python package using pip:
pip install wnetimport numpy as np
from wnet import WassersteinDistance, Distribution
from wnet.distances import DistanceMetric
positions1 = np.array([[0, 1, 5, 10], [0, 0, 0, 3]])
intensities1 = np.array([10, 5, 5, 5])
positions2 = np.array([[1, 10], [0, 0]])
intensities2 = np.array([20, 5])
S1 = Distribution(positions1, intensities1)
S2 = Distribution(positions2, intensities2)
print(WassersteinDistance(S1, S2, DistanceMetric.L1))
# 45.0By default the distance is the 1-Wasserstein distance. Pass p to use the order-p
Wasserstein distance, where each unit of mass moved a ground distance d costs d**p
and the returned value is the p-th root of the optimal transport cost:
from wnet import TruncatedWassersteinDistance
# Quadratic (W2) Wasserstein distance with a Euclidean ground metric
print(WassersteinDistance(S1, S2, DistanceMetric.L2, p=2))
print(TruncatedWassersteinDistance(S1, S2, DistanceMetric.L2, max_distance=3.0, p=2))
# Fractional orders work too (e.g. p = 1.5)
print(WassersteinDistance(S1, S2, DistanceMetric.L2, p=1.5))p can be any real number ≥ 1. The ground metric (L1/L2/L∞) is chosen independently of p.
For p != 1 the dense transport network is used by default: the chain solver's hop costs are
additive along the chain, which exponentiated step costs are not. In 1D there is a chain-native
alternative, ConvexSweep, which sweeps the sorted positions and prices pairs directly; select
it with split_distance= and solver=ConvexSweep().
p == 1 is bit-exact with the classic 1-Wasserstein distance. For p != 1 the cost
d**p is fractional, so the integer min-cost-flow solver works in automatically scaled
units (round(scale_factor() * d**p)); the public results divide that scale back out, so
no tuning is needed.
At the
WassersteinNetworklevel,total_cost()and all derivatives are inW_p**punits (the sum ofd**p · flow); take the p-th root for the literalW_pdistance, asWassersteinDistance()does.
Mass that cannot be matched within max_distance is discarded at a fixed cost rather than transported arbitrarily far:
from wnet import TruncatedWassersteinDistance
print(TruncatedWassersteinDistance(S1, S2, DistanceMetric.L2, max_distance=3.0))signal_part_derivatives() returns the marginal cost of increasing each theoretical peak's intensity by 1 — useful for scoring how well each peak is explained. It needs an escape route, so add one of the trash edges before build():
from wnet import WassersteinNetwork
W = WassersteinNetwork(S1, [S2], DistanceMetric.L2, max_distance=10.0)
W.add_simple_trash(10.0)
W.build()
W.solve()
derivs = W.signal_part_derivatives() # {spectrum_id: {peak_index: derivative}}
print(derivs[0]) # marginals for the peaks of S2
# {0: 10.0, 1: 10.0}The keys are the theoretical spectra in the order they were passed, so derivs[k][i] is the marginal for peak i of the k-th target. spectrum_proportion_derivatives() returns the gradient with respect to scaling each spectrum's proportion, as an np.ndarray indexed by spectrum.
Both are in W_p**p units, like total_cost().
After an initial solve, positions can be updated and re-solved cheaply via a warm restart. update_positions_and_get_gradient(new_base, new_targets) takes replacement Distribution objects with the same peak counts, re-solves, and returns ∂cost/∂position for all peaks so you can feed them into a gradient-based optimiser:
W = WassersteinNetwork(S1, [S2], DistanceMetric.L2, max_distance=10.0)
W.add_simple_trash(10.0)
W.build()
W.solve()
positions = S1.positions.copy() # [DIM, N]; the property is a read-only view
for _ in range(100):
moved = Distribution(positions, S1.intensities)
grad_empirical, grad_theoretical = W.update_positions_and_get_gradient(moved, [S2])
# gradients are [N, DIM], positions are [DIM, N]
positions -= 0.01 * grad_empirical.Tgrad_theoretical is a list holding one [N_k, DIM] array per target spectrum. Gradients are of total_cost(), i.e. the W_p**p objective; for the literal W_p multiply by (1/p) * total_cost()**(1/p - 1).
If you only want the re-solve and not the gradient, update_positions_and_solve(new_base, new_targets) is the cheaper call. Both keep the graph topology fixed, so the peak counts may not change; on 1D chain networks the peaks may also not cross one another.
MIT Licence
If you use this software, please cite:
Król J, Bochenek M, Jopa S, Kazimierczuk K, Gambin A, Startek MP (2026). WNetAlign: fast and accurate spectra alignment using truncated Wasserstein distance and network simplex. Briefings in Bioinformatics, 27(3), bbag247. https://doi.org/10.1093/bib/bbag247
@article{krol2026wnetalign,
title = {WNetAlign: fast and accurate spectra alignment using truncated Wasserstein distance and network simplex},
author = {Kr{\'o}l, Justyna and Bochenek, Maria and Jopa, Sylwia and Kazimierczuk, Krzysztof and Gambin, Anna and Startek, Micha{\l} Piotr},
journal = {Briefings in Bioinformatics},
volume = {27},
number = {3},
pages = {bbag247},
year = {2026},
doi = {10.1093/bib/bbag247}
}