Python 3.10+ · JAX · MIT License

finslerax: Differentiable Finsler geometry in JAX

finslerax is a JAX-native library for differentiable Finsler geometry. Its basic object is a cost $F(x,v)$ , attached to a position and a direction of travel. From it the library computes geodesics, arrival times, parallel transport and curvature; when $F$ contains trainable arrays, gradients return to them.

Everything descends from the energy $E = \tfrac{1}{2}F^2$ by automatic differentiation. Its velocity Hessian is the fundamental tensor; the Euler–Lagrange equations give the spray; differentiating the spray gives the nonlinear connection and curvature. The common pipeline begins atmetric_fn, while named metric families provide cheaper formulas and parametrisations that preserve their mathematical constraints.

bash

pip install finslerax

Get started Source Concepts

A first example

In still water the Euclidean straight line is both shortest and quickest. A uniform current leaves that line straight but changes its travel time, and a varying current may bend the quickest route as well. In either case the cost of a step depends on its direction, so reversing a path need not preserve its cost.

python

import finslerax as fax
import jax.numpy as jnp

# Still sea, steady eastward wind at 0.3 of the craft's own speed.
metric = fax.Randers(fax.EuclideanSpace(dim=2),
                     h_net=lambda x: jnp.eye(2),
                     w_net=lambda x: jnp.array([0.3, 0.0]))

traj = fax.AVBDSolver(iterations=50).solve(
    metric, jnp.array([0., 0.]), jnp.array([1., 1.]), n_steps=20)

print(metric.arc_length(traj.xs),        # 1.1890  downwind
      metric.arc_length(traj.xs[::-1]))  # 1.8484  upwind

The same straight path costs $1.189$ downwind and $1.848$ in reverse. Both the still-medium metric $h$ and the drift $W$ are callables, so they may depend on position or contain trainable parameters. The figure shows the local reason for the difference: the unit-cost indicatrix is translated by the drift.

Loading randers-indicatrix…

randers-indicatrix

Drag the wind. The arrows are a uniform drift field, and the set of directions reachable at unit cost is carried off the origin by exactly that vector.

What it computes

Metrics and domains: FinslerMetric takes any $F(x,v)$ and derives the rest. Randers, Riemannian, Euclidean in closed form; DiscreteRanders on graphs. Over EuclideanSpace, Sphere, Torus, FlatTorus, Hyperboloid, Paraboloid, TriangularMesh.
Derived geometry: Fundamental tensor, geodesic spray, nonlinear-connection and Berwald coefficients, the curvature tensor and flag-curvature samples. They are obtained by differentiating $E = \tfrac{1}{2}F^2$ , without a separate implementation of Christoffel symbols for each metric family.
Geodesics: ExponentialMap for initial-value problems. AVBDSolver and GaussNewtonGeodesic solve discretised boundary-value problems by local relaxation and a global block-tridiagonal Newton step, respectively; solve_continuation supplies warm starts for stiff problems.
Arrival time: EikonalSolver on two-dimensional grids, MeshEikonalSolver on triangulations, and VolumetricEikonalSolver on three-dimensional grids. Each accepts Randers drift and differentiates the converged fixed point with memory independent of the number of sweeps.
Learned metrics: Metrics are Equinox modules, so their array-valued parameters participate in JAX transformations. The library includes neural Riemannian and Randers metrics, decoder pullbacks for latent spaces, covariate-conditioned metrics, kernel-smoothed wind fields, and geometry-aware training losses.

The concept pages introduce these objects in dependency order and pair each one with a figure computed by the pinned finslerax release.

Where it applies

Navigation — a craft of fixed own-speed in a current, which is Zermelo's 1931 problem
Spreading fronts — fire, reaction and contagion, where wind displaces the arrival-time equation
Robotics — configuration-space motion with state- and direction-dependent travel costs
Representation learning — a decoder pullback, asymmetric where transitions do not reverse

Related software

Geomstats, pymanopt and geoopt cover Riemannian manifolds and optimisation on them, where the metric is symmetric by construction. Mirebeau’s AGD/HFM is a reference implementation for anisotropic and Randers eikonal methods. The forward sweep and implicit adjoint in finslerax are based on Gahtan, Shpund and Bronstein’s differentiable Randers–Finsler solver. The JAX implementation and its integration with the geometry library are here. The Diffrax, Equinox and Optax are the JAX stack finslerax is built on. What finslerax adds is the apparatus around an asymmetric metric — spray, geodesics, transport, curvature, arrival time — with the metric itself trainable.