Parameter Tensor Field

tensorfield

Composable and differentiable parameter tensor fields.

This module provides ABCs and implementations for the creation of differentiable parameter fields used in Eikonax. Recall that for the Eikonax solver, and particularly parameteric derivatives, we require an input tensor field \(\mathbf{M}: \mathbb{R}^M \times \mathbb{N}_0 \to \mathbb{S}_+^{d\times d}\). This means that the tensor field is a mapping \(\mathbf{M}(\mathbf{m},s)\) that assigns, given a global parameter vector \(\mathbf{m}\), an s.p.d tensor to every simplex \(s\) in the mesh. To allow for sufficient flexibility in the choice of tensor field, we implement it as a composition of two main components.

1. AbstractVectorToSimplicesMap provides the interface for a mapping from the global parameter vector \(\mathbf{m}\) to the local parameter values \(\mathbf{m}_s\) required to assemble the tensor \(\mathbf{M}_s\) for simplex \(s\).
2. AbstractSimplexTensor provides the interface for the assembly of the local tensor \(\mathbf{M}_s\), given the local contributions \(\mathbf{m}_s\) and a simplex s.

Concrete implementations of both components are used to initialize the TensorField object, which vectorizes and differentiates them using JAX, to provide the mapping \(\mathbf{M}(\mathbf{m})\) and its Jacobian tensor \(\frac{d \mathbf{M}}{d \mathbf{m}}\).

Classes:

Name	Description
AbstractVectorToSimplicesMap	ABC interface contract for vector-to-simplices maps
LinearScalarMap	Simple one-to-one map from global to simplex parameters
AbstractSimplexTensor	ABC interface contract for assembly of the tensor field
LinearScalarSimplexTensor	SimplexTensor implementation relying on one parameter per simplex
InvLinearScalarSimplexTensor	SimplexTensor implementation relying on one parameter per simplex
TensorField	Tensor field component

eikonax.tensorfield.AbstractVectorToSimplicesMap

Bases: eqx.Module

ABC interface contract for vector-to-simplices maps.

Every component derived from this class needs to implement the map and derivative methods. The map method is responsible for returning the relevant parameters for a given simplex from the global parameter vector. The derivative method computes the derivative of the mapping with respect to the global parameters.

Methods:

Name	Description
map	Interface for vector-to-simplex mapping
derivative	Interface for vector-to-simplex mapping derivative

map abstractmethod

map(simplex_ind: jtInt[jax.Array, ''], parameters: jtReal[jax.Array, num_parameters]) -> jtReal[jax.Array, num_parameters_local]

Interface for vector-to-simplex mapping.

For the given simplex_ind, return those parameters from the global parameter vector that are relevant for the simplex. This methods need to be broadcastable over simplex_ind by JAX (with vmap).

Parameters:

Name	Type	Description	Default
simplex_ind	int	Index of the simplex under consideration	required
parameters	jax.Array	Global parameter vector	required

Raises:

Type	Description
NotImplementedError	ABC error indicating that the method needs to be implemented in subclasses

Returns:

Type	Description
jtReal[jax.Array, num_parameters_local]	jax.Array: Relevant parameters for the simplex

derivative abstractmethod

derivative(simplex_ind: jtInt[jax.Array, ''], parameters: jtReal[jax.Array, num_parameters]) -> tuple[jtReal[jax.Array, num_parameters_local], jtReal[jax.Array, num_parameters_local]]

Interface for vector-to-simplex mapping derivative.

The derivative of the simplex mapping for a given simplex_ind is a num_local_parameters \(\times\) num_parameters Jacobian matrix, where the only non-zero entries are in the columns corresponding to the values in the global parameter that are used for the current simplex. The corresponding entries are typically identical to one. This method needs to return the Jacobian matrix in sparse COO format, i.e. one vector comprising the column numbers, and one vector comprising the matrix entries.

Parameters:

Name	Type	Description	Default
simplex_ind	int	Index of the simplex under consideration	required
parameters	jax.Array	Global parameter vector	required

Raises:

Type	Description
NotImplementedError	ABC error indicating that the method needs to be implemented in subclasses

Returns:

Type	Description
tuple[jtReal[jax.Array, num_parameters_local], jtReal[jax.Array, num_parameters_local]]	tuple[jax.Array, jax.Array]: Jacobi matrix, expressed via relevant global indices and matrix entries

eikonax.tensorfield.LinearScalarMap

Bases: AbstractVectorToSimplicesMap

Simple one-to-one map from global to simplex parameters.

Every simplex takes exactly one parameter \(m_s\), which is sorted in the global parameter in the same order as the simplices, meaning that \(m_s = \mathbf{m}[s]\).

map

map(simplex_ind: jtInt[jax.Array, ''], parameters: jtReal[jax.Array, num_parameters]) -> jtReal[jax.Array, num_parameters_local]

Return relevant parameters for a given simplex.

Parameters:

Name	Type	Description	Default
simplex_ind	int	Index of the simplex under consideration	required
parameters	jax.Array	Global parameter vector	required

Returns:

Type	Description
jtReal[jax.Array, num_parameters_local]	jax.Array: relevant parameter (only one)

derivative

derivative(simplex_ind: jtInt[jax.Array, ''], _parameters: jtReal[jax.Array, num_parameters]) -> tuple[jtReal[jax.Array, num_parameters_local], jtReal[jax.Array, num_parameters_local]]

Return sparse representation of Jacobi matrix, in this case two arrays of size one.

Parameters:

Name	Type	Description	Default
simplex_ind	int	Index of the simplex under consideration	required
_parameters	jax.Array	Global parameter vector (not used)	required

Returns:

Type	Description
tuple[jtReal[jax.Array, num_parameters_local], jtReal[jax.Array, num_parameters_local]]	tuple[jax.Array, jax.Array]: Jacobi matrix, expressed via relevant global indices and matrix entries

eikonax.tensorfield.AbstractSimplexTensor

Bases: eqx.Module

ABC interface contract for assembly of the tensor field.

SimplexTensor components assemble the tensor field for a given simplex and a set of parameters for that simplex. The relevant parameters are provided by the VectorToSimplicesMap component from the global parameter vector.

Note

This class provides the metric tensor as used in the inner product for the update stencil of the eikonal equation. This is the inverse of the conductivity tensor, which is the actual tensor field in the eikonal equation.

Methods:

Name	Description
assemble	Assemble the tensor field for a given simplex and parameters
derivative	Parametric derivative of the assemble method

assemble abstractmethod

assemble(simplex_ind: jtInt[jax.Array, ''], parameters: jtFloat[jax.Array, num_parameters_local]) -> jtFloat[jax.Array, 'dim dim']

Assemble the tensor field for given simplex and parameters.

Given a parameter array of size \(m_s\), the method returns a tensor of size \(d\times d\). The method needs to be broadcastable over simplex_ind by JAX (with vmap).

Parameters:

Name	Type	Description	Default
simplex_ind	int	Index of the simplex under consideration	required
parameters	jax.Array	Parameters for the simplex	required

Raises:

Type	Description
NotImplementedError	ABC error indicating that the method needs to be implemented in subclasses

Returns:

Type	Description
jtFloat[jax.Array, 'dim dim']	jax.Array: Tensor field for the simplex under consideration

derivative abstractmethod

derivative(simplex_ind: jtInt[jax.Array, ''], parameters: jtFloat[jax.Array, num_parameters_local]) -> jtFloat[jax.Array, 'dim dim num_local_parameters']

Parametric derivative of the assemble method.

Given a parameter array of size \(m_s\), the methods returns a Jacobian tensor of size \(d\times d\times m_s\). The method needs to be broadcastable over simplex_ind by JAX (with vmap).

Parameters:

Name	Type	Description	Default
simplex_ind	int	Index of the simplex under consideration	required
parameters	jax.Array	Parameters for the simplex	required

Raises:

Type	Description
NotImplementedError	ABC error indicating that the method needs to be implemented in subclasses

Returns:

Type	Description
jtFloat[jax.Array, 'dim dim num_local_parameters']	jax.Array: Jacobian tensor for the simplex under consideration

eikonax.tensorfield.LinearScalarSimplexTensor

Bases: AbstractSimplexTensor

SimplexTensor implementation relying on one parameter per simplex.

Given a scalar parameter \(m_s\), the tensor field is assembled as \(m_s \cdot \mathbf{I}\), where \(\mathbf{I}\) is the identity matrix.

Methods:

Name	Description
assemble	Assemble the tensor field for a parameter vector
derivative	Parametric derivative of the assemble method

assemble

assemble(_simplex_ind: jtInt[jax.Array, ''], parameters: jtFloat[jax.Array, num_parameters_local]) -> jtFloat[jax.Array, 'dim dim']

Assemble tensor for given simplex.

the parameters argument is a scalar here, and _simplex_ind is not used.

Parameters:

Name	Type	Description	Default
_simplex_ind	int	Index of simplex under consideration (not used)	required
parameters	jax.Array	Parameter (scalar) for tensor assembly	required

Returns:

Type	Description
jtFloat[jax.Array, 'dim dim']	jax.Array: Tensor for the simplex

derivative

derivative(_simplex_ind: jtInt[jax.Array, ''], _parameters: jtFloat[jax.Array, num_parameters_local]) -> jtFloat[jax.Array, 'dim dim num_local_parameters']

Parametric derivative of the assemble method.

Parameters:

Name	Type	Description	Default
_simplex_ind	int	Index of simplex under consideration (not used)	required
_parameters	jax.Array	Parameter (scalar) for tensor assembly	required

Returns:

Type	Description
jtFloat[jax.Array, 'dim dim num_local_parameters']	jax.Array: Jacobian tensor for the simplex under consideration

eikonax.tensorfield.InvLinearScalarSimplexTensor

Bases: AbstractSimplexTensor

SimplexTensor implementation relying on one parameter per simplex.

Given a scalar parameter \(m_s\), the tensor field is assembled as \(\frac{1}{m_s} \cdot \mathbf{I}\), where \(\mathbf{I}\) is the identity matrix.

Methods:

Name	Description
assemble	Assemble the tensor field for a parameter vector
derivative	Parametric derivative of the assemble method

assemble

assemble(_simplex_ind: jtInt[jax.Array, ''], parameters: jtFloat[jax.Array, num_parameters_local]) -> jtFloat[jax.Array, 'dim dim']

Assemble tensor for given simplex.

The parameters argument is a scalar here, and _simplex_ind is not used.

Parameters:

Name	Type	Description	Default
_simplex_ind	int	Index of simplex under consideration (not used)	required
parameters	jax.Array	Parameter (scalar) for tensor assembly	required

Returns:

Type	Description
jtFloat[jax.Array, 'dim dim']	jax.Array: Tensor for the simplex

derivative

derivative(_simplex_ind: jtInt[jax.Array, ''], parameters: jtFloat[jax.Array, num_local_parameters]) -> jtFloat[jax.Array, 'dim dim num_local_parameters']

Parametric derivative of the assemble method.

Parameters:

Name	Type	Description	Default
_simplex_ind	int	Index of simplex under consideration (not used)	required
parameters	jax.Array	Parameter (scalar) for tensor assembly	required

Returns:

Type	Description
jtFloat[jax.Array, 'dim dim num_local_parameters']	jax.Array: Jacobian tensor for the simplex under consideration

eikonax.tensorfield.TensorField

Bases: eqx.Module

Tensor field component.

Tensor fields combine the functionality of vector-to-simplices maps and simplex tensors according to the composition over inheritance principle. They constitute the full mapping \(\mathbf{M}(\mathbf{m})\) from the global parameter vector to the tensor field over all mesh faces (simplices). In addition, they provide the parametric derivative \(\frac{d\mathbf{M}}{\mathbf{m}}\) of that mapping. Tensor fields are completely independent from the Eikonax solver and derivator, but the output of these two components can be used to compute the partial derivative \(\mathbf{G}_m = \frac{du}{d\mathbf{M}}\frac{d\mathbf{M}}{\mathbf{m}}\).

Methods:

Name	Description
assemble_field	Assemble the tensor field for the given parameter vector
assemble_jacobian	Assemble the parametric derivative \(\frac{d\mathbf{M}}{\mathbf{m}}\) of the tensor field

__init__

__init__(num_simplices: int, vector_to_simplices_map: AbstractVectorToSimplicesMap, simplex_tensor: AbstractSimplexTensor) -> None

Constructor.

Takes information about the mesh simplices, a vector-to-simplices map, and a simplex tensor map.

Parameters:

Name	Type	Description	Default
num_simplices	int	Number of simplices in the mesh	required
vector_to_simplices_map	AbstractVectorToSimplicesMap	Mapping from global to simplex parameters	required
simplex_tensor	AbstractSimplexTensor	Tensor field assembly for a given simplex	required

assemble_field

assemble_field(parameter_vector: jtFloat[jax.Array | npt.NDArray, num_parameters]) -> jtFloat[jax.Array, 'num_simplices dim dim']

Assemble global tensor field from global parameter vector.

This method simply chains calls to the vector-to-simplices map and the simplex tensor objects, vectorized over all simplices.

Parameters:

Name	Type	Description	Default
parameter_vector	jax.Array | npt.NDArray	Global parameter vector	required

Returns:

Type	Description
jtFloat[jax.Array, 'num_simplices dim dim']	jax.Array: Global tensor field

assemble_jacobian

assemble_jacobian(parameter_vector: jtFloat[jax.Array | npt.NDArray, num_parameters]) -> spa.COO

Assemble the Jacobian \(\frac{d\mathbf{M}}{d\mathbf{m}}\).

The assembly of the Jacobian matrix works via a local chaining of the derivative calls to the simplex tensor and the vector-to-simplices map. A sparse COO representation of the call output is then vectorized over all simplices with JAX's vmap.

Parameters:

Name	Type	Description	Default
parameter_vector	jax.Array | npt.NDArray	Global parameter vector	required

Returns:

Type	Description
spa.COO	spa.COO: Jacobian tensor in sparse COO format, has dimension \(N_S \times d \times d \times M\)

_assemble_jacobian_global

_assemble_jacobian_global(parameter_vector: jtFloat[jax.Array, num_parameters]) -> tuple[list[jtInt[jax.Array, num_matrix_entries]], jtFloat[jax.Array, num_matrix_entries]]

Intermediate call for vectorization with vmap in jit context.

_assemble_jacobian_local

_assemble_jacobian_local(simplex_ind: jtInt[jax.Array, ''], parameter_vector: jtFloat[jax.Array, num_parameters]) -> tuple[jtInt[jax.Array, num_matrix_entries], jtInt[jax.Array, num_matrix_entries], jtInt[jax.Array, num_matrix_entries], jtInt[jax.Array, num_matrix_entries], jtFloat[jax.Array, 'derivative_dim_1 derivative_dim_2']]

Assembly of sparse jacobian representation for single simplex.