Spatial Discretisation¶

In this notebook, the spatial discretisation and the process for constructing it in the code will be shown.

We recap the total operator equation:

\[ (\mathbf{I} - \mathbf{M}\mathbf{T}^{-1}\mathbf{H}\mathbf{D}\mathbf{S}) [\phi_{g\ell\epsilon i}] = \mathbf{M}\mathbf{T}^{-1}\mathbf{H}[q_{gk\epsilon i}] \]

with the following operator definitions:

\[\begin{split} \begin{aligned} \mathbf{M}[a_{\dots k \dots}] &= \sum_k w_k Y_{\ell k} a_{\dots k \dots}, \\[0.5em] \mathbf{D}[b_{\dots \ell \dots}] &= \sum_\ell Y_{\ell k} b_{\dots \ell \dots}, \\[0.5em] \mathbf{H}[c_{\dots \epsilon i}] &= \sum_i c_{\dots \epsilon i} \int_{\epsilon} H_{\epsilon j} H_{\epsilon i}\, dV, \\[0.5em] \mathbf{S}[d_{\tilde g \ell \epsilon \dots}] &= \sum_{\tilde g} \sigma^s_{\ell g \tilde g \epsilon} d_{\tilde g \ell \epsilon \dots}, \\[0.5em] \mathbf{T}[\psi_{gk\epsilon i}] &= \mathbf{G}_{gk\epsilon}[\psi_{gk\epsilon i}] - \sum_{\beta\in\beta^-_{\epsilon k}} \mathbf{F}_{k\epsilon,\epsilon(\beta)} [\psi_{gk\epsilon(\beta)i}] \\ &= \sum_i \psi_{gk\epsilon i} \left[ \int_{\epsilon} H_{\epsilon i} \left( \sigma_g H_{\epsilon j} - \Omega_k\cdot\nabla H_{\epsilon j} \right) \,dV + \sum_{\beta\in\beta^+_{\epsilon k}} |\Omega_k\cdot\mathbf{n}_\beta| \int_{\beta} H_{\epsilon i}H_{\epsilon j}\,dA \right] \\ &\quad - \sum_{\beta\in\beta^-_{\epsilon k}} \sum_i \psi_{gk\epsilon(\beta)i} |\Omega_k\cdot\mathbf{n}_\beta| \int_{\beta} H_{\epsilon(\beta)i} H_{\epsilon j}\,dA. \end{aligned} \end{split}\]

This assumed a particular spatial discretisation with elements \(\epsilon\) with a number of faces \(\beta_\epsilon\) and basis functions \(H_{\epsilon i}\) defined on each element (local to each element).

In particular, the sweep algorithm requires the faces to be unambiguously upwind or downwind (\(\Omega_k \cdot \mathbf{n}_\beta\) positive or negative) and the element to be convex. Here, we show the currently standard Basix Lagrangian simplex implementation, although different kinds of elements could be implemented as long as they satisfy the above assumptions.

Note that of these operators, only the \(\mathbf{T}\) and \(\mathbf{H}\) operators depend on the spatial discretisation: they are the only ones involving basis functions \(H_{\epsilon i}\) and integrals over the elements \(\epsilon\) and faces \(\beta\). The other operators are purely algebraic and do not depend on the spatial discretisation, so we’ll focus our attention on these.

Simplex Elements¶

Lagrangian elements mean that there are a set of nodes \(\mathbf{x}_{\epsilon i}\) associated with each element \(\epsilon\), and the basis functions \(H_{\epsilon i}\) are defined such that they are equal to 1 at their associated node and 0 at all other nodes. The basis functions are polynomials of a given order \(p\). Let’s plot them for both 1D and 2D.

import jax
jax.config.update("jax_enable_x64", True)
jax.config.update("jax_platforms",   "cpu")
from jax_sn.solution_domain import BasixLagrangianSimplex
from jax_sn.tools.plotly import plot_element_basis_functions
plot_element_basis_functions(element_1 := BasixLagrangianSimplex(degree = 4, dimension = 1))

plot_element_basis_functions(element_2 := BasixLagrangianSimplex(degree = 3, dimension = 2), mode='all', n_subdivisions=50)

These basis functions don’t actually show up anywhere, except in the integrals. For simplex elements, the integrals are even easier: they have a linear Jacobian (mapping a real element to the reference element), so the integrals can be pre-computed and just multiplied by the Jacobian, as explained here:

\[ \int H_{\epsilon i} H_{\epsilon j}\, dV = J_\epsilon \int H_{i} H_{j}\, dV. \]

print("Do the elements have a linear Jacobian?")
print(f"Element 1: {element_1.has_linear_jacobian}")
print(f"Element 2: {element_2.has_linear_jacobian}")

Do the elements have a linear Jacobian?
Element 1: True
Element 2: True

Geometric Domain¶

These integrals therefore require some precomputed data. Since this depends on the particular type of element, each element has a way to calculate it from a GeometricDomain. Of course, they have to be compatible: you cannot compute it for a triangle using a tetrahedral domain. A GeometricDomain only stores adjacency information, cross sections, and nodes / connectivity — it doesn’t know about which type of element it is.

A detail is that a domain requires a face template: which nodes form which face. You can input it normally, but this has to be compatible with the face template of the finite element later. Commonly, we just create them both at the same time to avoid mismatching.

The face template doesn’t depend on the degree of the finite element!

We create a simple 2D domain:

import jax.numpy as jnp
vertices = jnp.array([
    [0.0, 0.0],  # 0 \
    [1.0, 0.0],  # 1  | box corners
    [1.0, 1.0],  # 2  |
    [0.0, 1.0],  # 3 /
    [2.0, 0.0],  # 4 \ top triangle
    [0.0, 2.0],  # 5 \ right triangle
])

connectivity = jnp.array([
    [0, 1, 2],  # box triangle A
    [0, 2, 3],  # box triangle B
    [1, 2 ,4],  # top triangle
    [2, 3, 5],  # right triangle
])

solution_element = BasixLagrangianSimplex(degree=2, dimension=2)

print("Face template:", solution_element.face_template)

from jax_sn.domain import GeometricDomain
geometric_domain = GeometricDomain.from_mesh((vertices, connectivity), solution_element.face_template)

Face template: [[1 2]
 [0 2]
 [0 1]]

and plot it:

from jax_sn.tools.plotly import plot_domain
plot_domain(geometric_domain, color_by_material= False, show_element_numbers = True)

The geometric domain has the following data structure

print(geometric_domain)

GeometricDomain(
  vertices=f64[6,2],
  connectivity=i64[4,3],
  face_normals=f64[4,3,3],
  etoe=weak_i64[4,3],
  etof=weak_i64[4,3],
  face_perm_idx=i64[4,3],
  face_template=i64[3,2]
)

face_normals are just the normals of each face: 4 elements, 3 faces per element, 3 dimensional. (minus one means not connected)
etoe is which element is connected to which element through its faces
etof is which face number each element is connected to (of the other element)
face_perm_idx is the permutation of the face nodes to match the other element’s face nodes.

These are calculated automatically by the above constructor.

print(geometric_domain.etoe)

[[ 2  1 -1]
 [ 3 -1  0]
 [-1 -1  0]
 [-1 -1  1]]

Here we see element 0 is connected to element 2 through face 0 and element 1 through face 1, with no elements connected through face 2.

Precomputing integral quantities¶

Using this geometric domain, we can precompute all geometric Jacobians (including edge Jacobians, inverse Jacobians, and determinants):

solution_element.precompute_geometric_data(geometric_domain)

PrecomputedGeometricData(
  vol_jacobians=f64[4,2,2],
  vol_jacobians_inv=f64[4,2,2],
  vol_jacobians_det=f64[4],
  face_jacobians_det=f64[4,3]
)

For the computation of the integrals, besides the Jacobian, we also need the integrals over the reference element, e.g.

\[ \int H_i H_j\, dV \]

over the reference element (as well as the face / derivative versions).

These are stored in a FEMMatrices and can be computed for a particular element easily:

from jax_sn.solution_domain import FEMMatrices

fem_matrices = FEMMatrices.from_element(solution_element)
print(fem_matrices)

FEMMatrices(
  mass_matrix=f64[6,6],
  integration_coeff=f64[6],
  derivative_matrix=f64[2,6,6],
  outflow_matrices=f64[3,6,6],
  inflow_matrices=f64[3,2,3,6,6]
)

All matrices have a size of 6 because a 2nd degree triangle has 6 basis functions.

mass_matrix is the \(\int H_i H_j\, dV\) integral over the reference element
integration_coeff is the \(\int H_i \, dV\) integral over the reference element
derivative_matrix are the \(\int H_i \frac{\partial H_j}{\partial q_k} \, dV\) matrices (2 reference coordinate dimension)
outflow_matrices are the \(\int_{\beta} H_i H_j\, dA\) matrices for each face \(\beta\) of the reference element (3 faces for a triangle)
inflow_matrices are the \(\int_{\beta} H^p_i H_j\, dA\) matrices for each face in the connected element \(p\) (3 faces), each permutation of the face nodes (2), and each face of the reference element (3)

Domain and cross sections¶

To actually solve the equations, we need the cross sections in there as well (\(\sigma_g \epsilon \) and \(\sigma^s_{\ell g \tilde g \epsilon}\)). Therefore, to create the container for all precomputed data, the SolutionDomain, we need to construct the domain including cross section information.

Let’s first create some cross sections:

from jax_sn.domain import Domain, CrossSectionData, CrossSectionsIndexed

water = CrossSectionData(
    total = jnp.array([15.2090, 24.4140]),
    scattering = jnp.array([
        [
            [5.5069, 0.0],
            [6.6227, 11.4080]
        ],
        [
            [1.158567, 0.0],
            [1.1234767, 2.4733]
        ]
    ])
)
iron = CrossSectionData(
    total = jnp.array([30.3110, 26.8760]),
    scattering = jnp.array([
        [
            [18.2500, 0.0],
            [6.45294, 21.3570]
        ],
        [
            [4.36, 0.0],
            [-0.003508, 2.3457]
        ]        
    ])
)
void = CrossSectionData(
    total = jnp.array([1e-10, 1e-10]),
    scattering = jnp.zeros((2, 2, 2))
)

Now we point each element to a specific material using indexing:

water_idx = 0 
iron_idx  = 1
void_idx  = 2
material_idx = jnp.array([water_idx, water_idx, iron_idx, void_idx])
# Create a batched version of the data:
total_xs               = CrossSectionData.from_list([water, iron, void])
cross_sections_indexed = CrossSectionsIndexed(material_idx, total_xs)

print(cross_sections_indexed)

domain = Domain.from_geometric_domain_and_cross_sections(geometric_domain, cross_sections_indexed)

CrossSectionsIndexed(
  indices=i64[4],
  cross_section_data=CrossSectionData(total=f64[3,2], scattering=f64[3,2,2,2])
)

and plot the domain again, colorbing by material:

plot_domain(domain, material_names = ["Water", "Iron", "Void"], show_element_numbers = True)

Solution domain and spatial data¶

We can now actually create the solution domain:

from jax_sn.solution_domain import SolutionDomain
solution_domain = SolutionDomain.from_element_and_domain(solution_element, domain)
print(solution_domain)

SolutionDomain(
  spatial_data=SpatialData(
    domain=Domain(
      vertices=f64[6,2],
      connectivity=i64[4,3],
      face_normals=f64[4,3,3],
      etoe=weak_i64[4,3],
      etof=weak_i64[4,3],
      face_perm_idx=i64[4,3],
      face_template=i64[3,2],
      cross_sections_indexed=CrossSectionsIndexed(
        indices=i64[4],
        cross_section_data=CrossSectionData(
          total=f64[3,2], scattering=f64[3,2,2,2]
        )
      )
    ),
    fem_matrices=FEMMatrices(
      mass_matrix=f64[6,6],
      integration_coeff=f64[6],
      derivative_matrix=f64[2,6,6],
      outflow_matrices=f64[3,6,6],
      inflow_matrices=f64[3,2,3,6,6]
    ),
    precomputed_geometric_data=PrecomputedGeometricData(
      vol_jacobians=f64[4,2,2],
      vol_jacobians_inv=f64[4,2,2],
      vol_jacobians_det=f64[4],
      face_jacobians_det=f64[4,3]
    ),
    element=<jax_sn.solution_domain.basix_fem.BasixLagrangianSimplex object at 0x7fae167e5dc0>
  )
)

It’s simply a container for all the precomputed data, the geometric domain, and the element type.

One thing to note is that if you modify the vertices (e.g. using a eqx.tree_at), this might make the precomputed data inconsistent!

See the … for how to cleanly update the vertices instead.

Mass Operator¶

The mass operator \(\mathbf{H}\)

\[\begin{split}\mathbf{H}[c_{\dots \epsilon i}] = \sum_i c_{\dots \epsilon i} \int_{\epsilon} H_{\epsilon j} H_{\epsilon i}\, dV= \sum_i c_{\dots \epsilon i}\mathcal{J}_\epsilon \int_{\epsilon} H_{j} H_{ i}\, dV\\[0.5em]\end{split}\]

is quite simple; it’s just a batched matrix product and requires the FEM matrices and the precomputed data.

It needs additional shape information to have proper checking of input/output arguments (e.g. group and moment indices)

from jax_sn.operators.mass_operator import MassOperator , LaxMass, MassData

mass_data = MassData(solution_domain.fem_matrices, solution_domain.precomputed_geometric_data)
mass_operator = MassOperator(mass_data, additional_shape = (), strategy = LaxMass())

print(mass_operator)

MassOperator(
  data=MassData(
    fem_matrices=FEMMatrices(
      mass_matrix=f64[6,6],
      integration_coeff=f64[6],
      derivative_matrix=f64[2,6,6],
      outflow_matrices=f64[3,6,6],
      inflow_matrices=f64[3,2,3,6,6]
    ),
    precomputed_geometric_data=PrecomputedGeometricData(
      vol_jacobians=f64[4,2,2],
      vol_jacobians_inv=f64[4,2,2],
      vol_jacobians_det=f64[4],
      face_jacobians_det=f64[4,3]
    )
  ),
  additional_shape=(),
  strategy=LaxMass()
)

In practice you’d never construct it this way; it’s a part of a composed operator and there are convenience functions.

Connecting spatial discretisation to other operators¶

Every other operator also ‘sees’ the spatial discretisation, but in a simpler way: they are all local operators that just operator on each element and value separately.

In other words, they are just applied to each element and value directly, without considering any integrals or basis functions. They don’t see the spatial discretisation except for the additional dimensions they operate on.