Eikonax

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 makes 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

The JAX in Eikonax

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.

Installation and Development

Eikonax is deployed as a python package, simply install via

pip install eikonax

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

Usage

Under Usage, we provide walkthroughs of the functionalities of Eikonax. The Forward Solver tutorial explains in detail how to set up Eikonax for solving the Eikonal equation. Parametric Derivatives demonstrates how to differentiate the solver, given a computed forward solution.

API Reference

The API reference contains detailed explanations of all software components of Eikonax, and how to use them.

Examples

We provide runnable examples in our Github repository.

Acknowledgement and License

Eikonax is being developed in the research group Uncertainty Quantification at KIT. It is partially based on the excellent FIM-Python tool. Eikonax is distributed as free software under the MIT License.