jax_sn.variance_reduction

fw_cadis_scalar_flux(operator: MultiGroupOperator, rhs: RightHandSide, solver: AbstractLinearSolver)[source]

Runs the FW-CADIS algorithm to compute an importance map for variance reduction.

First the forward problem is solved to obtain the scalar flux:

\[\phi = A^{-1} q\]

where \(A\) is the multi-group operator and \(q\) is the right-hand side.

Then, the importance map is computed as:

\[\phi^\dagger = (A^\dagger)^{-1} \left( \frac{1}{\phi} \right)\]

where \(A^\dagger\) is the continuous adjoint (see solve_continuous_adjoint()).

For stability purposes, the scalar flux is first averaged over each element (due to the conservation property of DG, this cannot be zero, while the value at a particular basis function node can be extremely small and blow up).

Parameters:

  • operator (MultiGroupOperator) – The multi-group operator for which to compute the importance map.

  • rhs (RightHandSide) – The right-hand side of the equation to solve, typically representing a response function.

  • solver (AbstractLinearSolver) – The (iterative) linear solver used for both the forward and adjoint solves.

Returns:

The adjoint scalar flux importance map, shape [n_groups, n_moments, n_elements, n_basis].

Return type:

Array

solve_continuous_adjoint(operator: MultiGroupOperator, rhs, solver: AbstractLinearSolver)[source]

Solve \(\mathbf{A}^\dagger \psi^\dagger = \mathbf{C}^\dagger[b]\) for the continuous adjoint of operator, where the source is conditioned by the adjoint chain \(\mathbf{C}^\dagger = \mathbf{M}\,(\mathbf{T}^\dagger)^{-1}\mathbf{H}\mathbf{D}\) (built automatically because the chain strategy runs on the adjoint-transformed data).

The adjoint operator is the ordinary forward MultiGroupOperator run on _continuous_adjoint_data(); it is built inside the jit so gradients w.r.t. cross sections / geometry still flow through the adjoint solve.

Parameters:

  • operator (MultiGroupOperator) – The forward multi-group operator whose continuous adjoint is solved.

  • rhs (RightHandSide) – The (physical) adjoint source.

  • solver (AbstractLinearSolver) – An iterative solver (GMRES / BiCGStab). The direct block solver is rejected: for downscattering the continuous adjoint is upper-triangular, which that solver does not handle.

Return type:

lx.LinearSolveOutput

Submodules