Multigroup operator

import jax
jax.config.update("jax_enable_x64", True)
jax.config.update("jax_platforms", "cpu")
import jax.numpy as jnp

A multigroup operator again requires picking a strategy. However, to actually solve the multigroup 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{D}\mathbf{H}[q_{g\ell\epsilon i}] \]

there are actually two multigroup operators requires: one including scattering:

\[ \text{Multigroup operator}: (\mathbf{I} - \mathbf{M}\mathbf{T}^{-1}\mathbf{H}\mathbf{D}\mathbf{S})\]

and one that does not include scattering:

\[ \text{Multigroup chain operator}: \mathbf{M}\mathbf{T}^{-1}\mathbf{D}\mathbf{H}\]

They can both be constructed and have their own strategy types.

from jax_sn.operators.multi_group_operator import MultiGroupStrategy, MultiGroupChainStrategy

Currently, the main implementations just use a composed strategy that applies a single-group strategy and then maps it over the groups (using a serial scan or a vmap), or in other words:

\[ \text{for all } g: \quad \mathbf{I} - \underbrace{\mathbf{M}_g \mathbf{T}^{-1}_g \mathbf{H}_g \mathbf{D}_g}_{\text{chain}}\underbrace{\mathbf{S}_g}_{\text{gather-scatter}} \]

Note that \(\mathbf{S}\) requires the full multi-group unknown vector: it gathers the scattering output of all groups into the specific group \(g\) (i.e. maps \(G \times L \times E \times B \to L\times E \times B\)), while the other operators just act on a single-group unknown vector.

This gather scatter is simply

\[ \mathbf{S}_g[\phi_{\tilde{g}\ell ei}] = \sum_{\tilde{g}}\sigma_{\ell g \tilde{g} e} \phi_{\tilde{g\ell ei}}\]

while the single-group multi-angle transport operator, moment, mass, and discrete operators were all discussed in previous notebooks.

Constructing a single-group chain strategy is usually done by combining the suboperators, where each one has its own strategy too (i.e., a ‘tower’ of strategies)

from jax_sn.operators.multi_group_operator import MultiGroupChainScanStrategy, MultiGroupChainGatherScanStrategy
from jax_sn.operator_creation.default_strategies import *
mg_chain_gather_strategy = MultiGroupChainGatherScanStrategy(
    single_group_chain_strategy = ComposedSingleGroupChainStrategy(
        transport_strategy           = VmapAnglesForward(VmapForward()),
        transport_inversion_strategy = VmapAnglesInversion(VmapSweep()),
        moment_strategy              = EinsumFluxMoment(),
        mass_strategy                = LaxMass(),
        discrete_strategy            = EinsumFluxMoment()
    ),
    single_group_gather_scatter_strategy = BroadcastedSingleGroupGather()  
)

mg_chain_strategy = MultiGroupChainScanStrategy(
    single_group_chain_strategy = mg_chain_gather_strategy.single_group_chain_strategy,
)
print(mg_chain_gather_strategy)

MultiGroupChainGatherScanStrategy(
  single_group_chain_strategy=ComposedSingleGroupChainStrategy(
    transport_strategy=VmapAnglesForward(inner=VmapForward()),
    transport_inversion_strategy=VmapAnglesInversion(inner=VmapSweep()),
    moment_strategy=EinsumFluxMoment(),
    mass_strategy=LaxMass(),
    discrete_strategy=EinsumFluxMoment()
  ),
  single_group_gather_scatter_strategy=BroadcastedSingleGroupGather()
)

This is somewhat cumbersome, and in practice you’d probably not construct it like this. We just show it here to illustrate how strategies are used internally.

Instead, you could construct one of the defaults:

from jax_sn.operator_creation.default_strategies import DefaultMultiGroupStrategy

print("Vmap strategy:\n", DefaultMultiGroupStrategy.default_multi_group_strategy())
print("Warp strategy:\n", DefaultMultiGroupStrategy.default_multi_group_warp_strategy(batch_size = 128))

Vmap strategy:
 MultiGroupChainGatherScanStrategy(
  single_group_chain_strategy=ComposedSingleGroupChainStrategy(
    transport_strategy=VmapAnglesForward(inner=VmapForward()),
    transport_inversion_strategy=VmapAnglesInversion(inner=VmapSweep()),
    moment_strategy=EinsumFluxMoment(),
    mass_strategy=LaxMass(),
    discrete_strategy=EinsumFluxMoment()
  ),
  single_group_gather_scatter_strategy=BroadcastedSingleGroupGather(),
  checkpoint=True
)
Warp strategy:
 MultiGroupChainGatherScanStrategy(
  single_group_chain_strategy=ComposedSingleGroupChainStrategy(
    transport_strategy=VmapAnglesForward(inner=BatchedForward(batch_size=128)),
    transport_inversion_strategy=WarpLevelAngles(),
    moment_strategy=EinsumFluxMoment(),
    mass_strategy=WarpMass(),
    discrete_strategy=EinsumFluxMoment()
  ),
  single_group_gather_scatter_strategy=BroadcastedSingleGroupGather(),
  checkpoint=True
)

Recall that the WarpLevelAngles is a fused multi-angle strategy and therefore has no inner strategy!

Let’s now construct a MultiGroupOperator. First we again make a similar tokamak domain:

Show code cell source

Hide code cell source

from jax_sbgeom.flux_surfaces.miller_tokamak import MillerTokamak
import jax_sbgeom as jsb
from jax_sn.domain import Domain, CrossSectionsIndexed, CrossSectionData
from jax_sn.solution_domain import BasixLagrangianSimplex
from jax_sn.tools.plotly import plot_domain

tokamak                 = MillerTokamak(rminor = 1.0, rmajor = 1.2, elongation = 0.5, triangularity = 0.2)

def construct_domain_solution_element_quadrature_set(n_s, n_theta):


    mesh_tokamak            = jsb.flux_surfaces.flux_surface_meshing.mesh_poloidal_plane(tokamak, jnp.linspace(0.0,1.0,n_s), 0.0, n_theta)
    # Convert it to a 2D mesh!
    mesh_tokamak           = (jnp.stack([mesh_tokamak[0][:,0], mesh_tokamak[0][:,2]], axis=-1), mesh_tokamak[1])
    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))
    )

    water_idx, iron_idx, void_idx = 0, 1, 2
    cross_section_indices         = jnp.concatenate([jnp.full((mesh_tokamak[1].shape[0]//3,), water_idx), jnp.full((mesh_tokamak[1].shape[0]//3,), iron_idx), jnp.full((mesh_tokamak[1].shape[0] - 2*(mesh_tokamak[1].shape[0]//3),), void_idx)])
    cross_sections_tokamak        = CrossSectionsIndexed(indices = cross_section_indices, cross_section_data = CrossSectionData.from_list([water, iron, void]))


    element_order = 2
    solution_element = BasixLagrangianSimplex(element_order, 2)

    quadrature_order = 3
    from jax_sn.quadrature_set import create_tn_quadrature_set
    quadrature_set  = create_tn_quadrature_set(quadrature_order)
    domain_tokamak                = Domain.from_mesh_and_cross_sections(mesh_tokamak, solution_element.face_template, cross_sections_tokamak)
    return domain_tokamak, quadrature_set, solution_element



n_s     = 4 
n_theta = 7

domain_tokamak, quadrature_set, solution_element = construct_domain_solution_element_quadrature_set(n_s, n_theta)

plot_domain(domain_tokamak, width=800, height=800, material_names = ["water", "iron", "void"])

Then we construct the OperatorData object, which contains all data required for the multigroup operator. This contains both the DiscretisationData which includes all spatial/angular/energy discretisation information, and the NonDifferentiableData, which contains (non-operator-defining) data that cannot be differentiated and is more like setup data (i.e. topological ordering). The strategy builds that tier itself, since what it needs of it (here: a plain sweep ordering) is the strategy’s business:

from jax_sn.operators.operator_data import OperatorData, DiscretisationData, SpatialData, AngularData, NonDifferentiableData
from jax_sn.solution_domain import SolutionDomain
from jax_sn.operators.multi_group_operator import MultiGroupOperator


operator_data = OperatorData(
    DiscretisationData(
        SolutionDomain.from_element_and_domain(solution_element, domain_tokamak).spatial_data, 
        AngularData.from_quadrature_set(quadrature_set, domain_tokamak.cross_sections_indexed.l_max)
    ),
    mg_chain_gather_strategy.build_default_non_differentiable_data(domain_tokamak, quadrature_set.angles)
)

mg_operator = MultiGroupOperator(operator_data, mg_chain_gather_strategy)

print(mg_operator)

MultiGroupOperator(
  data=OperatorData(
    discretisation_data=DiscretisationData(
      spatial=SpatialData(
        domain=Domain(
          vertices=f64[22,2],
          connectivity=i64[35,3],
          face_normals=f64[35,3,3],
          etoe=weak_i64[35,3],
          etof=weak_i64[35,3],
          face_perm_idx=i64[35,3],
          face_template=i64[3,2],
          cross_sections_indexed=CrossSectionsIndexed(
            indices=weak_i64[35],
            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[35,2,2],
          vol_jacobians_inv=f64[35,2,2],
          vol_jacobians_det=f64[35],
          face_jacobians_det=f64[35,3]
        ),
        element=<jax_sn.solution_domain.basix_fem.BasixLagrangianSimplex object at 0x7faaf81cf620>
      ),
      angular=AngularData(
        quadrature_set=QuadratureSet(angles=f64[72,3], weights=f64[72]),
        wk_ylm=f64[72,4],
        ylm=f64[72,4]
      )
    ),
    non_differentiable_data=NonDifferentiableData(
      sweep_data=BasicSweepData(levels=i64[72,35], levels_argsort=i64[72,35])
    )
  ),
  strategy=MultiGroupChainGatherScanStrategy(
    single_group_chain_strategy=ComposedSingleGroupChainStrategy(
      transport_strategy=VmapAnglesForward(inner=VmapForward()),
      transport_inversion_strategy=VmapAnglesInversion(inner=VmapSweep()),
      moment_strategy=EinsumFluxMoment(),
      mass_strategy=LaxMass(),
      discrete_strategy=EinsumFluxMoment()
    ),
    single_group_gather_scatter_strategy=BroadcastedSingleGroupGather()
  )
)

We can apply the multigroup operator as any other operator:

print("Shape of application: [n_groups, n_moments, n_elements, n_basis]",
       mg_operator.mv(jnp.ones_like(mg_operator.in_structure())).shape)

Shape of application: [n_groups, n_moments, n_elements, n_basis] (2, 4, 35, 6)

The meaning of an application is not as clear here though: in some sense, it means, given a flux \(\phi_{g\ell ei}\), how much of that flux had to come from an external source (and not from scattering its own distribution).

The main use of the multigroup operator is solving the multigroup equation, after which its result is just the flux moments.

We have now shown every part of the multigroup operator construction. Before showing how to solve it and post-process it, we will show a much easier and more efficient way to construct it from the set of quantities that fully define it.