Distance metrics
jaxscape.distance.AbstractDistance
¤
Abstract base class for distance computations on graphs.
Provides a unified interface for computing distances with automatic handling of
coordinate-based (for GridGraph) or index-based node specification.
Arguments:
graph: Graph on which to compute distances.sources: Source nodes as vertex indices (1D array) or coordinates (Nx2 array forGridGraph).targets: Target nodes as vertex indices (1D array) or coordinates (Nx2 array forGridGraph).nodes: Nodes for pairwise distances as vertex indices (1D) or coordinates (Nx2).state: Optional solver setup state. Defaults to automatic initialization.
Specify either: nodes alone, sources and/or targets, or neither
(for all-pairs).
Returns:
Distance array with shape depending on the inputs.
Example
from jaxscape import LCPDistance, GridGraph
import jax.numpy as jnp
distance = LCPDistance()
permeability = jnp.ones((12, 12))
grid = GridGraph(permeability, fun=lambda x, y: (x + y) / 2)
# All-pairs distance
D = distance(grid) # Shape: (n_nodes, n_nodes)
# Using vertex indices
D = distance(
grid,
sources=jnp.array([0, 1]),
targets=jnp.array([10, 20]),
)
# Using coordinates (for GridGraph)
D = distance(
grid,
sources=jnp.array([[0, 0], [1, 1]]),
targets=jnp.array([[10, 10]]),
)
# Pairwise among subset
D = distance(grid, nodes=jnp.array([0, 5, 10])) # Shape: (3, 3)
jaxscape.euclidean_distance.EuclideanDistance
¤
Straight-line distance in grid coordinates. Only works with GridGraph.
Example
from jaxscape import EuclideanDistance
distance = EuclideanDistance()
dist = distance(grid, sources=source_coords, targets=target_coords)
__call__(graph: AbstractGraph, sources: jax.Array | None = None, targets: jax.Array | None = None, nodes: jax.Array | None = None, state: typing.Any = None) -> Array
¤
jaxscape.lcp_distance.LCPDistance
¤
Compute least-cost path distances with edge cost 1/a for conductance a,
using algorithm.
For V vertices, E stored edges and S requested sources:
- Bellman–Ford (default): Relaxes sparse edges until distances stop changing. Each round costs O(S(V + E)), with at most V − 1 rounds. Favor it for sparse graphs or a small number of sources.
- Floyd–Warshall: Computes a dense all-pairs distance matrix by considering each vertex as an intermediate stop. Costs O(V³) time, even when requesting only selected pairs. Favor it for small dense graphs when most or all pairs are needed.
For a unique shortest path, the derivative of its distance with respect to
a is -1/a**2 on that path and zero elsewhere: increasing conductance
decreases the distance. At exact ties, Bellman–Ford chooses a generalized
gradient by splitting sensitivity equally among optimal incoming edges at
each vertex; complete paths need not receive equal weight. Raster gradients
follow by the chain rule through GridGraph.fun.
Attributes:
algorithm:"bellman-ford"(default) or"floyd-warshall".
Note
Bellman–Ford requires positive, finite conductances; zero means no edge.
Gradients assume the graph connections stay fixed. jax.jvp, jax.jacfwd,
and higher-order derivatives are currently unsupported.
Example
from jaxscape import LCPDistance, GridGraph
import jax.numpy as jnp
grid = GridGraph(jnp.ones((5, 5)), fun=lambda x, y: (x + y) / 2)
distance = LCPDistance()
D = distance(grid, sources=jnp.array([0, 1]), targets=jnp.array([10, 20]))
__call__(graph: AbstractGraph, sources: jax.Array | None = None, targets: jax.Array | None = None, nodes: jax.Array | None = None, state: typing.Any = None) -> Array
¤
jaxscape.resistance_distance.ResistanceDistance
¤
Compute the resistance distances.
Attributes:
solver: Optionallineax.AbstractLinearSolver. Must be compatible with BCOO matrices. We currently supportjaxscape.solvers.CholmodSolver,jaxscape.solvers.PyAMGSolver, andjaxscape.solvers.AMJaxCGSolver. If None, uses the pseudo-inverse method, which is very memory intensive for exact all-pairs distances (densifies the Laplacian matrix), or dense solves for approximate distances.method: A resistance distance method. Defaults toExactResistance(). UseSpielmanApproximation(epsilon=..., seed=...)for the randomized Spielman-Srivastava approximation.
Example
import jax.numpy as jnp
from jaxscape import GridGraph, ResistanceDistance, SpielmanApproximation
from jaxscape.solvers import AMJaxCGSolver, PyAMGSolver
permeability = jnp.ones((3, 3))
grid = GridGraph(permeability, fun=lambda x, y: (x + y) / 2)
# Default: pseudo-inverse (small graphs)
distance = ResistanceDistance()
# With solver (large graphs)
distance = ResistanceDistance(solver=PyAMGSolver())
# With an initialized AMG-preconditioned CG solver
distance = ResistanceDistance(
solver=AMJaxCGSolver(rtol=1e-6, atol=1e-6, max_steps=1_000)
)
state = distance.init(grid)
dist = distance(grid, state=state)
# Or reuse only the AMJax preconditioner across related graphs.
preconditioner_state = distance.init_preconditioner(grid)
updated_permeability = permeability * 1.1
updated_grid = GridGraph(updated_permeability, fun=lambda x, y: (x + y) / 2)
dist = distance(updated_grid, state=preconditioner_state)
# Approximate resistance distance
distance = ResistanceDistance(
method=SpielmanApproximation(epsilon=0.05),
solver=PyAMGSolver(),
)
dist = distance(grid)
Warning
The graph must be undirected for resistance distance to be well-defined.
__call__(graph: AbstractGraph, sources: jax.Array | None = None, targets: jax.Array | None = None, nodes: jax.Array | None = None, state: typing.Any = None) -> Array
¤
jaxscape.resistance_distance.ExactResistance
¤
Exact resistance distance via pseudoinverse or grounded linear solves.
jaxscape.resistance_distance.SpielmanApproximation
¤
Spielman-Srivastava randomized resistance distance approximation.
Attributes:
epsilon: Accuracy parameter. Smaller values use more random projections:ceil(log(n_vertices) / epsilon**2), which increases memory use.seed: Random seed for the Rademacher edge projections.projection_batch_size: Number of projection rows evaluated together. Larger batches can improve throughput at the cost of more working memory; choose a size that fits your device.
jaxscape.rsp_distance.RSPDistance
¤
Randomized shortest path distance. Requires the temperature parameter
theta and cost, which can be either a jax.experimental.sparse.BCOO
matrix or a function that will be used to map all non zero element of
the adjacency matrix to create the cost matrix. cost defaults to the
well adapted movement cost function lambda x: -jnp.log(x)).
Warning
This distance metric is experimental and may change in future releases.
Example
from jaxscape import RSPDistance
distance = RSPDistance(theta=0.01, cost=lambda x: -jnp.log(x))
dist = distance(grid)