Eikonax: A Fully Differentiable Solver for the Anisotropic Eikonal Equation

Eikonax is a pure-Python implementation of a solver for the anisotropic eikonal equation on triangulated meshes. In particular, it focuses on domains $\Omega$ either in 2D Euclidean space, or 2D manifolds in 3D Euclidean space. For a given, space-dependent parameter tensor field $\mathbf{M}$, and a set $\Gamma$ of initially active points, Eikonax computes the arrival times $u$ according to

$$ \begin{gather*} \sqrt{\big(\nabla u(\mathbf{x}),\mathbf{M}(\mathbf{x})\nabla u(\mathbf{x})\big)} = 1,\quad \mathbf{x}\in\Omega, \\ \nabla u(\mathbf{x}) \cdot \mathbf{n}(\mathbf{x}) \geq 0,\quad \mathbf{x}\in\partial\Omega, \\ u(\mathbf{x}_0) = u_0,\quad \mathbf{x}_0 \in \Gamma. \end{gather*} $$

The iterative solver is based on Godunov-type upwinding and employs global Jacobi updates, which can be efficiently ported to SIMD architectures. In addition, Eikonax implements an efficient algorithm for the evaluation of parametric derivatives, meaning the derivative of the solution vector with respect to the parameter tensor field, $\frac{du}{d\mathbf{M}}$. More precisely, we assume that the tensor field is parameterized through some vector $\mathbf{m}$, s.th. we compute $\frac{du}{d\mathbf{m}} = \frac{du}{d\mathbf{M}}\frac{d\mathbf{M}}{d\mathbf{m}}$. This make Eikonax particularly suitable for the inverse problem setting, where derivative information is typically indispensable for efficient solution procedures. Through exploitation of causality in the forward solution, Eikonax can compute these derivatives through discrete adjoints on timescales much smaller than those for the forward solve.

Key Features

• Supports anisotropic conductivity tensors
• Works on irregular meshes
• GPU offloading of performance-relevant computations
• Super fast derivatives through causality-informed adjoints

Eikonax is mainly based on the JAX software library. This allows for GPU offloading of relevant computations. In addition, Eikonax makes extensive use of JAX`s just-in-time compilation and automatic differentiation capabilities.

Getting Started

Eikonax is deployed as a python package, simply install via

pip install eikonax[examples]

For development, we recommend using the great uv project management tool, for which Eikonax provides a universal lock file. To set up a reproducible environment, run

uv sync --all-groups

in the project root directory.

Documentation

The documentation provides further information regarding usage, theoretical background, technical setup and API. Alternatively, you can check out the notebooks under examples