Core Functions

corefunctions

Core functions for Eikonax forward solves and parametric derivatives.

This module contains atomic functions that make up the Eikonax solver routines. They (and their automatic derivatives computed with JAX) are further used to evaluate parametric derivatives.

Classes:

Name	Description
MeshData	Data characterizing a computational mesh from a vertex-centered perspective.
InitialSites	Initial site info.

Functions:

Name	Description
compute_softminmax	Smooth double ReLU-type approximation that restricts a variable to the interval [0, 1].
compute_edges	Compute the edges of a triangle from vertex indices and coordinates.
compute_optimal_update_parameters_soft	Compute position parameter for update of a node within a specific triangle.
compute_optimal_update_parameters_hard	Compute position parameter for update of a node within a specific triangle.
_compute_optimal_update_parameters	Compute the optimal update parameter for the solution of the Eikonal equation.
compute_fixed_update	Compute update for a given vertex, triangle, and update parameter.
compute_update_candidates_from_adjacent_simplex	Compute all possible update candidates from an adjacent triangle.
compute_vertex_update_candidates	Compute all update candidates for a given vertex.
grad_average	JAX-compatible computation of the gradient of the average function.

eikonax.corefunctions.grad_update_solution module-attribute

grad_update_solution = jax.grad(compute_fixed_update, argnums=0)

eikonax.corefunctions.grad_update_parameter module-attribute

grad_update_parameter = jax.grad(compute_fixed_update, argnums=1)

eikonax.corefunctions.grad_update_lambda module-attribute

grad_update_lambda = jax.grad(compute_fixed_update, argnums=2)

eikonax.corefunctions.jac_lambda_soft_solution module-attribute

jac_lambda_soft_solution = jax.jacobian(compute_optimal_update_parameters_soft, argnums=0)

eikonax.corefunctions.jac_lambda_hard_solution module-attribute

jac_lambda_hard_solution = jax.jacobian(compute_optimal_update_parameters_hard, argnums=0)

eikonax.corefunctions.jac_lambda_soft_parameter module-attribute

jac_lambda_soft_parameter = jax.jacobian(compute_optimal_update_parameters_soft, argnums=1)

eikonax.corefunctions.jac_lambda_hard_parameter module-attribute

jac_lambda_hard_parameter = jax.jacobian(compute_optimal_update_parameters_hard, argnums=1)

eikonax.corefunctions.MeshData dataclass

Data characterizing a computational mesh from a vertex-centered perspective.

Attributes:

Name	Type	Description
vertices	jax.Array | npt.NDArray	The coordinates of the vertices in the mesh. The dimension of this array is (num_vertices, dim), where num_vertices is the number of vertices in the mesh and dim is the dimension of the space in which the mesh is embedded.
adjacency_data	jax.Array | npt.NDArray	Adjacency data for each vertex. This is the list of adjacent triangles, together with the two vertices that span the respective triangle with the current vertex. The dimension of this array is (num_vertices, max_num_adjacent_simplices, 4), where max_num_adjacent_simplices is the maximum number of simplices that are adjacent to a vertex in the mesh. All entries have this maximum size, as JAX only operates on homogeneous data structures. If a vertex has fewer than max_num_adjacent_simplices adjacent simplices, the remaining entries are filled with -1.

__post_init__

__post_init__() -> None

Convert to jax arrays.

eikonax.corefunctions.InitialSites dataclass

Initial site info.

For a unique solution of the state-constrained Eikonal equation, the solution values need to be given a number of initial points (at least one). Multiple initial sites need to be compatible, in the sense that the arrival time from another source is not smaller than the initial value itself.

Attributes:

Name	Type	Description
inds	jax.Array | npt.NDArray	The indices of the nodes where the initial sites are placed.
values	jax.Array | npt.NDArray	The values of the initial sites.

__post_init__

__post_init__() -> None

Convert to jax arrays.

eikonax.corefunctions.compute_softminmax

compute_softminmax(value: jtReal[jax.Array, ...], order: int) -> jtFloat[jax.Array, ...]

Smooth double ReLU-type approximation that restricts a variable to the interval \([0, 1]\).

Unstable AD

While the method itself is stable, its derivative computed with JAX is not for higher orders \(\kappa\)

Given an input value \(x\) and an order parameter \(\kappa > 0\), the method performs a differentiable transformation according to

\[ \begin{gather*} f_{\text{lb}}(x) = \frac{1}{\kappa}\log\Big[ 1 + \exp\big( \kappa x \big) \Big], \\ f_{\text{ub}}(x) = 1 - \frac{1}{\kappa}\log\Big[ 1 + \exp\big( -\kappa(x-1) \big) \Big], \\ \tilde{\phi}(x) = f_{\text{ub}}(f_{\text{lb}}(x)), \\ \phi_{\text{lb}} = 1 - \frac{\log\big(1+\exp(\kappa)\big)}{\kappa} < 0, \\ \phi(x) = \frac{\tilde{\phi}(x) - \phi_{\text{lb}}}{1-\phi_{\text{lb}}}. \end{gather*} \]

The method is numerically stable, obeys the value range, and does not introduce any new extrema.

Parameters:

Name	Type	Description	Default
value	jax.Array	variable to restrict to range [0,1]	required
order	int	Approximation order of the smooth approximation	required

Returns:

Type	Description
jtFloat[jax.Array, ...]	jax.Array: Smoothed/restricted value

eikonax.corefunctions.compute_edges

compute_edges(i: jtInt[jax.Array, ''], j: jtInt[jax.Array, ''], k: jtInt[jax.Array, ''], vertices: jtFloat[jax.Array, 'num_vertices dim']) -> tuple[jtFloat[jax.Array, dim], jtFloat[jax.Array, dim], jtFloat[jax.Array, dim]]

Compute the edges of a triangle from vertex indices and coordinates.

Parameters:

Name	Type	Description	Default
i	jax.Array	First vertex index of a triangle	required
j	jax.Array	Second vertex index of a triangle	required
k	jax.Array	Third vertex index of a triangle	required
vertices	jax.Array	jax.Array of all vertex coordinates	required

Returns:

Type	Description
tuple[jtFloat[jax.Array, dim], jtFloat[jax.Array, dim], jtFloat[jax.Array, dim]]	tuple[jax.Array, jax.Array, jax.Array]: Triangle edge vectors

eikonax.corefunctions.compute_optimal_update_parameters_soft

compute_optimal_update_parameters_soft(solution_values: jtFloat[jax.Array, 2], parameter_tensor: jtFloat[jax.Array, 'dim dim'], edges: tuple[jtFloat[jax.Array, dim], jtFloat[jax.Array, dim], jtFloat[jax.Array, dim]], softminmax_order: int, softminmax_cutoff: Real) -> jtFloat[jax.Array, 4]

Compute position parameter for update of a node within a specific triangle.

For a given vertex \(i\) and adjacent triangle, we compute the update for the solution of the Eikonal as propagating from a point on the connecting edge of the opposite vertices \(j\) and \(k\). We thereby assume the solution value to vary linearly on that edge. The update parameter in \([0,1]\) indicates the optimal linear combination of the solution values at \(j\) and \(k\), in the sense that the solution value at \(i\) is minimized. As the underlying optimization problem is constrained, we compute the solutions of the unconstrained problem, as well as the boundary values. The former are constrained to the feasible region [0,1] by the soft minmax function implemented in compute_softminmax. We further clip values lying to far outside the feasible region, by masking them with value -1. This function is a wrapper, for the unconstrained solution values, it calls the implementation function _compute_optimal_update_parameters.

Parameters:

Name	Type	Description	Default
solution_values	jax.Array	Current solution values, as per the previous iteration	required
parameter_tensor	jax.Array	Parameter tensor for the current triangle	required
edges	tuple[jax.Array, jax.Array, jax.Array]	Edge vectors of the triangle under consideration	required
softminmax_order	int	Order of the soft minmax function to be employed	required
softminmax_cutoff	int	Cutoff value beyond parameter values are considered infeasible and masked with -1	required

Returns:

Type	Description
jtFloat[jax.Array, 4]	jax.Array: All possible candidates for the update parameter, according to the solution of the constrained optimization problem

eikonax.corefunctions.compute_optimal_update_parameters_hard

compute_optimal_update_parameters_hard(solution_values: jtFloat[jax.Array, 2], parameter_tensor: jtFloat[jax.Array, 'dim dim'], edges: tuple[jtFloat[jax.Array, dim], jtFloat[jax.Array, dim], jtFloat[jax.Array, dim]]) -> jtFloat[jax.Array, 4]

Compute position parameter for update of a node within a specific triangle.

For a given vertex \(i\) and adjacent triangle, we compute the update for the solution of the Eikonal as propagating from a point on the connecting edge of the opposite vertices \(j\) and \(k\). We thereby assume the solution value to vary linearly on that dge. The update parameter in \([0,1]\) indicates the optimal linear combination of the solution values at \(j\) and \(k\), in the sense that the solution value at \(i\) is minimized. As the underlying optimization problem is constrained, we compute the solutions of the unconstrained problem, as well as the boundary values. The former are constrained to the feasible region \([0,1]\) by a simple cutoff. This function is a wrapper, for the unconstrained solution values, it calls the implementation function _compute_optimal_update_parameters.

Parameters:

Name	Type	Description	Default
solution_values	jax.Array	Current solution values, as per the previous iteration	required
parameter_tensor	jax.Array	Parameter tensor for the current triangle	required
edges	tuple[jax.Array, jax.Array, jax.Array]	Edge vectors of the triangle under consideration	required

Returns:

Type	Description
jtFloat[jax.Array, 4]	jax.Array: All possible candidates for the update parameter, according to the solution of the constrained optimization problem

eikonax.corefunctions._compute_optimal_update_parameters

_compute_optimal_update_parameters(solution_values: jtFloat[jax.Array, 2], parameter_tensor: jtFloat[jax.Array, 'dim dim'], edges: tuple[jtFloat[jax.Array, dim], jtFloat[jax.Array, dim], jtFloat[jax.Array, dim]]) -> tuple[jtFloat[jax.Array, ''], jtFloat[jax.Array, '']]

Compute the optimal update parameter for the solution of the Eikonal equation.

The function works for the update within a given triangle. The solutions of the unconstrained minimization problem are given as the roots of a quadratic polynomial. For the local metric tensor \(M_s\) for the given simplex, the set of known solution values \(\mathbf{U}_{jk}\), and the simplex edges \(\mathbf{E}_{ijk}\), we have that

\[ \lambda_{1/2}(\mathbf{M}_{s},\mathbf{U}_{jk},\mathbf{E}_{ijk}) = \frac{(\mathbf{e}_{jk},\mathbf{M}_{s}^{-1}\mathbf{e}_{ki}) \pm c(\mathbf{M}_{s},\mathbf{U}_{jk},\mathbf{E}_{ijk})}{(\mathbf{e}_{jk}, \mathbf{M}_{s}^{-1}\mathbf{e}_{jk})}, \]

with

\[ c(\mathbf{M}_{s},\mathbf{U}_{jk},\mathbf{E}_{ijk}) = (u_k - u_j) \sqrt{\frac{(\mathbf{e}_{jk},\mathbf{M}_{s}^{-1}\mathbf{e}_{jk}) (\mathbf{e}_{ki},\mathbf{M}_{s}^{-1}\mathbf{e}_{ki}) - (\mathbf{e}_{jk},\mathbf{M}_{s}^{-1}\mathbf{e}_{ki})^2} {(\mathbf{e}_{jk},\mathbf{M}_{s}^{-1}\mathbf{e}_{jk}) - (u_k - u_j)^2}}. \]

These roots may or may not lie inside the feasible region \([0,1]\).The function returns both solutions, which are then further processed in the calling wrapper.

eikonax.corefunctions.compute_fixed_update

compute_fixed_update(solution_values: jtFloat[jax.Array, 2], parameter_tensor: jtFloat[jax.Array, 'dim dim'], lambda_value: jtFloat[jax.Array, ''], edges: tuple[jtFloat[jax.Array, dim], jtFloat[jax.Array, dim], jtFloat[jax.Array, dim]]) -> jtFloat[jax.Array, '']

Compute update for a given vertex, triangle, and update parameter.

The update value is given by the solution at a point on the edge between the opposite vertices, plus the travel time from that point to the vertices under consideration. For a local metric tensor \(M_s\) for the given simplex, the set of known solution values \(\mathbf{U}_{jk}\), the simplex edges \(\mathbf{E}_{ijk}\), and an optimal update parameter \(\lambda\), it reads

\[ G_{\text{fix}}(\lambda,\mathbf{M}_s,\mathbf{U}_{jk}, \mathbf{E}_{ijk}) = u_j + \lambda (u_k-u_j) + \sqrt{\big((\mathbf{e}_{ji} - \lambda\mathbf{e}_{jk}), \mathbf{M}_{s}^{-1}(\mathbf{e}_{ji} - \lambda\mathbf{e}_{jk})\big)} \]

Parameters:

Name	Type	Description	Default
solution_values	jax.Array	Current solution values at opposite vertices j and k, as per the previous iteration	required
parameter_tensor	jax.Array	Conductivity tensor for the current triangle	required
lambda_value	jax.Array	Optimal update parameter	required
edges	tuple[jax.Array, jax.Array, jax.Array]	Edge vectors of the triangle under consideration	required

Returns:

Type	Description
jtFloat[jax.Array, '']	jax.Array: Updated solution value for the vertex under consideration

eikonax.corefunctions.compute_update_candidates_from_adjacent_simplex

compute_update_candidates_from_adjacent_simplex(old_solution_vector: jtFloat[jax.Array, num_vertices], tensor_field: jtFloat[jax.Array, 'num_simplices dim dim'], adjacency_data: jtInt[jax.Array, max_num_adjacent_simplices], vertices: jtFloat[jax.Array, 'num_vertices dim'], use_soft_update: bool, softminmax_order: int | None, softminmax_cutoff: Real | None) -> tuple[jtFloat[jax.Array, 4], jtFloat[jax.Array, 4]]

Compute all possible update candidates from an adjacent triangle.

Given a vertex and an adjacent triangle, this method computes all optimal update parameter candidates and the corresponding update values. To obey JAX's homogeneous array requirement, update values are also computed for infeasible parameter values, and have to be masked in the calling routine. This methods basically collects all results from the compute_optimal_update_parameters and compute_fixed_update functions.

Parameters:

Name	Type	Description	Default
old_solution_vector	jax.Array	Given solution vector, as per a previous iteration	required
tensor_field	jax.Array	Array of all tensor fields	required
adjacency_data	jax.Array	Index of one adjaccent triangle and corresponding vertices	required
vertices	jax.Array	Array of all vertex coordinates	required
use_soft_update	bool	Flag to indicate whether to use a soft update or a hard update	required
softminmax_order	int | None	Order of the soft minmax function for the update parameter, see compute_softminmax. Only required for use_soft_update=True	required
softminmax_cutoff	Real | None	Cutoff value for the soft minmax computation, see compute_optimal_update_parameters_soft. Only required for use_soft_update=True	required

Returns:

Type	Description
tuple[jtFloat[jax.Array, 4], jtFloat[jax.Array, 4]]	tuple[jax.Array, jax.Array]: Update values and parameter candidates from the given triangle

eikonax.corefunctions.compute_vertex_update_candidates

compute_vertex_update_candidates(old_solution_vector: jtFloat[jax.Array, num_vertices], tensor_field: jtFloat[jax.Array, 'num_simplices dim dim'], adjacency_data: jtInt[jax.Array, 'max_num_adjacent_simplices 4'], vertices: jtFloat[jax.Array, 'num_vertices dim'], use_soft_update: bool, softminmax_order: int, softminmax_cutoff: Real) -> jtFloat[jax.Array, 'max_num_adjacent_simplices 4']

Compute all update candidates for a given vertex.

This method combines all updates from adjacent triangles to a given vertex, as computed in the function compute_update_candidates_from_adjacent_simplex. Infeasible candidates are masked with jnp.inf.

Parameters:

Name	Type	Description	Default
old_solution_vector	jax.Array	Given solution vector, as per a previous iteration	required
tensor_field	jax.Array	jax.Array of all tensor fields	required
adjacency_data	jax.Array	Data of all adjacent triangles and corresponding vertices	required
vertices	jax.Array	jax.Array of all vertex coordinates	required
use_soft_update	bool	Flag to indicate whether to use a soft update or a hard update	required
softminmax_order	int | None	Order of the soft minmax function for the update parameter, see compute_softminmax. Only required for use_soft_update=True	required
softminmax_cutoff	Real | None	Cutoff value for the soft minmax computation, see compute_optimal_update_parameters_soft. Only required for use_soft_update=True	required

Returns:

Type	Description
jtFloat[jax.Array, 'max_num_adjacent_simplices 4']	jax.Array: All possible update values for the given vertex, infeasible vertices are masked with jnp.inf

eikonax.corefunctions.grad_average

grad_average(args: jtFloat[jax.Array, num_args], min_arg: jtFloat[jax.Array, '']) -> jtFloat[jax.Array, num_args]

JAX-compatible computation of the gradient of the average function.

This function is applied to actual minimum values, meaning it does not have a purpose on the evaluation level. It renders the minimum computation differentiable, however. Specifically, consider the scenario where a vertex can be updated identically from two different directions. Clearly, the mapping from the parameter field to the solution vector is not continuously differentiable at this point. Instead, we can only formulate a possible set of sub-gradients. One strategy to cope with sub-gradients is to employ simple averaging over all possible candidates.