Theory VI: Graph operations

We’ll look at the theory on how to solve such questions with a slightly less pathological example: a Fusion system code, inspired by Freidberg’s systems code[1].

This is such a simple systems code that it can be solved by hand, but we’ll not do that; instead, we’ll use the cottax machinery to solve.

Show the imports and physical constants

Hide the imports and physical constants

from functools import partial
from math import pi

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)
from cottax.core.frozen import OrderedMap
from cottax.core import Plan, Combine, Nest, GraphWithValues
from cottax.relational import RelationalCondition
from cottax.core.realisation import ImplementDrivers
from cottax.consumers.execution import run
from cottax.consumers.visualization.sequencing import entries

import cottax.assemblies.jax as cx
from cottax.assemblies.jax import N, flat
from cottax.consumers.visualization.xdsm import PlainFormat, spell_plain
from cottax.relational.contracts import (
    FixedPointIteration, FixedPointIterator, RootFinding, RootFinder,
)
from cottax.consumers.execution.jax_jit import JitGraph
from cottax.relational.consumers.execution.jax_solvers import Newton, Picard
from cottax.consumers.visualization import render_html

xdsm = partial(render_html, formatter = PlainFormat(),
               xdsm = {"input_label": "Input", "output_label": "Output", "collapse_names" :True}, views=("n2", "xdsm"))


MU0      = 4e-7 * pi                    # vacuum permeability          [H/m]
E_FUSION = 17.6e6 * 1.602176634e-19     # energy per DT reaction       [J]
KEV      = 1.602176634e-16              # one keV                      [J]
ION_MASS = 2.5                          # average DT ion mass          [amu]

Show the model definitions

Hide the model definitions

@flat.function('plasma_volume', 'wall_area', 'inverse_aspect_ratio')
def Geometry(major_radius, minor_radius, elongation):
    """Plasma volume, first-wall area and inverse aspect ratio."""
    return (2*pi**2*major_radius*minor_radius**2*elongation,      # m^3
            4*pi**2*major_radius*minor_radius*elongation,         # m^2
            minor_radius/major_radius)


@flat.function('plasma_current')
def Current(driven_current, bootstrap_fraction):
    """Total plasma current: what is driven externally, topped up by what bootstraps."""
    return driven_current/(1.0 - bootstrap_fraction)               # MA


@flat.function('beta_poloidal', 'beta_normalised', 'safety_factor')
def MHD(density, temperature, field_on_axis, plasma_current,
        minor_radius, major_radius, elongation):
    """Poloidal and normalised beta, and the kink safety factor."""
    pressure  = 2.0*density*temperature*KEV                        # Pa
    shaping   = 0.5*(1 + elongation**2)
    field_pol = (MU0*plasma_current*1e6
                 / (2*pi*minor_radius*jnp.sqrt(shaping)))          # T
    beta_tor  = 2*MU0*pressure/field_on_axis**2

    return (2*MU0*pressure/field_pol**2,
            100*beta_tor*minor_radius*field_on_axis/plasma_current,
            2*pi*minor_radius**2*field_on_axis*shaping
            / (MU0*major_radius*plasma_current*1e6))


@flat.function('bootstrap_fraction')
def Bootstrap(inverse_aspect_ratio, beta_poloidal, bootstrap_coeff):
    """The fraction of the plasma current the plasma drives itself."""
    return bootstrap_coeff*jnp.sqrt(inverse_aspect_ratio)*beta_poloidal


@flat.function('fusion_power', 'alpha_power', 'neutron_power')
def Fusion(density, temperature, plasma_volume):
    """DT fusion power, split into its alpha and neutron halves.

    `<sigma v> = 1.1e-24 T^2` is the usual fit over 10 - 20 keV.
    """
    reactivity   = 1.1e-24*temperature**2                          # m^3/s
    fusion_power = (0.25*density**2*reactivity
                    * E_FUSION*plasma_volume/1e6)                  # MW
    return fusion_power, fusion_power/5.0, 4.0*fusion_power/5.0


@flat.function('radiated_power')
def Radiation(density, temperature, plasma_volume, Z_eff):
    """Bremsstrahlung loss out of the plasma."""
    return (5.35e-37*Z_eff*density**2
            * jnp.sqrt(temperature)*plasma_volume/1e6)             # MW


@flat.function('heating_power')
def Heating(alpha_power, radiated_power, auxiliary_power):
    """What is left over to heat the plasma."""
    return alpha_power + auxiliary_power - radiated_power          # MW


@flat.function('confinement_time')
def Confinement(plasma_current, field_on_axis, density, major_radius,
                inverse_aspect_ratio, heating_power, elongation, h_factor):
    """Energy confinement time, ITER98(y,2) ELMy H-mode scaling."""
    return (0.0562*h_factor
            * plasma_current**0.93 * field_on_axis**0.15
            * heating_power**-0.69 * (density/1e19)**0.41
            * ION_MASS**0.19 * major_radius**1.97
            * inverse_aspect_ratio**0.58 * elongation**0.78)       # s


@flat.function('temperature')
def Transport(density, plasma_volume, confinement_time, heating_power):
    """The temperature the plasma settles at.

    In steady state the stored energy `3 n T V` is what the confinement time holds
    against the heating power, so the temperature follows from the other two.
    """
    stored_energy = heating_power*1e6*confinement_time             # J
    return stored_energy/(3.0*density*plasma_volume)/KEV           # keV


@flat.function('wall_load')
def WallLoad(neutron_power, wall_area):
    """Neutron power per unit area of first wall."""
    return neutron_power/wall_area                                 # MW/m^2


@flat.residual('blanket_thickness')
def Shield(blanket_thickness, wall_load,
           attenuation_length, allowable_coil_heating):
    """How thick the blanket and shield have to be.

    Neutron flux attenuates as `exp(-thickness/attenuation_length)`, and the shield
    is sized so that what reaches the TF coil is exactly what the coil tolerates.
    """
    reaching_coil = wall_load*jnp.exp(-blanket_thickness/attenuation_length)
    return reaching_coil/allowable_coil_heating - 1.0

@flat.function('peak_field', 'magnet_energy')
def Magnet(major_radius, minor_radius, blanket_thickness, field_on_axis):
    """Peak field at the TF coil, and the field energy it has to store.    
    """
    bore        = minor_radius + blanket_thickness                 # m
    peak_field  = field_on_axis*major_radius/(major_radius - bore) # T
    bore_volume = 2*pi**2*major_radius*bore**2                     # m^3
    return peak_field, peak_field**2*bore_volume/(2*MU0)           # J


@flat.function('electric_power')
def Plant(fusion_power, auxiliary_power,
          eta_thermal, eta_current_drive):
    """Net electric power out of the plant: what the turbines make, less what the
    current drive draws back in."""
    return (eta_thermal*(fusion_power + auxiliary_power)
            - auxiliary_power/eta_current_drive)                   # MW


def spell_flatter(iters):
    return [spell_plain(i) for i in iters]

def print_spell_flatter(iters, end = " "):
    for i in spell_flatter(iters):
        print(i, end=end)
    print('')

def print_graph_nodes(graph, start = '\t'):
    print(start, end = "")
    print_spell_flatter(list(graph.definitions.keys()))

We’ll immediately generate an N2 that is SCC’d (a xdsm starts to get complex with this size)

reactor = cx.absolute.declare(Geometry, Current, MHD, Bootstrap, Fusion, Radiation,
                      Heating, Confinement, Transport, WallLoad, Shield,
                      Magnet, Plant).graph

xdsm(reactor, blocks=True, views =("n2", "xdsm"),
     xdsm = {"input_label": "Input", "output_label": "Output", "collapse_names": True})

Beyond the numerical ordering, this also gives a view of the reactor: what depends on what? Apparently the magnet can only be computed after you’ve computed all other parameters!

We can through the graph also see what is actually an input the user has to provide (and isn’t an output of the model):

print("Unowned inputs:")
for i in reactor.graph.boundary_inputs:
    print("\t", spell_plain(i))

Unowned inputs:
	 major_radius
	 minor_radius
	 elongation
	 driven_current
	 density
	 field_on_axis
	 bootstrap_coeff
	 Z_eff
	 auxiliary_power
	 h_factor
	 attenuation_length
	 allowable_coil_heating
	 eta_thermal
	 eta_current_drive

What do we need to make this thing solvable? Well, every block needs to be solvable!

def all_acyclic(graph):
    for graph_i in (e.subgraph for e in entries(reactor)):
        print("Block:")
        print_graph_nodes(graph_i)
        print("\tIs acyclic: ", graph_i.graph.is_acyclic)


all_acyclic(reactor)

Block:
	Geometry 
	Is acyclic:  True
Block:
	Current MHD Bootstrap Fusion Radiation Heating Confinement Transport 
	Is acyclic:  False
Block:
	WallLoad 
	Is acyclic:  True
Block:
	Shield ^solve.Shield 
	Is acyclic:  False
Block:
	Magnet 
	Is acyclic:  True
Block:
	Plant 
	Is acyclic:  True

Two of them are not, so those are the two that need solving. The graph can say which variables a solver could iterate on: a variable is feedback when some reader of it can reach its own producer, so assuming it is what opens the loop.

for block in (e.subgraph for e in entries(reactor)):
    if block.graph.is_acyclic:
        continue
    print_graph_nodes(block, start ="Block: ")
    for var in block.graph.feedback_variables:
        readers = block.graph.closing_readers(var)
        print(f"    {spell_plain(var):22s} closed by ", end ="")
        print_spell_flatter(readers)

Block: Current MHD Bootstrap Fusion Radiation Heating Confinement Transport 
    plasma_current         closed by MHD Confinement 
    beta_poloidal          closed by Bootstrap 
    bootstrap_fraction     closed by Current 
    alpha_power            closed by Heating 
    radiated_power         closed by Heating 
    heating_power          closed by Confinement Transport 
    confinement_time       closed by Transport 
    temperature            closed by MHD Fusion Radiation 
Block: Shield ^solve.Shield 
    ^cond.blanket_thickness closed by ^solve.Shield 
    blanket_thickness      closed by Shield 

What we now have to do is Cut the cycle in some way. This means that you assume the variable, and let a reader use that while the solver iterates for a condition.

Of course, the block itself that the solver drives then has to be acyclic: otherwise it cannot run it.

Let’s now do it with the large cycle block. The basic Cut type is a FixedPointCut: you assume a value and the solver has to drive that assumption and what is actually computed to the same value.

Applying a GraphOp (which in the case of a FixedPointCut is multiple operations) to a Graph yields a Plan: a Graph and its operations.

This Plan is useful because it keeps a reference ot the original graph and checks whether the operations are viable and can then cheaply built the graph with the operations applied.

plan_fpc = Plan(reactor) + cx.absolute.fixed_point_cut(flat.place("bootstrap_fraction"), [N["Current"]])

# Plotting just the cycle block!
cycle_entry = entries(plan_fpc.graph)[1]
cycle_block = cycle_entry.subgraph
xdsm( cycle_block, blocks=True)

Is this enough? We can check that, but from the figure we know the answer: the block that solve(bootstrap_fraction) drives isn’t acyclic, it still has coupling below the diagonal…

cycle_entry.interior.graph.is_acyclic

False

(Solve.interior is that particular block with its problem type removed: in other words, exactly what the solver requires)

We have to make a second cut!

Let’s use the interior block’s feedback variables again to determine what we can cut:

print_spell_flatter(cycle_entry.interior.graph.feedback_variables)

alpha_power radiated_power heating_power confinement_time temperature 

Let’s pick temperature, and find out which readers need to get the assumed value:

print_spell_flatter(cycle_entry.interior.graph.closing_readers(cycle_entry.interior.graph.feedback_variables[-1]))

Fusion Radiation 

Both Fusion and Radiation! This can also be seen in the graph: cutting just Fusion actually doesn’t remove the cycle.

plan_2fpc = plan_fpc + cx.absolute.fixed_point_cut(flat.place("temperature"), [N["Fusion"], N["Radiation"]])

xdsm(plan_2fpc.graph, blocks=True)

As we can see in the XDSM, the large block has no longer any feedback variable that isn’t answered by a solver.

However, we have decision to make: do we nest these solvers, and if so which way round, or are the two really one solve?

Option 1: combining the two problems

Combining these problems is possible: both are FixedPoints and thus they combine normally.

#We name it 'mda'
combine = Combine(cx.absolute.node(cx.solve(N["mda"])), (plan_2fpc.ops[0].at, plan_2fpc.ops[1].at), RelationalCondition)

plan_combined = plan_2fpc + combine

xdsm(plan_combined.graph, blocks=True)

There’s now only one problem in the cycle, and that works![2]

As a side-note, this could have been done in one go as well:

mda = flat.fixed_point_cut({"bootstrap_fraction": N["Current"],
                            "temperature":        [N["Fusion"], N["Radiation"]]},
                           at = cx.solve(N["mda"]))

plan_double = Plan(reactor) + mda

print("Same graph: ", plan_combined.graph == plan_double.graph)

Same graph:  True

solvers = ImplementDrivers({FixedPointIterator: Picard(rtol = 1e-10, atol = 1e-10),
                            RootFinder:         Newton(rtol = 1e-10, atol = 1e-10)})

plan_total = plan_double + cx.absolute.assign(mda.at, FixedPointIteration(steps = False))
plan_total = plan_total + cx.absolute.assign(cx.solve(N["Shield"]), RootFinding(steps = False))
plan_total = plan_total + solvers

xdsm(plan_total.graph)

A Plan is the whole derivation, so it can be drawn as one: the slider above the diagram steps from the bare reactor through each operation, and highlights what that operation changed.

xdsm(plan_total, blocks=True)

The solvers require some extra data: a initial guess (see start(bootstrap_fraction)). In cottax, these are first-class Graph citizens: we could hand-initialise, but also wire the output of a different solver to it (for example a lower fidelity solver or surrogate)

Let’s see what inputs we now actually need to give. We’ll first give the Plan to a JitGraph, which checks that we can actually run it as is (internally, it’s just a blocking with more checks!) and provides distinct steps [3]

# A `JitGraph` is the graph proved runnable under jit: one step per component.
jit_combined = JitGraph(plan_total.graph)

print("Steps:")
for step in jit_combined.steps:
    print(step)

Steps:
Run(node=NodePath(['Geometry']))
Solve(problem=NodePath(^solve['mda']), interior=JitGraph(['Current'] -> ['Fusion'] -> ['Radiation'] -> ['Heating'] -> ['Confinement'] -> ['Transport'] -> ['MHD'] -> ['Bootstrap']))
Run(node=NodePath(['WallLoad']))
Solve(problem=NodePath(^solve['Shield']), interior=JitGraph(['Shield']))
Run(node=NodePath(['Magnet']))
Run(node=NodePath(['Plant']))

This notation is cottax internal notation: ^ is a new (minted) location in the namespace that only cottax can create; that way, we avoid name-clashes (e.g. a hat location in a user-defined PyTree). ^copy is a cut’s copy, the value assumed (drawn with a hat), ^cond means a condition that has to be satisfied, and ^solve means a problem has to be solved. The dictionary-like names for everything is because we are using a flat, dictionary-like interface to cottax.

Let’s now check what inputs we need to give:

print_spell_flatter(jit_combined.graph.graph.boundary_inputs, end='\n')

major_radius
minor_radius
elongation
driven_current
density
field_on_axis
bootstrap_coeff
Z_eff
auxiliary_power
h_factor
attenuation_length
allowable_coil_heating
^start.blanket_thickness
eta_thermal
eta_current_drive
^start.^copy.bootstrap_fraction
^start.^copy.temperature

We now create a initial input, including the guesses for the solvers:

design_point = dict(
    major_radius = 6.2, 
    minor_radius = 2.0, 
    elongation = 1.7,
    field_on_axis = 5.3,
    density = 1.0e20, 
    driven_current = 15.0, 
    bootstrap_coeff = 0.7, 
    Z_eff = 1.6,
    auxiliary_power = 50.0,
    h_factor = 1.0, 
    attenuation_length = 0.15,
    allowable_coil_heating = 1e-3, 
    eta_thermal = 0.35, 
    eta_current_drive = 0.4,
) # All boundary inputs

# The guesses:
# A closure's unknown is its copy (^copy), so its start is the copy's: ^start.^copy.temperature.
place = flat.place
guesses = {place(cx.start(cx.copy('temperature'))):        12.0,
           place(cx.start(cx.copy('bootstrap_fraction'))): 0.3,
           place(cx.start('blanket_thickness')):           1.0}   # an implicit model's own unknown

design_point_varpath = {flat.place(k): v for k, v in design_point.items()}

And the running of the schedule can remain as before (but now cottax handles the topological ordering and evaluation and solving)

def run_on(jit_graph, design_point, guesses):
    return run(GraphWithValues(jit_graph, OrderedMap({**design_point, **guesses})))


def print_results(total_datastructure):
    for name in ('temperature', 'bootstrap_fraction', 'blanket_thickness',
             'fusion_power', 'electric_power', 'peak_field'):
        print(f"{name:22s} {total_datastructure[flat.place(name)]:.4g}")

print_results(run_on(jit_combined, design_point_varpath, guesses))

temperature            14.96
bootstrap_fraction     0.2395
blanket_thickness      1.085
fusion_power           1444
electric_power         397.9
peak_field             10.55

Done, works! [4]

Option 3: not combining at all

Alternatively, we can assign both FixedPoint problems a solver and nest the them:

bootstrap_current_problem = plan_2fpc.ops[0].at
temperature_problem      = plan_2fpc.ops[1].at

plan_nested = plan_2fpc + cx.absolute.assign(bootstrap_current_problem, FixedPointIteration(steps = False))
plan_nested = plan_nested + cx.absolute.assign(cx.solve(N["Shield"]), RootFinding(steps = False))
plan_nested = plan_nested + cx.absolute.assign(temperature_problem, FixedPointIteration(steps = False))
# The bootstrap problem is the outer one: every other problem on its cycle is answered
# inside its iteration, which `Nest` records on the graph, one pair at a time.
block = next(e for e in entries(plan_nested.graph) if bootstrap_current_problem in e.nodes)
for p in block.subgraph.outermost_problems:
    if p != bootstrap_current_problem:
        plan_nested = plan_nested + Nest(p, bootstrap_current_problem)
jit_nested = JitGraph((plan_nested + solvers).graph)

xdsm(jit_nested, blocks=True)

The nested route, step by step:

xdsm(plan_nested + solvers, blocks=True)

We can see we have a nested evaluation and solver.

We can re-use the above functions to run this schedule:

print_results(run_on(jit_nested, design_point_varpath, guesses))

temperature            14.96
bootstrap_fraction     0.2395
blanket_thickness      1.085
fusion_power           1444
electric_power         397.9
peak_field             10.55

Naturally, this gives the exact same results.

Summary

In this notebook, a simple cyclic tokamak model with cycles has been solved. Graph operations were used to convert these cycles to fixed point problems, and they can be assigned drivers, combined, and solved.

---

[1]

See https://pubs.aip.org/aip/pop/article/22/7/070901/264353/Designing-a-tokamak-fusion-reactor-How-does-plasma

[2]

You could argue this is a lot of effort for just a fixed point iteration on a block. That’s true; but it can be automated: a total fixed point i sjust all variables assumed (wasteful). If you don’t want to do it yourself, agentic AI can do this very well as well!

[3]

This might seem a lot of obfuscation: a blocking already gives the same information. True, but a ExecutableGraph can be build piece-by-piece and is not necessarily a validated runner. Furthermore, a JitGraph nests: a Solve step’s interior is one too.

[4]

This doesn’t JIT jet; the problem is that the VarPath PyTree index (which all dictionaries have here) are not sortable and therefore not flattenable. This is on purpose: the above function modifies total_dict in place, which can be dangerous in some cases. It would not JIT anyway. Prefer to use a boundary evaluation instead.