API Reference

Package

EnergyFlow — Module
EnergyFlow

A 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

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.

source

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$); columns 2:(1+d) = ground-space coordinates (typically rapidity and azimuth).

Keywords

  • backend: solver backend, either a compatibility symbol such as :ns64 or a typed backend strategy such as EnergyFlow.NS64.
  • R: distance scale; ground distances are divided by R before exponentiation. With norm=false, R controls the relative cost of transporting vs. creating/destroying energy (choose R ≥ half the maximum ground distance for a true metric).
  • beta: exponent applied to the scaled ground distance (angular exponent).
  • norm: if true, normalize each event's weights to sum to 1 before solving. If false (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, a GroundMetric instance (default EuclideanMetric()).
  • return_flow: if true, also return the optimal transport plan matrix.
  • strict: if true, 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 φ handling

See also emds, emd!.

source
EnergyFlow.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: an EMDWorkspace (for the network-simplex backends) or a SinkhornWorkspace (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 columns 2:(1+d) = coordinates.

Keywords

  • backend: compatibility symbol or typed strategy. It must match the workspace family, and an exact backend's precision must match the EMDWorkspace element type. A SinkhornWorkspace defaults 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 a SinkhornWorkspace; configure Sinkhorn iterations on the workspace.
  • metric: optional ground-metric override for an EMDWorkspace.
  • return_flow: supported with a SinkhornWorkspace; if true, also return the transport plan.
  • strict: if true, 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]]

See also emd, emds!.

source
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 (each M × (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): a Vector of length n*(n-1)/2 containing the strict upper triangle of the distance matrix in SciPy pdist (row-major) order — i.e. (1,2), (1,3), …, (1,n), (2,3), ….
  • Cross-pairwise: a Matrix D of size (length(events0), length(events1)) with D[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 Matrix

See also emd, emds!.

source
EnergyFlow.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)
source

Backend selection

EnergyFlow.set_backend — Function
set_backend(backend::Union{Symbol,EMDBackend}) -> Symbol

Set the global default EMD solver backend used by emd, emd!, emds and emds! when no backend keyword is given.

Available backends:

BackendSolverPrecisionExact?
:ns64Network simplex (default)Float64yes
:ot64Network simplex with arc mixingFloat64yes
:ns32Network simplexFloat32yes
:ot32Network simplex with arc mixingFloat32yes
:sinkhornEntropy-regularised OT (Sinkhorn)Float64no

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 override

See also get_backend.

source

Ground metrics

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.

source
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.

source
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.

source
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.

source
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)
source

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.

Note

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: NetworkSimplexSolver workspace
  • beta, R: EMD parameters
  • norm: if true, normalize event weights to sum to 1 before solving
  • metric: ground distance metric (GroundMetric); its concrete type is the workspace's M parameter
  • parallel_threshold: fill the cost matrix in parallel when n0 * n1 >= parallel_threshold (default 40_000)

See also SinkhornWorkspace.

source
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; see emd.
  • 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: if true (default), start from a large ε and geometrically anneal down to epsilon, warm-starting each stage — usually faster and more stable than solving directly at small ε.
Note

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.

source

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.

source
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)
source
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.

source

Event Isotropy

EnergyFlow.ring_reference — Function
ring_reference(n) -> n×2 Matrix

Quasi-uniform ring reference event: n equal-weight particles at φ_j = 2π(j - 1/2)/n. Columns are [weight, φ].

source
EnergyFlow.cylinder_reference — Function
cylinder_reference(nphi, ymax) -> m×3 Matrix

Quasi-uniform cylinder reference event: nphi slices in φ and floor(ymax * nphi / π) slices in y over |y| ≤ ymax. Columns are [weight, y, φ].

source
EnergyFlow.sphere_reference — Function
sphere_reference(nval) -> m×4 Matrix

Quasi-uniform spherical reference event built from the centers of the 12·(2^nval)² HEALPix pixels. Columns are [weight, x, y, z].

source
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.

source
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.

source
EnergyFlow.cylinder_metric — Function
cylinder_metric(ymax)

Ground distance on the cylinder, 12/(π² + 16·ymax²)·(Δy² + Δφ²), with the normalization baked in.

source
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.

source
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.

source
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.

source
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 φ (n points)
  • :cylinder — hadron-collider cylinder in (y, φ) (nphi × floor(ymax*nphi/π) grid)
  • :sphere — lepton-collider sphere via HEALPix (nval, so nside = 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.

source

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) -> Float64

Compute 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 (default EuclideanMetric()).
  • return_flow: also return the transport plan when true.
  • strict: raise an error rather than warn if the solver does not converge.

Returns

  • A Float64 EMD value, or (value, plan) when return_flow=true.
source
EnergyFlow.emd_ns64! — Function
emd_ns64!(ws, ev0, ev1; gdim=nothing, n_iter_max=100_000,
          metric=ws.metric, strict=false) -> Float64

Compute 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

  • Float64 EMD value.
source
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: nothing for self-pairwise, or a second Vector of events for cross-pairwise.
  • R, beta, norm, gdim, n_iter_max, metric, strict: same as emd_ns64.

Returns

  • Self-pairwise (events1=nothing): Vector{Float64} of length n*(n-1)/2 in SciPy pdist (upper-triangular row-major) order.
  • Cross-pairwise: Matrix{Float64} of shape (length(events0), length(events1)).
source
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.

source

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) -> Float64

Compute 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.

source
EnergyFlow.emd_ot64! — Function
emd_ot64!(ws, ev0, ev1; gdim=nothing, n_iter_max=100_000,
          metric=ws.metric, strict=false) -> Float64

Compute 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.

source
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.

source
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.

source

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) -> Float32

Compute 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.

source
EnergyFlow.emd_ns32! — Function
emd_ns32!(ws, ev0, ev1; gdim=nothing, n_iter_max=100_000,
          metric=ws.metric, strict=false) -> Float32

Compute EMD using the Network Simplex Float32 backend with a pre-allocated EMDWorkspace{Float32}.

source
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).

source
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.

source

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) -> Float32

Compute 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.

source
EnergyFlow.emd_ot32! — Function
emd_ot32!(ws, ev0, ev1; gdim=nothing, n_iter_max=100_000,
          metric=ws.metric, strict=false) -> Float32

Compute EMD using the OT-style Float32 backend (arc mixing enabled).

source
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).

source
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).

source

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) -> Float64

Compute 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 in emd.
  • 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 (default true).
  • n_iter_max: accepted for compatibility with emd and ignored; use max_iter_sinkhorn to control Sinkhorn iterations.
  • return_flow: if true, also return the transport plan.
  • strict: if true, 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.

source
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.

source
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.

source

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.
source
EnergyFlow.network_simplex! — Function
network_simplex!(ns, source_weights, target_weights; max_iter=100_000) -> Symbol

Solve 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].

source