Multigroup solution

In this notebook, we will show how to solve a multigroup neutron transport problem, using both the total-group solver and the per-group solver.

We’ll re-use the same materials as before and create the domain from the following regions:

0   - 0.2 m: Water
0.2 - 0.3 m: Water with source
0.3 - 0.8 m: Water
0.8 - 1.0 m: Iron
1.0 - 1.1 m: Void
1.1 - 1.2 m: Water
1.2 - 1.3 m: Void

We just use a simple discretization with a set number of elements per 10 cm and build from there (we hide the details of creating the mesh and material indices)

Show code cell source

Hide code cell source

import jax
import jax.numpy as jnp
jax.config.update("jax_platforms", "cpu")
jax.config.update("jax_enable_x64", True)
import jax_sn
import equinox as eqx
from typing import List
from jax_sn.domain.cross_sections import 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))
)


def base_region_stats(n_elements_per_10_cm : int):
    base_region_boundaries  = jnp.array([0.0, 0.2, 0.3, 0.8, 1.0, 1.1, 1.2, 1.3])
    base_region_sizes       = [a - b for a,b in zip(base_region_boundaries[1:], base_region_boundaries[:-1])]
    base_elements_per_region =  [int(round( size * 10 * n_elements_per_10_cm)) for size in base_region_sizes]
    return base_region_boundaries, base_region_sizes, base_elements_per_region

def create_vertices(region_boundaries, elements_per_region):
    region_sizes          = [a - b for a,b in zip(region_boundaries[1:], region_boundaries[:-1])]    

    vertices = jnp.concatenate([
        jnp.linspace(start, start + size, n_elements, endpoint=False) 
        for start, size, n_elements in zip(region_boundaries[:-1], region_sizes, elements_per_region)
    ] + [jnp.array([region_boundaries[-1]])])[:, None]
    return vertices

def create_mesh_and_source(n_elements_per_10_cm : int):
    region_boundaries, region_sizes, elements_per_region = base_region_stats(n_elements_per_10_cm)
    vertices     = create_vertices(region_boundaries, elements_per_region)
    connectivity = jnp.array([[i, i+1] for i in range(sum(elements_per_region))])

    water_idx, iron_idx, void_idx = 0, 1, 2    
    material_indices = [water_idx, water_idx, water_idx, iron_idx, void_idx, water_idx, void_idx]

    material_indices_elements =jnp.concatenate([
        jnp.full(n_elements, material) for n_elements, material in zip(elements_per_region, material_indices)
    ])
    source_regions = [0, 1, 0, 0, 0, 0, 0]

    source_base = jnp.concatenate([
        jnp.full(n_elements, source) for n_elements, source in zip(elements_per_region, source_regions)
    ])
    
    source_values = jnp.array([9.7702e4, 4.5451e5])        
    source = source_base[None, :] * source_values[:, None] 
    cross_sections = CrossSectionsIndexed(
        indices = material_indices_elements,
        cross_section_data = CrossSectionData.from_list([water, iron, void])        
    )
    return vertices, connectivity, cross_sections, source

from jax_sn.operator_creation import MultiGroupBuilder, DefaultMultiGroupStrategy
from jax_sn.solution_domain import BasixLagrangianSimplex
from jax_sn.quadrature_set import create_tn_quadrature_set
n_elements_per_10_cm = 2
degree  = 1
tn_order = 3


vertices, connectivity, cross_sections, source = create_mesh_and_source(n_elements_per_10_cm )
strategy             = DefaultMultiGroupStrategy.default_multi_group_strategy()
element              = BasixLagrangianSimplex(degree = degree, dimension = 1)
quadrature_set       = create_tn_quadrature_set(tn_order)

mg_builder = MultiGroupBuilder(vertices, connectivity, element, cross_sections,  quadrature_set, strategy)

from jax_sn.tools.plotly import plot_domain
plot_domain(mg_builder.domain, material_names = ['Water', 'Iron', 'Void'], width =800)

Before solving we create an isotropic source (we repeat along the basis to have a flat source within each element):

from jax_sn.solvers.right_hand_side import TotalSource

def create_source(mg_builder, elementwise_source):
    isotropic_src = jnp.repeat(elementwise_source[..., None], mg_builder.operator_data.n_basis, axis=-1)
    total_source = TotalSource.from_operator_and_isotropic_source(mg_builder.multi_group_operator, isotropic_src)
    return total_source

source_p1 = create_source(mg_builder, source)

Solving

Solving is easy in this formulation! We first create the correct right-hand-side:

\[ \mathbf{b} = \mathbf{M} \mathbf{T}^{-1} \mathbf{D} \mathbf{H} [q_{g\ell ei}] \]

This is precisely the chain strategy in multigroup form. We could create a separate multigroup version of it, but here we just scan the single-group chain strategy we already made.

from jax_sn.operators.multi_group_operator.composed_multi_group_strategies import MultiGroupChainScanStrategy
mg_chain = MultiGroupChainScanStrategy(strategy.single_group_chain_strategy)

We will create a top-level function to solve, including the right-hand-side operator, since this allows for a consistent derivative later:

import lineax as lx
import equinox as eqx

@eqx.filter_jit
def solve_anisotropic(mgop , total_source : TotalSource, solver : lx.AbstractLinearSolver, mg_chain  : MultiGroupChainScanStrategy):        
    rhs   = mg_chain.forward(mgop.data, total_source)
    return lx.linear_solve(operator = mgop, vector = rhs, solver = solver)

solution = solve_anisotropic(mg_builder.multi_group_operator,source_p1, lx.BiCGStab(rtol = 1e-10, atol = 1e-10, max_steps = 20), mg_chain)
print("Solution converged in", solution.stats['num_steps'], "steps")

Solution converged in 15 steps

Let’s plot it:

Show code cell source

Hide code cell source

from jax_sn.tools.plotly import plot_1d_solution

group_names  = ["E₀", "E₁"]
# the domain's regions, shaded behind the flux so the material layout is readable off the plot itself
regions      = [(0.0, 0.2, "Water"), (0.2, 0.3, "Source"), (0.3, 0.8, "Water"), (0.8, 1.0, "Iron"),
                (1.0, 1.1, "Void"),  (1.1, 1.2, "Water"),  (1.2, 1.3, "Void")]
region_style = {"Water" : dict(fillcolor = "#8C9BAB", opacity = 0.13),
                "Source": dict(fillcolor = "#00A67E", opacity = 0.22),
                "Iron"  : dict(fillcolor = "#4A5560", opacity = 0.30),
                "Void"  : dict(fillcolor = "#FFFFFF", opacity = 0.00)}

def interpolate_and_plot(solution_domain, solution, fig = None, dash = None, label = ""):
    '''Scalar flux (the l=0 moment) per energy group. `plot_1d_solution` samples every element
    separately, so inter-element jumps show as real gaps rather than being smoothed over.
    Pass an existing `fig` (and a `dash`/`label`) to overlay a second discretisation.'''
    new_fig = plot_1d_solution(solution_domain, solution.value[:, 0], xaxis_title = "x [m]", yaxis_title = "ϕ",
                               solution_names = [name + label for name in group_names], width = 800)
    new_fig.update_traces(line = dict(dash = dash, width = 2.5))
    if fig is not None:
        return fig.add_traces(new_fig.data)

    for x0, x1, name in regions:
        new_fig.add_vrect(x0 = x0, x1 = x1, line_width = 0, layer = "below", annotation_text = name,
                          annotation_position = "outside top", annotation_font_size = 10,
                          **region_style[name])
    new_fig.update_xaxes(range = [0.0, 1.3])
    new_fig.update_yaxes(type = "log", exponentformat = "power", showgrid = True)
    new_fig.update_layout(template = "plotly_white", width = 800, height = 480, hovermode = "x unified",
                          legend = dict(x = 0.99, y = 0.97, xanchor = "right", yanchor = "top",
                                        bgcolor = "rgba(255,255,255,0.75)", bordercolor = "#C8C8C8",
                                        borderwidth = 1),
                          margin = dict(l = 60, r = 20, t = 45, b = 50))
    return new_fig

interpolate_and_plot(mg_builder.solution_domain, solution)

Clearly, the linear elements are not so smooth in this regime! Let’s do third order elements, re-using the same strategies.

mg_p3 = mg_builder.updated(element =  BasixLagrangianSimplex(degree = 3, dimension=1))
solution_p3 = solve_anisotropic(mg_p3.multi_group_operator, create_source(mg_p3, source), lx.BiCGStab(rtol = 1e-10, atol = 1e-10, max_steps = 20), mg_chain)

fig = interpolate_and_plot(mg_builder.solution_domain, solution, dash = "dot", label = " p=1")
interpolate_and_plot(mg_p3.solution_domain, solution_p3, fig = fig, label = " p=3")

Naturally the third order elements are significantly smoother and have less discontinuity.

Exploiting downscatter structure

The above solved the multigroup equation by applying the multigroup operator. These cross sections do however not include any upscattering, and as such, the resulting multigroup structure is lower triangular:

cross_sections.cross_section_data.downscattering_only

Array(True, dtype=bool)

This can be exploited by using a more specialized lower-triangular block solver.

from jax_sn.solvers import MultiGroupLowerTriangularBlockSolver, SolverDiagnostics

For each particular single group block, the equation is:

\[( \mathbf{I} - \mathbf{M}_g \mathbf{T}_g^{-1} \mathbf{H}_g \mathbf{D}_g \mathbf{S}_{gg})[\phi_{g\ell ei}] = \mathbf{M}_g \mathbf{T}_g^{-1} \mathbf{H}_g \mathbf{D}_g[q_{g\ell ei} + \mathbf{S}_{g}[\phi_{g\ell ei}]] \]

where \(\mathbf{S}_g\) is the familiar gather-scatter operator (which gathers contributions from all, in this case only higher, energy groups.) and \(\mathbf{S}_{gg}\) is the same operator, but only for the current group (or equivalently, the inscatter operator)

Therefore, to use a lower-triangular block solver, we need strategies for the single-group versions. These can be the same as for the MultiGroupOperator, or just arbitrary single-group strategies. The solver is not differentiated, so the single-group strategies don’t even need to be differentiable, as long as they are consistent with the MultiGroupOperator!

It needs three:

• SingleGroupStrategy: \(\mathbf{I}- \mathbf{M}_g \mathbf{T}_g^{-1} \mathbf{H}_g \mathbf{D}_g \mathbf{S}_{gg}\)

• SingleGroupChain: \(\mathbf{M}_g \mathbf{T}_g^{-1} \mathbf{H}_g \mathbf{D}_g\)

• SingleGroupGatherScatter: \(\mathbf{S}_g\)

In practice, we can just create it from a composed multigroup operator easily, which just then requires the inscatter strategy (this has no multigroup analogue).

from jax_sn.operators.scatter_operator import IndexedScatter

single_group_solver = lx.BiCGStab(rtol = 1e-10, atol = 1e-10, max_steps = 50)

block_solver = MultiGroupLowerTriangularBlockSolver.from_composed(mg_builder.strategy, single_group_solver, 
                                                   scatter_strategy = IndexedScatter(), 
                                                   diagnostics = SolverDiagnostics(True, True))

print(block_solver)

MultiGroupLowerTriangularBlockSolver(
  single_group_solver=BiCGStab(rtol=1e-10, atol=1e-10, max_steps=50),
  diagonal_strategy=ComposedSingleGroupStrategy(
    transport_strategy=VmapAnglesForward(inner=VmapForward()),
    transport_inversion_strategy=VmapAnglesInversion(inner=VmapSweep()),
    moment_strategy=EinsumFluxMoment(),
    mass_strategy=LaxMass(),
    discrete_strategy=EinsumFluxMoment(),
    scatter_strategy=IndexedScatter()
  ),
  chain_strategy=ComposedSingleGroupChainStrategy(
    transport_strategy=VmapAnglesForward(inner=VmapForward()),
    transport_inversion_strategy=VmapAnglesInversion(inner=VmapSweep()),
    moment_strategy=EinsumFluxMoment(),
    mass_strategy=LaxMass(),
    discrete_strategy=EinsumFluxMoment()
  ),
  gather=BroadcastedSingleGroupGather(),
  diagnostics=SolverDiagnostics(print_sweep=True, print_single_group=True)
)

The single-group strategies are then created from the multigroup strategy!

solution_block = solve_anisotropic(mg_builder.multi_group_operator, create_source(mg_builder, source), block_solver, mg_chain)

Using lower triangular block solver for multi-group operator with 2 energy groups.
Solved energy group 0 in 11 steps
Solved energy group 1 in 14 steps

Naturally, this are the same (up to machine precision; per-group tolerances are different then total tolerance) as the original result

import numpy as onp
onp.testing.assert_allclose(solution_block.value, solution.value, atol = 1e-12)

Taking derivatives

Let’s now take some derivatives. We parametrize the region boundaries and create a function that solves and sums the square of the result[1]:

from typing import Callable
from jax_sn.operators.multi_group_operator import MultiGroupChainStrategy
n_elements_per_10_cm = 2
base_region_boundaries, _ ,elements_per_region = base_region_stats(n_elements_per_10_cm)


@eqx.filter_jit
def solve_and_sum(theta : jax.Array, f_vertices : Callable, builder : MultiGroupBuilder,
                   elementwise_source : jax.Array, 
                  solver : lx.AbstractLinearSolver, 
                  mg_chain : MultiGroupChainStrategy):
    new_vertices = f_vertices(theta)
    new_builder  = builder.updated(vertices = new_vertices)
    new_source   = create_source(new_builder, elementwise_source)
    new_solution = solve_anisotropic(new_builder.multi_group_operator, new_source, solver, mg_chain)
    return jnp.sum(new_solution.value**2)

f_vertices         = lambda theta : create_vertices(theta, elements_per_region)
grad_solve_and_sum = eqx.filter_jit(eqx.filter_grad(solve_and_sum))

These can accept both the total solver and the block solver:

args_total_solver = (base_region_boundaries, f_vertices, mg_builder, source, lx.BiCGStab(rtol = 1e-10, atol = 1e-10, max_steps = 20), mg_chain)

print("Result", solve_and_sum(*args_total_solver))

print("Gradient", grad_solve_and_sum(*args_total_solver))

Result 4273209339.5416546

Gradient [ 1.66226512e+09 -2.86813814e+10  2.76858109e+10 -6.66694616e+08
 -1.85055666e+01  1.92998983e-02 -1.92998983e-02 -1.22749803e-14]

args_block_solver = (base_region_boundaries, f_vertices, mg_builder, source, block_solver, mg_chain)
print("Result", solve_and_sum(*args_block_solver))
print("Gradient", grad_solve_and_sum(*args_block_solver))

Using lower triangular block solver for multi-group operator with 2 energy groups.
Solved energy group 0 in 11 steps
Solved energy group 1 in 14 steps
Result 4273209339.541652

Using lower triangular block solver for multi-group operator with 2 energy groups.
Using upper triangular transpose solver for multi-group operator with 2 energy groups.
Solved energy group 0 in 11 steps
Solved energy group 1 in 14 steps
Solved energy group 1 in 14 steps
Solved energy group 0 in 11 steps
Gradient [ 1.66226512e+09 -2.86813814e+10  2.76858109e+10 -6.66694616e+08
 -1.85055666e+01  1.92998983e-02 -1.92998983e-02 -1.22746009e-14]

Again, up to numerical precision, both results are identical! We can see that the adjoint solve of the block solver goes the other direction than the forward solve (solving the upper triangular transpose).

This concludes the brief tour of the entire code. The Examples directory contains some more involved examples.

---

[1]

We square to make the function non-linear; the compiler will optimize a linear function (since the cotangent values then don’t actually depend on the primals) to not solve the adjoint problem after the forward solve but concurrently (which isn’t pedagogically useful since it’s not the general situation)