Multi-angle sweeping

This notebook shows how to do multi-angle sweeping, including fused strategies versus composed strategies.

We first create the same circular domain as before:

Hide code cell source

 
import jax
jax.config.update("jax_enable_x64", True)
jax.config.update("jax_platforms", "cpu")
import jax.numpy as jnp
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)

mesh_tokamak = jsb.flux_surfaces.flux_surface_meshing.mesh_poloidal_plane(tokamak, jnp.linspace(0.0,1.0,4), 0.0, 7)
# Convert it to a 2D mesh!
mesh_tokamak           = (jnp.stack([mesh_tokamak[0][:,0], mesh_tokamak[0][:,2]], axis=-1), mesh_tokamak[1])

cross_section          = CrossSectionData(
                            total     = jnp.array([[5.0]]), 
                            scattering= jnp.array([[[[0.0]]]]))

cross_sections_tokamak = CrossSectionsIndexed(indices = jnp.full((mesh_tokamak[1].shape[0],), 0), cross_section_data = cross_section)

domain_tokamak         = Domain.from_mesh_and_cross_sections(mesh_tokamak, BasixLagrangianSimplex(1, 2).face_template, cross_sections_tokamak)
plot_domain(domain_tokamak, width=800, height=800)

Let’s create a quadrature set:

from jax_sn.quadrature_set import create_tn_quadrature_set

quadrature_set = create_tn_quadrature_set(3)

In general, composed operators have two possibilities (also for single group and multigroup operators): either we use a sub-operator and compose it using Python functions (e.g. vmap or scan), or we implement the entire operator immediately. We will now show how to construct multi-angle transport operator inversion using either approach, but always use Vmap forward strategy.

First, we will just vmap the a BatchedSweep (which instead of first inverting all element matrices and then collapsing, does this batch-by-batch), and print what it requires of the sweep data (nothing: a plain topological sort will do):

from jax_sn.operators.transport_operator import TransportOperatorMultiAngle, VmapSweep, VmapAnglesInversion, VmapForward, VmapAnglesForward, BatchedSweep
multi_angle_forward_strategy   = VmapAnglesForward(VmapForward())
multi_angle_inversion_strategy = VmapAnglesInversion(BatchedSweep(5))
print("Required sweep-data capabilities:", multi_angle_inversion_strategy.required_capabilities())

Required sweep-data capabilities: ()

The second sweep type is a fused Warp kernel, which requires Levels data. This means data to parallelize over the levels of the elements as well (instead of just the topological sort)

from jax_sn.operators.transport_operator import WarpLevelAngles

multi_angle_warp_strategy = WarpLevelAngles()
print("Required sweep-data capabilities:", multi_angle_warp_strategy.required_capabilities())

plot_domain(domain_tokamak, show_levels=True, level_mode = "levels", graph_angle = jnp.array([1/ jnp.sqrt(2), 1/ jnp.sqrt(2)]), show_graph = True, width=800, height=800)

Required sweep-data capabilities: (<class 'jax_sn.operators.transport_operator.sweep_data_implementations.ProvidesLevels'>,)

Elements 0 and 1 can be computed in parallel; this Warp-based kernel parallelizes that too to achieve higher parallelism (mostly useful on GPU).

Let’s actually construct them (warp requires \(\leq 4\) basis functions, so limited to linear triangles)

from jax_sn.solution_domain import SolutionDomain
from jax_sn.operators.transport_operator import TransportData
solution_domain = SolutionDomain.from_element_and_domain(BasixLagrangianSimplex(1, 2), domain_tokamak)

transport_data_basic = TransportData(solution_domain.spatial_data, 0, quadrature_set.angles,
                                      multi_angle_inversion_strategy.build_default_sweep_data(domain_tokamak, quadrature_set.angles))
transport_data_level = TransportData(solution_domain.spatial_data, 0, quadrature_set.angles, 
                                     multi_angle_warp_strategy.build_default_sweep_data(domain_tokamak, quadrature_set.angles))

The added sweep data is permissive: we can construct a multi-angle strategy requiring the basic data with the more specialized data (but obviously not vice-versa):

from jax_sn.operators.transport_operator import TransportOperatorMultiAngle

TransportOperatorMultiAngle(transport_data_level, multi_angle_forward_strategy, multi_angle_inversion_strategy)
try: 
    TransportOperatorMultiAngle(transport_data_basic, multi_angle_forward_strategy, multi_angle_warp_strategy)
except Exception as e:
    print("Expected error:", e)

Expected error: WarpLevelAngles requires sweep data providing ['ProvidesLevels'], but got BasicSweepData. Build it via WarpLevelAngles.build_default_sweep_data(geometric_domain, angle).

transport_operator_vmap = TransportOperatorMultiAngle(transport_data_basic, 
                                                      multi_angle_forward_strategy, 
                                                      multi_angle_inversion_strategy)
transport_operator_warp = TransportOperatorMultiAngle(transport_data_level, 
                                                      multi_angle_forward_strategy, 
                                                      multi_angle_warp_strategy)

They have the same inversion call as single-angle transport operators:

rhs = jnp.ones_like(transport_operator_vmap.out_structure())

sol_1 = transport_operator_vmap.inv(rhs)
sol_2 = transport_operator_warp.inv(rhs)

import numpy as np

print("Shape [n_angles, n_elements, n_basis]:", sol_1.shape)
print("Difference between vmap and warp solutions:", np.max(np.abs(sol_1 - sol_2)))

Warp CUDA warning: Could not find or load the NVIDIA CUDA driver. GPU execution will not be available.

Warp 1.18.0 initialized:
   CUDA Toolkit 13.4, CUDA driver not available (NVRTC compilation available)
   Devices:
     "cpu"      : "x86_64"
   Kernel cache:
     /home/runner/.cache/warp/1.18.0

Module jax_sn.operators.transport_operator.implementations.simplex_warp_kernel 74dc481 load on device 'cpu' took 1842.52 ms  (compiled)

Module jax_sn.operators.transport_operator.implementations.simplex_warp_kernel_nb3_nf3 dd152f7 load on device 'cpu' took 2253.32 ms  (compiled)
Shape [n_angles, n_elements, n_basis]: (72, 35, 3)
Difference between vmap and warp solutions: 4.263256414560601e-14

If the operator is composed or has a ‘natural’ single-angle version, it supports indexing to get the single-angle operator for a given angle:

transport_operator_vmap[0]

TransportOperatorSingleAngle(
  data=TransportData(
    spatial_data=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[1,1], scattering=f64[1,1,1,1]
          )
        )
      ),
      fem_matrices=FEMMatrices(
        mass_matrix=f64[3,3],
        integration_coeff=f64[3],
        derivative_matrix=f64[2,3,3],
        outflow_matrices=f64[3,3,3],
        inflow_matrices=f64[3,2,3,3,3]
      ),
      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 0x7f50d0b8b050>
    ),
    energy_group=0,
    angle=f64[3],
    sweep_data=BasicSweepData(levels=i64[35], levels_argsort=i64[35])
  ),
  strategy=VmapForward(),
  inversion_strategy=BatchedSweep(batch_size=5)
)

transport_operator_warp[0]

TransportOperatorSingleAngle(
  data=TransportData(
    spatial_data=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[1,1], scattering=f64[1,1,1,1]
          )
        )
      ),
      fem_matrices=FEMMatrices(
        mass_matrix=f64[3,3],
        integration_coeff=f64[3],
        derivative_matrix=f64[2,3,3],
        outflow_matrices=f64[3,3,3],
        inflow_matrices=f64[3,2,3,3,3]
      ),
      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 0x7f50d0b8b050>
    ),
    energy_group=0,
    angle=f64[3],
    sweep_data=LeveledSweepData(
      levels=i64[35],
      levels_argsort=i64[35],
      level_fwd=LevelBasedParallelism(
        first_occurences=i64[15], ends=i64[15], sizes=i64[15]
      ),
      level_rev=LevelBasedParallelism(
        first_occurences=i64[15], ends=i64[15], sizes=i64[15]
      )
    )
  ),
  strategy=VmapForward(),
  inversion_strategy=WarpLevelSweep()
)

To recap, a composed strategy in this case is just either two nested strategies or one bigger:

from jax_sn.operators.transport_operator import MultiAngleInversionStrategy
multi_angle_inversion_strategy = VmapAnglesInversion(BatchedSweep(5))
multi_angle_warp_strategy      = WarpLevelAngles()

print(isinstance(multi_angle_inversion_strategy, MultiAngleInversionStrategy))
print(isinstance(multi_angle_warp_strategy, MultiAngleInversionStrategy))

True
True

Naturally, in tests, we test that all such implementations are equivalent up to numerical precision, including VJP/JVPs and adjoint/transport/inversion idempotence for a set of representative dimension/domain combinations (and similarly for other operators!)