Code Structure¶
From the theory, the discretised version of the linear neutron transport equation is
where the bold-faced letters indicate linear operators: \(\mathbf{I}\) is the identity operator, \(\mathbf{M}\) the operator converting values defined on the discrete angles to values on the unit sphere (moments), \(\mathbf{D}\) the operator doing the reverse, \(\mathbf{T}^{-1}\) the inverted transport operator (using an efficient sweep algorithm such that it is as expensive as the normal forward operator), \(\mathbf{H}\) the finite element mass operator, and \(\mathbf{S}\) the scattering operator. For only using the code, their precise definition is not important.
This linear equation is solved by using a iterative solver which only requires applying the operators to an iteration vector \(\phi_{g\ell ei}^{(q)}\), we don’t need to have a matrix representation of the full system. For differentiation, we should be able to solve the transpose system.
Below we will describe the code structure in a simplified way.
Energy Set¶
Energy sets don’t actually matter for the discretisation; they are implicit in the cross sections. This submodule just has some common energy sets and a way to collapse group fluxes to a different energy set.
Quadrature Set¶
This submodule defines the angular quadrature sets and provides functions for computing the required data for the moment and discrete operators \(\mathbf{M}, \mathbf{D}\). It also provides a way to reduce the dimensionality of a 3D quadrature set for lower dimensional problems.
GeometricDomain and Domain¶
A GeometricDomain defines how the domain is geometrically constructed: its vertices and the connection between
each vertices. Can compute the required discontinuous-Galerkin connected to matrices (which face connects to which other face or element)
and Jacobian matrices. Also handles interpolation location (using BVH trees).
A Domain also adds cross section information: which element has which cross section in CrossSectionsIndexed.
SolutionDomain¶
SolutionDomain defines which type of Element is used on a Domain (e.g. first order / second order) and
has functions for computing the relevant finite element matrices. Also checks compatibility.
Operators and Strategies¶
The multigroup system is essentially an operator equation, which may consist of sub-operators. This module defines all operators and sub-operators needed for constructing it.
An operator is defined solely by its meaning (what does it calculate). To reflect this (see developer guide for a more indepth argument), each operator is implemented as follows:
class Operator(BaseLinearOperator)
operator_data : OperatorData
strategy : Strategy
def mv(self, x ):
return strategy.forward(self.operator_data, x)
where the strategy only defines the numerics (how it calculates). Constructing the multigroup operator thus requires picking
a implementation of MultiGroupStrategy and constructing the operator data (which defines the mesh, and other discretisation parameters).
Then, it can be passed to a Lineax solver, which automatically handles the differentiation as well.
Strategies can be composed: in the multigroup system, we can either compute the entire action \(\mathbf{A}\) using a single strategy (i.e. writing the whole kernel in one go), or we can decompose it into each sub-operator. This looks schematically as
class FusedMultiGroupStrategy(MultiGroupStrategy):
def forward(data, x):
return compute_fused_kernel(data, x)
class MultiGroupChainGatherStrategy(StrategyA):
strategy_m : StrategyM
strategy_tinv : StrategyTinv
strategy_h : StrategyH
strategy_d : StrategyD
strategy_s : StrategyS
def forward(data, x):
return x - strategy_m.forward(data,
strategy_tinv.forward(data,
strategy_h.forward(data,
strategy_d.forward(data,
strategy_s.forward(data, x) ) ) ) )
In fact, since each operator can also be applied either for a single energy group (instead of multiple groups as above),
this can even be further decomposed into single-group strategies, and a group lifting (how to parallelize / serialize).
Even more, each angle can be further decomposed for a single group (do you fuse in a CUDA kernel or just `jax.vmap`?).
We want to remain flexible in these decisions since we don’t know what is the best a priori for a given problem.
Solvers¶
This submodule defines some useful solvers, especially for fusion problems with only downscatter where a lower triangular block solver can be used. This saves a lot of memory since the iteration vector is much smaller than in the full system.
Parametrisation¶
We are interested in parametrising arbitrary quantities. This is supported, but you do need to recalculate various quantities if you update
the vertices (e.g. the Jacobian matrices or the sweep ordering for \(T^{-1}\)). The builders in operator_creation do exactly that:
each operator is declared as a graph of derivations over a set of primitives, so updating a primitive recomputes everything downstream of it
and nothing else. The mechanism itself lives in jax_sn.tools.parametrisation.
Variance Reduction¶
Running stochastic simulations can be very expensive; this submodule provides a way to construct weight windows from deterministic simulations using the continuous adjoint FW-CADIS method.