API Reference
Package
EnergyFlow — Module
EnergyFlowA Julia package for computing the Energy Mover's Distance (EMD) between collider events, with multiple solver backends and ground metrics.
The EMD (Komiske, Metodiev, Thaler, 2019) is an optimal-transport distance between events viewed as energy distributions over a metric space. Its balanced transport term is
\[\mathrm{EMD}_{\beta,R}(\mathcal{E}_0, \mathcal{E}_1) = \min_{f_{ij} \ge 0} \sum_{ij} f_{ij} \left(\frac{d_{ij}}{R}\right)^{\beta}\]
subject to transport constraints, where $d_{ij}$ is the ground distance between particles $i$ and $j$. When norm=false and total event weights differ, a unit-cost fictitious particle balances the transport problem.
Main entry points
emd/emd!: single-pair EMD (allocating / workspace-reusing)emds/emds!: pairwise EMDs over collections of eventsset_backend/get_backend: choose the default solver backendload_hepmc3_events: read events from a HepMC3 ASCII file
Event format
Events are M × (1 + d) matrices with one row per particle: column 1 holds the particle weight (typically $p_T$ or energy) and columns 2:(1+d) hold the ground-space coordinates (typically rapidity $y$ and azimuth $\phi$).
Backends
:ns64 (default), :ot64, :ns32, :ot32 (exact network simplex variants) and :sinkhorn (approximate entropic OT). The high-level APIs also accept the corresponding typed strategies, such as EnergyFlow.NS64. See set_backend.
Ground metrics
EuclideanMetric (default), SquaredEuclideanMetric, EtaPhiMetric, PrecomputedMetric, CustomMetric.
High-level API
The recommended entry points. These dispatch to the currently active backend (see set_backend).
EnergyFlow.emd — Function
emd(ev0, ev1; backend=get_backend(), R=1.0, beta=1.0, norm=false,
gdim=nothing, n_iter_max=100_000, metric=EuclideanMetric(),
return_flow=false, strict=false)Compute the Energy Mover's Distance between two events. Allocates a fresh workspace each call; for repeated calls prefer emd! with a pre-allocated workspace.
Arguments
ev0,ev1: event matrices,M × (1 + d)with one row per particle. Column 1 = particle weights (typically $p_T$); columns2:(1+d)= ground-space coordinates (typically rapidity and azimuth).
Keywords
backend: solver backend, either a compatibility symbol such as:ns64or a typed backend strategy such asEnergyFlow.NS64.R: distance scale; ground distances are divided byRbefore exponentiation. Withnorm=false,Rcontrols the relative cost of transporting vs. creating/destroying energy (chooseR ≥ halfthe maximum ground distance for a true metric).beta: exponent applied to the scaled ground distance (angular exponent).norm: iftrue, normalize each event's weights to sum to 1 before solving. Iffalse(default) and the total weights differ, the weight difference is handled by a fictitious particle at unit cost, matching the Python EnergyFlow convention.gdim: number of coordinate dimensions to use (nothing= all columns after the first).n_iter_max: maximum network-simplex iterations; accepted but ignored by the Sinkhorn backend.metric: ground distance metric, aGroundMetricinstance (defaultEuclideanMetric()).return_flow: iftrue, also return the optimal transport plan matrix.strict: iftrue, raise an error when the selected solver does not converge.
Returns
The EMD value by default. If return_flow=true, returns (emd_value, transport_plan). When norm=true, the transport plan is normalized so its row/column sums are 1 and the marginals sum to 1. When norm=false, the internal solve is performed on a scaled balanced problem, and the returned plan is rescaled back to the original physical weights, including any fictitious source/target row or column introduced by the solver so its marginals match those original weights.
Example
ev0 = [1.0 0.5 0.1; 0.8 -0.3 0.4] # 2 particles: pT, y, φ
ev1 = [0.9 0.4 0.0; 0.7 -0.2 0.3]
emd(ev0, ev1) # default backend/metric
emd(ev0, ev1; R=0.4, beta=2.0, norm=true) # EMD parameters
emd(ev0, ev1; backend=:sinkhorn) # approximate solver
emd(ev0, ev1; backend=EnergyFlow.NS64) # typed strategy
emd(ev0, ev1; metric=EtaPhiMetric()) # periodic φ handlingEnergyFlow.emd! — Function
emd!(ws, ev0, ev1; backend=get_backend(), gdim=nothing,
n_iter_max=100_000, metric=nothing, return_flow=false, strict=false)Compute the EMD between two events using a pre-allocated workspace. This reuses the solver's large internal buffers and reduces per-call allocations when computing many EMDs in a loop. Temporary arrays are still allocated while the public API unpacks each event matrix.
The EMD parameters (beta, R, and norm) are read from the workspace. For an EMDWorkspace, metric=nothing uses the workspace metric and an explicit metric temporarily overrides it for that call. Sinkhorn currently supports only EuclideanMetric.
Arguments
ws: anEMDWorkspace(for the network-simplex backends) or aSinkhornWorkspace(for the Sinkhorn backend). Must be large enough for the events being compared.ev0,ev1: event matrices,M × (1 + d)with column 1 = particle weights and columns2:(1+d)= coordinates.
Keywords
backend: compatibility symbol or typed strategy. It must match the workspace family, and an exact backend's precision must match theEMDWorkspaceelement type. ASinkhornWorkspacedefaults to:sinkhorn.gdim: number of coordinate dimensions to use (nothing= all columns after the first).n_iter_max: maximum network-simplex iterations. It is accepted but ignored for aSinkhornWorkspace; configure Sinkhorn iterations on the workspace.metric: optional ground-metric override for anEMDWorkspace.return_flow: supported with aSinkhornWorkspace; iftrue, also return the transport plan.strict: iftrue, raise an error when the solver does not converge.
Returns
The EMD value (Float64 or Float32 matching the workspace precision). With a SinkhornWorkspace and return_flow=true, returns (emd_value, transport_plan). The exact EMDWorkspace overload returns only the value; use emd when an exact transport plan is needed.
Example
n = maximum(size(e, 1) for e in events)
ws = EMDWorkspace(n, n; beta=1.0, R=1.0, norm=true)
vals = [emd!(ws, events[1], e) for e in events[2:end]]EnergyFlow.emds — Function
emds(events0, events1=nothing; backend=get_backend(), R=1.0, beta=1.0,
norm=false, gdim=nothing, n_iter_max=100_000,
metric=EuclideanMetric(), strict=false)Compute pairwise Energy Mover's Distances between collections of events. Computation is multithreaded over pairs (start Julia with -t auto or set JULIA_NUM_THREADS to benefit).
Arguments
events0: vector of event matrices (eachM × (1 + d)).events1:nothing(default) for self-pairwise mode, or a second vector of events for cross-pairwise mode.
Keywords
Same as emd: backend, R, beta, norm, gdim, n_iter_max, metric, and strict. The backend may be a compatibility symbol or a typed strategy. Sinkhorn supports only EuclideanMetric through this interface.
Returns
- Self-pairwise (
events1 === nothing): aVectorof lengthn*(n-1)/2containing the strict upper triangle of the distance matrix in SciPypdist(row-major) order — i.e.(1,2), (1,3), …, (1,n), (2,3), …. - Cross-pairwise: a
MatrixDof size(length(events0), length(events1))withD[i, j] = emd(events0[i], events1[j]).
Element type is Float64, or Float32 for the :ns32/:ot32 backends and their corresponding typed strategies.
Example
events = load_hepmc3_events("events.hepmc"; maxevents=20)
dists = emds(events[1:10]; norm=true) # 45-element Vector
D = emds(events[1:10], events[11:20]; norm=true) # 10×10 MatrixEnergyFlow.emds! — Function
emds!(results, events0; backend=get_backend(), R=1.0, beta=1.0, norm=false,
gdim=nothing, n_iter_max=100_000, metric=EuclideanMetric(),
strict=false)Preallocated-output self-pairwise EMD: computes all n*(n-1)/2 pairwise distances among events0 and writes them into results (flat upper-triangular, SciPy pdist row-major order — see emds). Internal worker workspaces are created for each call.
Computation is multithreaded over pairs. Keywords are as in emds. The results element type must match the exact backend precision. Sinkhorn is not supported by emds!; use emds instead.
Example
n = length(events)
results = Vector{Float64}(undef, n * (n - 1) ÷ 2)
emds!(results, events; norm=true)Backend selection
EnergyFlow.set_backend — Function
set_backend(backend::Union{Symbol,EMDBackend}) -> SymbolSet the global default EMD solver backend used by emd, emd!, emds and emds! when no backend keyword is given.
Available backends:
| Backend | Solver | Precision | Exact? |
|---|---|---|---|
:ns64 | Network simplex (default) | Float64 | yes |
:ot64 | Network simplex with arc mixing | Float64 | yes |
:ns32 | Network simplex | Float32 | yes |
:ot32 | Network simplex with arc mixing | Float32 | yes |
:sinkhorn | Entropy-regularised OT (Sinkhorn) | Float64 | no |
The compatibility symbols map to the typed strategies EnergyFlow.NS64, EnergyFlow.OT64, EnergyFlow.NS32, EnergyFlow.OT32, and EnergyFlow.Sinkhorn, respectively. Either form may be passed. The canonical symbol is stored and returned. An unknown symbol raises ArgumentError.
Example
set_backend(:ot64)
emd(ev0, ev1) # now uses :ot64
emd(ev0, ev1; backend=:ns64) # per-call override
emd(ev0, ev1; backend=EnergyFlow.NS64) # typed per-call overrideSee also get_backend.
EnergyFlow.get_backend — Function
get_backend() -> SymbolReturn the current global default EMD backend (initially :ns64).
See also set_backend.
Ground metrics
EnergyFlow.GroundMetric — Type
GroundMetricAbstract supertype for ground distance metrics used to fill the pairwise particle cost matrix in an EMD computation.
Concrete subtypes: EuclideanMetric, SquaredEuclideanMetric, EtaPhiMetric, PrecomputedMetric, CustomMetric.
Pass an instance via the metric keyword of emd / emds, or store it in an EMDWorkspace for the in-place API.
EnergyFlow.EuclideanMetric — Type
EuclideanMetric()Euclidean ground distance (the default metric).
The cost of moving weight between particles i and j is
cost_ij = (√(Σₖ Δxₖ²) / R)^βcomputed over all coordinate dimensions of the event.
EnergyFlow.SquaredEuclideanMetric — Type
SquaredEuclideanMetric()Squared Euclidean ground distance (no square root):
cost_ij = (Σₖ Δxₖ² / R)^βEquivalent to EuclideanMetric with beta=2 when R = 1. Useful when a Wasserstein-2-style cost is desired, or to skip the sqrt for speed.
EnergyFlow.EtaPhiMetric — Type
EtaPhiMetric()Collider-aware (η, φ) ground distance with periodic azimuthal angle:
cost_ij = (√(Δη² + Δφ_wrap²) / R)^β, Δφ_wrap = π - |mod(|Δφ|, 2π) - π|Assumes coordinate dimension 1 is η (rapidity or pseudorapidity) and dimension 2 is φ (azimuthal angle, radians). This is the appropriate metric for events in [pT, η, φ] format when particles may lie near the φ = ±π boundary, where a plain EuclideanMetric would overestimate distances.
EnergyFlow.PrecomputedMetric — Type
PrecomputedMetric(costs::AbstractMatrix)Use a caller-supplied n0 × n1 cost matrix directly, where costs[i, j] is the cost of transporting weight from particle i of the first event to particle j of the second.
Values are used as-is — R and beta are not applied. The matrix must be at least as large as the particle counts of the event pair it is used with, so this metric is mainly useful for single-pair computations.
EnergyFlow.CustomMetric — Type
CustomMetric(dist_fn)Compute ground distances with a user-supplied function.
dist_fn(p, q) receives views of the coordinate columns of one particle from each event (length-d vectors) and must return a scalar distance. R and beta are applied on top of the returned value:
cost_ij = (dist_fn(pᵢ, qⱼ) / R)^βExample
# Manhattan (L1) ground distance
m = CustomMetric((p, q) -> sum(abs, p .- q))
emd(ev0, ev1; metric=m)Workspaces
EnergyFlow.EMDWorkspace — Type
EMDWorkspace{V<:AbstractFloat, M<:GroundMetric}Pre-allocated workspace for repeated EMD computations with the network-simplex backends. Construct once, then pass to emd! (or the backend-specific emd_ns64! etc.) to reuse the solver's internal buffers and reduce per-call allocations. The public event-matrix APIs still allocate temporary arrays while unpacking their inputs.
The EMD parameters (beta, R, norm, metric) are stored in the workspace and used by every solve.
The workspace is specialized on its ground-metric type M (the second type parameter). Updating beta, R, or norm in place is fine, but switching to a different metric type requires constructing a new workspace — assigning a metric of another type to ws.metric will error, because the field is typed M.
Fields
ns:NetworkSimplexSolverworkspacebeta,R: EMD parametersnorm: iftrue, normalize event weights to sum to 1 before solvingmetric: ground distance metric (GroundMetric); its concrete type is the workspace'sMparameterparallel_threshold: fill the cost matrix in parallel whenn0 * n1 >= parallel_threshold(default40_000)
See also SinkhornWorkspace.
EnergyFlow.SinkhornWorkspace — Type
SinkhornWorkspace{V<:AbstractFloat}
SinkhornWorkspace(max_n0, max_n1; beta=1.0, R=1.0, norm=true,
epsilon=0.01, max_iter=5000, tol=1e-9, annealing=true)Pre-allocated workspace for approximate EMD computation with the Sinkhorn (entropy-regularised OT) backend. Pass to emd! or emd_sinkhorn!. The workspace reuses the solver's internal buffers, although the public event-matrix APIs still allocate temporary arrays while unpacking inputs. The unparameterized constructor defaults to Float64.
Arguments
max_n0,max_n1: maximum particle counts of the two events in any subsequent solve.
Keywords
beta,R,norm: EMD parameters; seeemd.epsilon: entropic regularisation strength. Smaller values approximate the exact EMD more closely but require more iterations.max_iter: maximum Sinkhorn iterations per ε-level.tol: convergence tolerance on the marginal violation.annealing: iftrue(default), start from a large ε and geometrically anneal down toepsilon, warm-starting each stage — usually faster and more stable than solving directly at small ε.
Sinkhorn solves the regularised transport problem, so the returned cost carries an entropic bias that shrinks with epsilon but never fully vanishes. Use the network-simplex backends for exact values.
After a solve, the fields total_cost, n_iters and converged hold diagnostics for the last computation.
Event I/O
EnergyFlow.HepMC3 — Module
Minimal utilities for reading HepMC3 ASCII files.
Most users should call EnergyFlow.load_hepmc3_events for [pT, η, φ] matrices or EnergyFlow.load_hepmc3_momenta for [E, px, py, pz] matrices. The lower-level read_events interface passes vectors of Particle objects to a callback.
EnergyFlow.load_hepmc3_events — Function
load_hepmc3_events(filepath; maxevents=-1, skipevents=0, status=1, min_pt=0.0)
-> Vector{Matrix{Float64}}Load events from a HepMC3 ASCII file and convert to EnergyFlow event matrices. Each event is an M×3 Matrix{Float64} with columns [pT, η, φ].
Only final-state particles with the given status code (default 1) are kept. Particles with pT ≤ min_pt (including beam-axis particles with pT = 0) are skipped. This is a standalone parser; it does not require LorentzVectorHEP.
HepMC3 ASCII P-line format
P barcode vertex pdgid px py pz energy mass status [...]
Coordinate conventions
- pT = √(px² + py²) — particle weight
- η = asinh(pz / pT) — pseudorapidity
- φ = atan(py, px) ∈ (-π,π] — azimuthal angle (radians)
EnergyFlow.HepMC3.load_hepmc3_momenta — Function
load_hepmc3_momenta(filepath; maxevents=-1, status=1) -> Vector{Matrix{Float64}}Load events from a HepMC3 ASCII file as M×4 matrices with columns [E, px, py, pz], suitable for spherical isotropy. maxevents=-1 reads all events and status=1 keeps final-state particles.
Event Isotropy
EnergyFlow.ring_reference — Function
ring_reference(n) -> n×2 MatrixQuasi-uniform ring reference event: n equal-weight particles at φ_j = 2π(j - 1/2)/n. Columns are [weight, φ].
EnergyFlow.cylinder_reference — Function
cylinder_reference(nphi, ymax) -> m×3 MatrixQuasi-uniform cylinder reference event: nphi slices in φ and floor(ymax * nphi / π) slices in y over |y| ≤ ymax. Columns are [weight, y, φ].
EnergyFlow.sphere_reference — Function
sphere_reference(nval) -> m×4 MatrixQuasi-uniform spherical reference event built from the centers of the 12·(2^nval)² HEALPix pixels. Columns are [weight, x, y, z].
EnergyFlow.healpix_pix2vec_ring — Function
healpix_pix2vec_ring(nside) -> Vector{NTuple{3,Float64}}Unit 3-vectors at the centers of the 12·nside² HEALPix pixels in the RING scheme. Reproduces astropy_healpix.pix2vec closely and is used internally by sphere_reference. nside must be positive.
EnergyFlow.ring_cos_metric — Function
ring_cos_metric()Ground distance on the ring, (π/(π-2))·(1 - cos Δφ), with the normalization baked in so isotropy uses R=1, beta=1, norm=true.
EnergyFlow.ring_phi_metric — Function
ring_phi_metric()Ground distance on the ring, (4/π)·Δφ, with the normalization baked in.
EnergyFlow.cylinder_metric — Function
cylinder_metric(ymax)Ground distance on the cylinder, 12/(π² + 16·ymax²)·(Δy² + Δφ²), with the normalization baked in.
EnergyFlow.sphere_cos_metric — Function
sphere_cos_metric()Ground distance on the sphere, 2·(1 - cos θ) between 3-momentum directions.
EnergyFlow.sphere_angular_metric — Function
sphere_angular_metric()Angular ground distance on the sphere, arccos(cos θ).
EnergyFlow.select_rapidity — Function
select_rapidity(events, ymax; min_particles=2) -> Vector{Matrix}Physics preselection for hadron-collider isotropy: keep particles with |y| ≤ ymax and drop events with fewer than min_particles accepted particles. This is a deliberate acceptance choice and should be applied consistently when comparing implementations.
EnergyFlow.select_sphere_events — Function
select_sphere_events(events; min_particles=2) -> Vector{Matrix}Physics preselection for spherical isotropy: drop zero-momentum particles and any event with fewer than min_particles remaining. This is the acceptance choice for the lepton-collider version of the observable.
EnergyFlow.event_isotropy — Function
event_isotropy(event, ref, metric; backend=:ns64, n_iter_max=100_000,
recoil=false)Compute event isotropy as the EMD between event and a quasi-uniform reference event ref using the supplied ground metric. The isotropy normalization is baked into the reference and metric, so the EMD is evaluated with R=1, beta=1, norm=true. If recoil=true, a single pseudo-particle is added to balance the event's net momentum before evaluating the EMD. backend may be a compatibility symbol such as :ns64 or a typed strategy such as EnergyFlow.NS64.
event_isotropy(event; geometry=:ring, n=128, nphi=16, ymax=4.0, nval=2,
backend=:ns64, n_iter_max=100_000, recoil=false)Convenience wrapper that builds the standard isotropy reference and metric for the requested geometry and then computes the observable.
Supported geometries:
:ring— hadron-collider ring inφ(npoints):cylinder— hadron-collider cylinder in(y, φ)(nphi × floor(ymax*nphi/π)grid):sphere— lepton-collider sphere via HEALPix (nval, sonside = 2^nval)
For :ring, 3-column hadron-collider events are projected to [pT, φ]. For :cylinder, particles with |y| > ymax are dropped before evaluation. For :sphere, zero-momentum particles are dropped before evaluation. Set recoil=true to append a balancing pseudo-particle before the EMD is computed, which is useful when the input event has residual net momentum. backend accepts either a compatibility symbol or a typed backend strategy.
Backend-specific functions
Direct entry points for each solver backend, bypassing backend dispatch.
Network simplex, Float64 (:ns64)
EnergyFlow.emd_ns64 — Function
emd_ns64(ev0, ev1; R=1.0, beta=1.0, norm=false, gdim=nothing,
n_iter_max=100_000, metric=EuclideanMetric(),
return_flow=false, strict=false) -> Float64Compute EMD using the Network Simplex Float64 backend. Allocates a fresh workspace each call. For repeated calls, prefer emd_ns64!.
Equivalent to emd with backend=:ns64.
Arguments
ev0,ev1: M×(1+gdim) matrices. Column 1 = particle weights (pT); columns 2:(1+gdim) = spatial coordinates.R: distance scale parameter (default 1.0).beta: distance exponent (default 1.0).norm: normalize weights to sum to 1 before solving (default false).gdim: number of coordinate dimensions to use.nothing= use all remaining columns.n_iter_max: max network-simplex iterations (default 100_000).metric: ground distance metric (defaultEuclideanMetric()).return_flow: also return the transport plan whentrue.strict: raise an error rather than warn if the solver does not converge.
Returns
- A
Float64EMD value, or(value, plan)whenreturn_flow=true.
EnergyFlow.emd_ns64! — Function
emd_ns64!(ws, ev0, ev1; gdim=nothing, n_iter_max=100_000,
metric=ws.metric, strict=false) -> Float64Compute EMD using the Network Simplex Float64 backend with a pre-allocated EMDWorkspace. The workspace's beta, R, norm settings are used.
Arguments
ws::EMDWorkspace: pre-allocated workspace.ev0,ev1: EnergyFlow-format event matrices (M×(1+gdim)). Column 1 = particle weights (pT); columns 2:(1+gdim) = coordinates.gdim: coordinate dimensions to use (nothing= all).n_iter_max: maximum network-simplex iterations.metric: ground metric for this call (default: the workspace metric).strict: raise an error rather than warn if the solver does not converge.
Returns
Float64EMD value.
EnergyFlow.emds_ns64 — Function
emds_ns64(events0, events1=nothing; R=1.0, beta=1.0, norm=false,
gdim=nothing, n_iter_max=100_000, metric=EuclideanMetric(),
strict=false)Compute pairwise EMDs using the Network Simplex Float64 backend.
Arguments
events0: Vector of event matrices (each M x (1+gdim)).events1:nothingfor self-pairwise, or a second Vector of events for cross-pairwise.R,beta,norm,gdim,n_iter_max,metric,strict: same asemd_ns64.
Returns
- Self-pairwise (
events1=nothing):Vector{Float64}of lengthn*(n-1)/2in SciPy pdist (upper-triangular row-major) order. - Cross-pairwise:
Matrix{Float64}of shape(length(events0), length(events1)).
EnergyFlow.emds_ns64! — Function
emds_ns64!(results, events0; R=1.0, beta=1.0, norm=false, gdim=nothing,
n_iter_max=100_000, metric=EuclideanMetric(), strict=false)Writes pairwise EMD values to the preallocated results vector (flat upper-triangular, length n*(n-1)/2). Internal worker workspaces are created for each call.
OT-style arc mixing, Float64 (:ot64)
EnergyFlow.emd_ot64 — Function
emd_ot64(ev0, ev1; R=1.0, beta=1.0, norm=false, gdim=nothing,
n_iter_max=100_000, metric=EuclideanMetric(),
return_flow=false, strict=false) -> Float64Compute EMD using the OT-style Float64 backend (arc mixing enabled). Allocates a fresh workspace each call. For repeated calls, prefer emd_ot64!. Arguments are as in emd_ns64.
EnergyFlow.emd_ot64! — Function
emd_ot64!(ws, ev0, ev1; gdim=nothing, n_iter_max=100_000,
metric=ws.metric, strict=false) -> Float64Compute EMD using the OT-style Float64 backend (arc mixing enabled). Uses the same NetworkSimplex solver but with POT-style arc interleaving for improved performance on unbalanced transport problems.
EnergyFlow.emds_ot64 — Function
emds_ot64(events0, events1=nothing; R=1.0, beta=1.0, norm=false,
gdim=nothing, n_iter_max=100_000, metric=EuclideanMetric(),
strict=false)Pairwise EMDs using the OT-style Float64 backend (arc mixing enabled). See emds_ns64 for argument documentation.
EnergyFlow.emds_ot64! — Function
emds_ot64!(results, events0; R=1.0, beta=1.0, norm=false, gdim=nothing,
n_iter_max=100_000, metric=EuclideanMetric(), strict=false)In-place pairwise EMDs using the OT-style Float64 backend (arc mixing enabled). See emds_ns64! for argument documentation.
Network simplex, Float32 (:ns32)
EnergyFlow.emd_ns32 — Function
emd_ns32(ev0, ev1; R=1.0, beta=1.0, norm=false, gdim=nothing,
n_iter_max=100_000, metric=EuclideanMetric(),
return_flow=false, strict=false) -> Float32Compute EMD using the Network Simplex Float32 backend. Allocates a fresh workspace each call. For repeated calls, prefer emd_ns32!. Arguments are as in emd_ns64.
EnergyFlow.emd_ns32! — Function
emd_ns32!(ws, ev0, ev1; gdim=nothing, n_iter_max=100_000,
metric=ws.metric, strict=false) -> Float32Compute EMD using the Network Simplex Float32 backend with a pre-allocated EMDWorkspace{Float32}.
EnergyFlow.emds_ns32 — Function
emds_ns32(events0, events1=nothing; R=1.0, beta=1.0, norm=false,
gdim=nothing, n_iter_max=100_000, metric=EuclideanMetric(),
strict=false)Pairwise EMDs using the Network Simplex Float32 backend. Returns Vector{Float32} (self) or Matrix{Float32} (cross).
EnergyFlow.emds_ns32! — Function
emds_ns32!(results, events0; R=1.0, beta=1.0, norm=false, gdim=nothing,
n_iter_max=100_000, metric=EuclideanMetric(), strict=false)In-place pairwise EMDs using the Network Simplex Float32 backend.
OT-style arc mixing, Float32 (:ot32)
EnergyFlow.emd_ot32 — Function
emd_ot32(ev0, ev1; R=1.0, beta=1.0, norm=false, gdim=nothing,
n_iter_max=100_000, metric=EuclideanMetric(),
return_flow=false, strict=false) -> Float32Compute EMD using the OT-style Float32 backend (arc mixing enabled). Allocates a fresh workspace each call. For repeated calls, prefer emd_ot32!. Arguments are as in emd_ns64.
EnergyFlow.emd_ot32! — Function
emd_ot32!(ws, ev0, ev1; gdim=nothing, n_iter_max=100_000,
metric=ws.metric, strict=false) -> Float32Compute EMD using the OT-style Float32 backend (arc mixing enabled).
EnergyFlow.emds_ot32 — Function
emds_ot32(events0, events1=nothing; R=1.0, beta=1.0, norm=false,
gdim=nothing, n_iter_max=100_000, metric=EuclideanMetric(),
strict=false)Pairwise EMDs using the OT-style Float32 backend (arc mixing enabled).
EnergyFlow.emds_ot32! — Function
emds_ot32!(results, events0; R=1.0, beta=1.0, norm=false, gdim=nothing,
n_iter_max=100_000, metric=EuclideanMetric(), strict=false)In-place pairwise EMDs using the OT-style Float32 backend (arc mixing enabled).
Sinkhorn (:sinkhorn)
EnergyFlow.emd_sinkhorn — Function
emd_sinkhorn(ev0, ev1; R=1.0, beta=1.0, norm=false, gdim=nothing,
epsilon=0.01, max_iter_sinkhorn=5000, sinkhorn_tol=1e-9,
annealing=true, n_iter_max=100_000, return_flow=false,
strict=false) -> Float64Compute an approximate EMD between two events using the Sinkhorn (entropy-regularised OT) solver. Allocates a fresh workspace each call; for repeated calls prefer emd_sinkhorn!.
Keywords
R,beta,norm,gdim: as inemd.epsilon: entropic regularisation strength (smaller = closer to exact EMD, slower to converge).max_iter_sinkhorn: maximum Sinkhorn iterations per ε-level.sinkhorn_tol: convergence tolerance on the marginal violation.annealing: use ε-annealing with warm starts (defaulttrue).n_iter_max: accepted for compatibility withemdand ignored; usemax_iter_sinkhornto control Sinkhorn iterations.return_flow: iftrue, also return the transport plan.strict: iftrue, raise an error when the solver does not converge; otherwise a finite last iterate may be returned.
The returned value includes an entropic bias relative to the exact EMD; see SinkhornWorkspace.
EnergyFlow.emd_sinkhorn! — Function
emd_sinkhorn!(ws::SinkhornWorkspace, ev0, ev1; gdim=nothing,
return_flow=false, strict=false)Compute an approximate EMD using the Sinkhorn backend with a pre-allocated SinkhornWorkspace. All Sinkhorn parameters (beta, R, norm, epsilon, tol, annealing, max_iter) are read from the workspace.
Returns the regularised transport cost (same precision as the workspace). If return_flow=true, also returns the transport plan. A finite last iterate may be returned when convergence is not reached; inspect ws.converged, or set strict=true to raise an error on non-convergence.
EnergyFlow.emds_sinkhorn — Function
emds_sinkhorn(events0, events1=nothing; R=1.0, beta=1.0, norm=false,
gdim=nothing, epsilon=0.01, max_iter_sinkhorn=5000,
sinkhorn_tol=1e-9, annealing=true, n_iter_max=100_000,
strict=false, metric=EuclideanMetric())Pairwise approximate EMDs using the Sinkhorn backend, multithreaded over pairs. Return conventions match emds: a flat upper-triangular Vector{Float64} for self-pairwise mode (events1 === nothing), or a Matrix{Float64} for cross-pairwise mode.
Sinkhorn-specific keywords are as in emd_sinkhorn. n_iter_max is accepted for compatibility with emds and ignored; use max_iter_sinkhorn to control Sinkhorn iterations.
Low-level solver
The network simplex solver underlying the exact backends. Most users should not need these directly.
EnergyFlow.NetworkSimplexSolver — Type
NetworkSimplexSolver{V<:AbstractFloat}
NetworkSimplexSolver(max_n0, max_n1)Pre-allocated workspace for the network simplex algorithm on a complete bipartite transportation graph with up to max_n0 sources and max_n1 targets. All arrays are allocated once for the maximum problem size; the solver itself (network_simplex!) performs no hot-path allocation.
The unparameterized constructor defaults to Float64.
This is the low-level solver underlying the :ns64/:ot64/:ns32/:ot32 backends; most users should use emd/emds instead. After a solve, total_cost, status, and n_iters hold the result and diagnostics.
Based on LEMON's network simplex implementation, with optional parallel entering-arc search following Kara & Özturan (2022).
The entering-arc rule is selected by the pivot_mode field:
:serial(default) — single-threaded block search.:parallel_block— block search with the block split across a persistent worker pool (Kara & Özturan). The mode to use when threads are available.:full_parallel— every pivot scans all arcs in parallel and takes the globally most-negative reduced cost. Kept as a simple baseline for benchmarking; it is expected to be slower than:parallel_block.
EnergyFlow.network_simplex! — Function
network_simplex!(ns, source_weights, target_weights; max_iter=100_000) -> SymbolSolve the minimum-cost transportation problem between source_weights and target_weights on the complete bipartite graph, using the pre-allocated NetworkSimplexSolver ns.
Arc costs must be filled in ns.costs[1:n0*n1] beforehand, in row-major order: the arc from source i to target j is at index (i-1)*n1 + j. Source and target weights must sum to the same total (within tolerance).
Returns a status Symbol, one of :optimal, :infeasible, :unbounded, :supply_mismatch, or :max_iter. On success the objective value is in ns.total_cost and the optimal flows in ns.flows[1:n0*n1].