Managing complexity for integrated fusion reactor design using a differentiable graph-based MDAO framework: PROCESS as a case study

Abstract

Integrated design of nuclear fusion reactors requires combining an increasingly large number of models from different disciplines self-consistently, with model selection and parametrisation changing between design studies. Manual definition of execution order and solvers for cyclic dependencies is error-prone, inefficient, and inflexible. We propose to use a generic graph-based coordination and execution core that enables automated execution order, cycle detection, and solver architecture specification for any set of models or parametrisation. To this end, a computational dependency graph of purely functional models is extended by a conditional node and nesting edges, treating solvers and optimisers in the same way. A minimal set of graph operations then allows for fine-tuned control on feedback resolution, making both optimiser and solver architectures easily reconfigurable. Existing monolithic multidisciplinary design analysis and optimisation (MDAO) architectures can be recovered by sequentially applying these operations. The cottax package implements this extended dependency graph, graph operations, executability checks, and a JAX-based execution layer. This execution layer eliminates runtime overhead and provides native gradient computation through the whole graph. As a proof of principle, cottax is then applied to the PROCESS systems code. Graph sequencing shows that for PROCESS only three small coupled components need to be solved every optimiser iteration. Monolithic MDAO architectures are then assigned to seven different PROCESS configurations, demonstrating the ease of (re)configuring its architecture. Compared to PROCESS, the sequencing and execution back-end together yield a two to three orders of magnitude reduction in model runtime during optimisation. cottax is currently being considered as the basis for developing an integrated stellarator MDAO framework.