Parametrization

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

Recall how we constructed a OperatorData in the previous notebook:

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)

Besides verbose, it’s also not optimal: if you just change the cross sections, doing the above again recomputes everything, from topological connections to angular quadrature data. Instead, we can use the parametrization module to construct the operator data efficiently from a set of primitives which define the entire operator:

  • Vertices

  • Connectivity

  • Cross sections

  • Finite Element

  • Quadrature set / angles

  • Strategy

  • keyword arguments to non-differentiable data (in some cases needed; maximum levels in sweep data)

All other data and operators follow directly from these seven primitives. However, not all data depends on all seven; the topology of the mesh (which element connects to which) only depends on the connectivity for example.

This dependency can be visualized in a graph using a builder helper:

from jax_sn.operator_creation.builders import *

from jax_sn.tools.plotly import plot_derivation_graph

plot_derivation_graph(MultiGroupBuilder)

From this picture, it’s clearly visible that if we just vary the cross sections, there’s no need to compute face normals or the mesh topology or even the precomputed geometric data.

But this is also just automated:

plot_derivation_graph(MultiGroupBuilder.varying("cross_sections_indexed"))

We can see that only varying that parameter has the effect of storing all the parameters that do not depend on the cross sections! Some terminology:

  • “retained primitive” means that in principle it’s free as long as it satisfies some properties. In this case, if the \(L\) order of the cross section would change, this would require updating the angular data as well (it’s dependen on \(L\)!). But in most cases, you don’t change the shape, and just change the value, in which case it doesn’t have to be recomputed.

  • “frozen primitive” means that this primitive cannot be changed given the stored data.

  • “stored” means stored in the object itself (i.e. computed on initialisation of the builder)

  • “derived” means computed on changing the non-frozen primitives: in this case, it’s all just containers anyway.

  • “projected” means that this object is simply contained in another larger object; e.g. the face normals are present in the stored geometric domain (to avoid storing all of them multiple times!)

Actually constructing them requires having the primitives. We will re-use the same domain from the previous notebook and also re-use the default strategy.

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, element = construct_domain_solution_element_quadrature_set(n_s, n_theta)
vertices = domain_tokamak.vertices
connectivity = domain_tokamak.connectivity
cross_sections_indexed = domain_tokamak.cross_sections_indexed

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

from jax_sn.operator_creation.default_strategies import *
mg_strategy = DefaultMultiGroupStrategy.default_multi_group_strategy()

mg_builder = MultiGroupBuilder(vertices = vertices, 
                  connectivity = connectivity, 
                  element = element, 
                  cross_sections_indexed = cross_sections_indexed, 
                   quadrature_set = quadrature_set,                                
                   strategy = mg_strategy,
                   )

We can immediately create the multigroup operator:

_ = mg_builder.multi_group_operator

Another advantage of the parametrization is that in this first form, we can’t actually JIT the creation, let alone take gradients later:

try:
    eqx.filter_jit(lambda mg_builder: mg_builder.multi_group_operator)(mg_builder)
except Exception as e:
    print("Error during JIT compilation:", e)

Error during JIT compilation: The numpy.ndarray conversion method __array__() was called on traced array with shape int64[2]
The error occurred while tracing the function <lambda> at /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/equinox/_jit.py:43 for jit. This value became a tracer due to JAX operations on these lines:

  operation a:bool[35,3] = lt b 0:i64[]
    from line /tmp/ipykernel_2612/1038254634.py:2:38 (<lambda>)

  operation a:i64[35,3] = add b 22:i64[]
    from line /tmp/ipykernel_2612/1038254634.py:2:38 (<lambda>)
See https://docs.jax.dev/en/latest/101/errors.html#jax.errors.TracerArrayConversionError

but if we freeze the element (and thus already compute face templates and fem matrices), this is no longer an issue (and of course the result is the same as the earlier builder, up to numerical precision when JIT’ed)

eqx.tree_equal(mg_builder.frozen("element").multi_group_operator, mg_builder.multi_group_operator)

Array(True, dtype=bool)

eqx.tree_equal(eqx.filter_jit(lambda a : a.multi_group_operator)(mg_builder.frozen("element")), mg_builder.multi_group_operator, atol= 1e-12, rtol =1e-12)

Array(True, dtype=bool)

The nice thing is that it allows us to just take gradients from the primitives easily without worrying about the specifics of what needs to be recomputed:

varying_vertices = mg_builder.varied("vertices")

def _updated_mg(theta, f_theta, mg_builder):
    new_vertices = f_theta(theta, mg_builder.domain.vertices)
    updated_mg = mg_builder.updated(vertices =  new_vertices)
    return updated_mg

@eqx.filter_jit
def objective(theta, f_theta, mg_builder):
    updated_mg = _updated_mg(theta, f_theta, mg_builder)
    rhs = jnp.ones_like(updated_mg.multi_group_operator.in_structure())
    return jnp.sum(updated_mg.multi_group_operator.mv(rhs))

f_theta = lambda theta, old_vertices : old_vertices.at[-1].add(theta)

_grad_objective = eqx.filter_jit(jax.grad(objective, argnums=0))

print("Finite differences", (objective(jnp.array([0.1,0.04]), f_theta, varying_vertices) - objective(jnp.array([0.1 + 1e-5,0.04]), f_theta, varying_vertices))/-1e-5)
print("Autodiff", _grad_objective(jnp.array([0.1,0.04]), f_theta, varying_vertices))

Finite differences 11.198486345165291

Autodiff [ 11.19827927 -13.19499493]

Note that changing the vertices actually has an impact on the topological sort, but this is handled properly by the NonDifferentiableData through stop_gradient and recomputed.

levels_updated  = _updated_mg(jnp.array([0.1,0.04]), f_theta, varying_vertices).sweep_data.levels
levels_original = mg_builder.sweep_data.levels
print("Levels equal: ", jnp.array_equal(levels_updated, levels_original))

Levels equal:  False

In practice, if your entire program is jax compatible, this kind of optimization can be performed by the compiler itself; it also has this computational graph entirely visible. However, in our case, there are two main reasons why this is still useful:

  1. As before, part of the computation might not be jittable in any case. Then the compiler cannot even compile it directly.

  2. In optimization, the optimizer might be written in a different language, and thus the compiler will not actually see the whole program. In that case, updating e.g. vertices from the base builders will invariably re-build everything.

Of course, you don’t have to use it and could still just manually construct and cache the operator data if you’d prefer; it’s just a small add-on.