Graph Neural Networks for Coupled Subsystems

The plant's physics couples magnet, plasma, blanket, and conversion train; a GNN over that coupling graph learns cross-subsystem dynamics the twin needs.

Coupling is a graph

The machines are strongly coupled: the MHD equilibrium sets the neutron source, neutronics deposits heat, thermomechanics deforms geometry, and (on the burner) direct energy conversion sets the potential that changes confinement. Represent each subsystem as a node and each physical coupling as an edge, and the whole plant is a graph a GNN can model as a dynamical system.

Coupled-subsystem graph (nodes / edges):

  nodes: {plasma/MHD, neutronics, thermomechanics, magnets,
          power/DEC, fuel-cycle}
  edges: physical couplings, each with a transfer relation

  State update per twin step (learned message passing):
    s_v(t+1) = f_v( s_v(t), AGG_{u~v} g_{uv}( s_u(t), s_v(t) ) )

Learning the coupling operators

Rather than hand-couple full solvers every step, the GNN learns the edge transfer relations from the offline high-fidelity coupled simulations. Each edge function g_uv approximates how one subsystem's state perturbs its neighbor. The node update integrates incoming influences, giving a fast surrogate for the weakly-coupled fixed-point solve KRONOS-CTRL otherwise runs.

# coupled-subsystem step as a GNN (schematic)
def twin_step(state, graph):
    msgs = {}
    for (u, v) in graph.edges:
        msgs.setdefault(v, []).append(g_edge[u,v](state[u], state[v]))
    new = {}
    for v in graph.nodes:
        new[v] = f_node[v](state[v], aggregate(msgs.get(v, [])))
    return new     # one coupled advance, faster than nested solves

Physics guardrails

A learned coupling can violate conservation. The stack constrains the GNN to respect the invariants that must hold: energy and particle balance across edges, and sign conventions on transfers. Where a learned edge disagrees with the reference solver beyond tolerance, the twin falls back to the true coupled solve for that step. This keeps speed without abandoning conservation.

• Nodes = subsystems, edges = physical couplings with transfer laws.
• Edge functions learned from offline coupled multiphysics runs.
• Conservation and sign constraints imposed as guardrails.
• Fallback to full solve when a learned edge exceeds tolerance.

On the burner the DEC-to-confinement edge carries the extrapolation caveat: its transfer relation is drawn from a regime 166-830x beyond any device, so that edge's uncertainty dominates the coupled state and is reported.