🔬 physics

TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

Ce document présente TORAX, un simulateur de transport de tokamak open-source et différentiable, construit avec JAX, qui exploite la différenciation automatique et la compilation JIT pour permettre une optimisation rapide basée sur le gradient pour la modélisation de scénarios de plasma, la conception d'impulsions et le contrôle.

Auteurs originaux : Jonathan Citrin, Ian Goodfellow, Akhil Raju, Jeremy Chen, Jonas Degrave, Craig Donner, Federico Felici, Philippe Hamel, Andrea Huber, Dmitry Nikulin, David Pfau, Brendan Tracey, Martin Riedmiller, Pus

Publié 2026-07-10

📖 7 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Jonathan Citrin, Ian Goodfellow, Akhil Raju, Jeremy Chen, Jonas Degrave, Craig Donner, Federico Felici, Philippe Hamel, Andrea Huber, Dmitry Nikulin, David Pfau, Brendan Tracey, Martin Riedmiller, Pushmeet Kohli

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de cuire la miche de pain parfaite, mais au lieu de la farine et de l'eau, vos ingrédients sont du plasma surchauffé, des champs magnétiques et de l'énergie de fusion. L'objectif ? Construire une machine (un tokamak) capable d'alimenter le monde. Le problème ? Prédire comment ce plasma se comporte à l'intérieur de la machine, c'est comme essayer de prévoir la météo dans un ouragan pendant que la tempête est en train de se produire. C'est chaotique, c'est rapide, et cela implique un mélange vertigineux de règles physiques qui changent chaque milliseconde.

Entrez en scène TORAX, un nouveau programme informatique open-source créé par des chercheurs de Google DeepMind. Considérez TORAX comme un « simulateur de vol intelligent et ultra-rapide » pour ces réacteurs à fusion. Mais voici le twist : contrairement aux anciens simulateurs qui sont comme des camions lourds et lents, TORAX est construit avec un moteur spécial appelé JAX. Cet moteur permet au simulateur non seulement de fonctionner incroyablement vite, mais aussi d'« apprendre » de ses erreurs instantanément en calculant des gradients (en gros, en comprenant exactement quel bouton tourner pour obtenir un meilleur résultat).

L'ancienne méthode vs La nouvelle méthode

Pendant des années, les scientifiques ont utilisé d'autres simulateurs (comme RAPTOR, JETTO ou ASTRA) pour modéliser ces réacteurs. Ces outils sont comme des calculatrices fiables et robustes. Ils fonctionnent bien, mais ils sont souvent écrits dans des langages plus anciens (comme le Fortran) qui rendent difficile l'intégration de modèles modernes d'Intelligence Artificielle (IA) ou la modification rapide des calculs pour trouver les réglages parfaits. Ils sont également lents ; faire tourner une simulation de quelques secondes de temps de réacteur pourrait prendre des heures, voire des jours, sur un ordinateur standard.

Les auteurs de cet article soutiennent que cette lenteur et ce manque de flexibilité freinent la recherche sur la fusion. Ils excluent explicitement l'idée que nous devions nous contenter de ces codes anciens et rigides si nous voulons utiliser l'apprentissage automatique moderne ou effectuer des optimisations complexes. Ils démontrent que, bien que les anciens codes soient bons, ils ne peuvent pas facilement gérer les scénarios de type « et si... » nécessaires pour concevoir les futurs réacteurs ou les contrôler en temps réel.

Comment fonctionne TORAX : La magie du « simulateur de vol »

TORAX simule le cœur du plasma en résolvant un ensemble de quatre équations principales. Vous pouvez les voir comme les règles régissant :

La chaleur des ions : Comment les particules lourdes chauffent.
La chaleur des électrons : Comment les particules légères chauffent.
Le transport des particules : Comment la densité du plasma change.
La diffusion du courant : Comment le courant électrique circule à travers le plasma.

Au lieu de résoudre ces équations une par une dans une marche lente et étape par étape, TORAX utilise une technique appelée Méthode des Volumes Finis. Imaginez découper le plasma en une pile de fines crêpes invisibles (ou cellules). Le programme calcule comment la chaleur et les particules passent d'une crêpe à l'autre. Il utilise une astuce mathématique ingénieuse, le « schéma de la loi de puissance », pour décider s'il doit observer le flux venant de la gauche ou de la droite, selon la vitesse de déplacement du plasma. Cela garantit que la simulation ne devient pas « instable » ou imprécise.

La véritable magie opère parce que TORAX est écrit en Python utilisant JAX. Cela signifie que :

Vitesse : Il compile le code en une version super efficace qui s'exécute sur des CPU, des GPU ou même des TPU. Lors des tests, il a fait tourner une simulation de démarrage de réacteur de 80 secondes en seulement 22 secondes sur un ordinateur standard. C'est plus rapide que le temps réel !
Différentiabilité : C'est le point crucial. TORAX peut calculer exactement comment un minuscule changement dans l'entrée (comme augmenter la puissance de chauffage de 1 %) affecte le résultat final. C'est essentiel pour utiliser l'IA afin de concevoir automatiquement la « impulsion » (la séquence d'événements) parfaite pour un réacteur.

L'astuce du « Substitut » : L'IA comme raccourci

L'une des caractéristiques les plus passionnantes est la façon dont TORAX gère la turbulence. La turbulence dans le plasma est incroyablement complexe, comme essayer de prédire chaque tourbillon dans une rivière. Habituellement, simuler cela prend un temps infini. TORAX peut utiliser des substituts d'apprentissage automatique (ML) à la place.

Considérez un substitut ML comme une « boule de cristal » entraînée sur des millions de simulations physiques complexes. Au lieu de faire tout le travail lourd de calcul de chaque tourbillon, la boule de cristal (le modèle d'IA) prédit instantanément le résultat. TORAX est conçu pour intégrer ces modèles d'IA de manière transparente. Par exemple, il utilise un modèle appelé QLKNN, qui est entraîné sur des données provenant d'un code plus complexe appelé QuaLiKiz. Cela permet à TORX de faire tourner des simulations avec des modèles de turbulence de haute fidélité en une fraction du temps qu'il fallait auparavant.

Est-ce que cela a fonctionné ? La preuve est dans le pudding

Les auteurs ne se sont pas contentés de construire l'outil ; ils l'ont testé. Ils ont fait tourner TORAX côte à côte avec RAPTOR, un simulateur bien établi et fiable. Ils ont simulé un scénario « de type ITER » (ITER est l'expérience de fusion internationale massive en construction en France).

Les résultats sont impressionnants :

Accord : Les profils de température, de densité et de magnétisme générés par TORAX correspondent presque parfaitement à ceux de RAPTOR. La différence est infime — moins de 1 % pour les états stationnaires et environ 2,5 % lors des changements rapides.
Vitesse : Bien que RAPTOR soit rapide, TORAX a été capable de faire tourner les mêmes scénarios encore plus vite, surtout lorsqu'il utilisait les substituts d'IA.
Vérification : L'article stipule explicitement que TORAX est « vérifié » par rapport à RAPTOR, ce qui signifie que les résultats sont dignes de confiance dans les scénarios testés.

Et maintenant ? (Et ce qui ne l'est pas encore)

L'article prend soin de ne pas survendre la situation. Il admet que TORAX est un outil « nouveau » et que certains aspects sont encore sur la feuille de route.

Ce qui manque : Il ne gère pas encore les formes 3D complexes des champs magnétiques de manière dynamique (il utilise des formes pré-calculées), et il n'inclut pas encore tous les types de pertes par rayonnement ou le transport des impuretés lourdes.
Ce qui arrive : Les auteurs prévoient d'ajouter plus de modèles physiques, comme ceux pour les « dents de scie » (un type d'instabilité magnétique) et de meilleures façons de gérer le bord du plasma (le « piédestal »).

L'essentiel à retenir

TORAX n'est pas une baguette magique qui résout la fusion du jour au lendemain. C'est un outil puissant, flexible et incroyablement rapide qui rend la tâche de conception des réacteurs à fusion beaucoup plus facile. En combinant la vitesse de l'informatique moderne avec l'intelligence de l'apprentissage automatique, il permet de lancer des milliers de scénarios de type « et si... » dans le temps qu'il fallait autrefois pour en exécuter un seul. C'est comme passer d'une calculatrice à manivelle à un superordinateur capable de s'enseigner à lui-même comment résoudre le problème plus efficacement.

L'article conclut que cet outil est prêt à aider à accélérer la recherche dans la modélisation des scénarios de tokamak, la conception d'impulsions et le contrôle, tout en soulignant qu'il s'agit d'une étape vers l'avant, et non de la destination finale. L'avenir de la fusion est brillant, et des outils comme TORAX aident à éclairer le chemin.

Résumé Technique : TORAX – Un simulateur de transport pour Tokamak rapide et différentiable en JAX

Problématique

La simulation intégrée précise du plasma est cruciale pour la conception et l'exploitation des dispositifs de confinement magnétique toroïdal tels que les tokamaks, servant des applications allant de l'interprétation expérimentale à la conception de contrôleurs. Cependant, les codes de transport de cœur existants (par exemple, JETTO, ASTRA, PTRANSP) sont confrontés à un compromis entre fidélité physique et vitesse de simulation ; les modèles de haute fidélité nécessitent souvent des heures ou des jours pour simuler quelques temps de confinement de l'énergie. De plus, les codes traditionnels, généralement écrits en Fortran ou MATLAB, manquent de différentiabilité native, ce qui entrave l'optimisation basée sur le gradient, l'analyse de sensibilité et l'intégration fluide avec les substituts d'apprentissage automatique (ML). Bien que certains codes dédiés comme RAPTOR offrent de la rapidité, ils reposent sur des Jacobiennes codées manuellement, ce qui limite l'extensibilité vers des modèles de ML complexes. Il existe un besoin distinct pour un simulateur combinant des temps d'exécution rapides, une différentiabilité automatique et un support natif pour l'intégration du ML.

Méthodologie

Les auteurs présentent TORAX, un simulateur de transport de cœur de tokamak en open-source implémenté en Python en utilisant le framework JAX.

Équations de base et discrétisation

TORAX résout un ensemble couplé d'équations aux dérivées partielles (EDP) de transport 1D dans des coordonnées de flux toroïdal normalisées ( $\hat{\rho}$ ). Le système régit :

Le transport thermique des ions ( $T_i$ )
Le transport thermique des électrons ( $T_e$ )
Le transport des particules électroniques ( $n_e$ )
La diffusion du courant (flux poloidal $\psi$ )

Les équations sont dérivées des moments moyennés par surface de flux des équations cinétiques et de la loi de Faraday. La discrétisation spatiale est effectuée par une méthode de volumes finis (FVM) sur une grille uniforme. Le schéma utilise une pondération par loi de puissance pour les flux de convection-diffusion afin de passer en douceur entre le diffèrent centré et l'upwinding en fonction du nombre de Péclet.

Implémentation du solveur

L'intégration temporelle utilise la méthode theta, une moyenne pondérée entre les schémas Euler explicite et implicite. TORAX propose trois options de solveur pour gérer les systèmes algébriques non linéaires résultants :

Solveur Linéaire (Prédicteur-Correcteur) : Utilise une itération de point fixe. Pour améliorer la stabilité avec des coefficients raides (ex: issus de modèles ML), il implémente optionnellement la méthode de Pereverzev-Corrigan, ajoutant une diffusion artificielle équilibrée par la convection.
Solveur Newton-Raphson : Une méthode de recherche de racines basée sur le gradient. Crucialement, il exploite la différentiation automatique de JAX pour calculer automatiquement la Jacobienne du vecteur résiduel, éliminant ainsi le besoin de codage manuel de la Jacobienne.
Solveur d'Optimisation : Recast le résidu en une fonction de perte minimisée via des algorithmes d'optimisation itératifs (utilisant jaxopt).

Modèles Physiques

TORAX dispose d'un cadre modulaire pour les composants physiques :

Géométrie : Supporte les données d'équilibre CHEASE et une géométrie circulaire simplifiée.
Transport : Inclut un modèle de transport constant, le Modèle de Gradient Critique (CGM), et QLKNN, un substitut d'apprentissage automatique entraîné sur le code gyrocinétique QuaLiKiz. Le modèle QLKNN est implémenté nativement en JAX en utilisant la bibliothèque Flax.
Sources : Gère l'échange de chaleur ion-électron, la puissance de fusion (paramétrage Bosch-Hale), le chauffage ohmique et les sources de particules (injection de gaz, injection de pellets, NBI). Les sources peuvent être traitées de manière implicite ou explicite.
Physique Néoclassique : Utilise le modèle de Sauter pour le courant bootstrap et la conductivité néoclassique.

Architecture Logicielle

Le code est structuré en modules distincts (Géométrie, Modèles de Transport, Sources, Steppers, Calculateurs de Pas de Temps, État). Il utilise la compilation Just-In-Time (JIT) de JAX pour accélérer les composants de calcul intensif (modèles de sources, coefficients de transport, algèbre linéaire) tout en gardant le contrôle du flux en Python standard pour la flexibilité. Cette conception permet une exécution sur CPU, GPU et TPU.

Résultats Clés

Les auteurs ont vérifié TORAX par rapport au code établi RAPTOR, qui a lui-même été validé par d'autres solveurs de transport.

Benchmark d'état stationnaire (ITER L-mode) :
- Simulation de 10 secondes d'un scénario de type ITER avec courant et chauffage constants.
- Les résultats montrent un excellent accord entre TORAX et RAPTOR pour les profils de température ionique/électronique, de densité, de flux poloidal et de facteur de sécurité.
- Les valeurs de l'erreur quadratique moyenne normalisée (NRMSD) sont d'environ 0,7 % à 1,1 % pour les profils stationnaires.
- Les déviations transitoires de température sont limitées à environ 2,5 %.
Benchmark de scénario dynamique :
- Simulation d'un scénario de 300 secondes avec courant de plasma variable (montée en courant, plateau, descente en courant) et puissance de chauffage modulée.
- TORAX et RAPTOR suivent étroitement la dynamique, avec des valeurs NRMSD de ~3,5 % à 4,0 % pour les profils de température sur les échelles de temps de confinement de l'énergie et de diffusion du courant.
Performance :
- Sur un nœud CPU de 48 cœurs, une simulation de 80 secondes de montée en courant d'ITER (utilisant QLKNN et Newton-Raphson) s'est terminée en 22 secondes de temps d'exécution, avec un temps de compilation de 15,6 secondes.
- Cela démontre une capacité de fonctionnement plus rapide que le temps réel pour des configurations spécifiques, bien que les auteurs notent que ceci n'est pas une conclusion générale pour tous les phénomènes transitoires.

Signification et Revendications

L'article positionne TORAX comme un outil polyvalent conçu pour accélérer la recherche dans la modélisation de scénarios de tokamak, la conception de pulses et le contrôle. Sa signification primaire réside dans trois domaines :

Différentiabilité : En exploitant JAX, TORAX fournit un résiduel d'EDP différentiable. Cela permet l'optimisation basée sur le gradient et l'analyse de sensibilité (via l'Analyse de Sensibilité Directe) sans le codage manuel de la Jacobienne requis par les codes précédents comme RAPTOR.
Intégration du ML : Le support natif des réseaux de neurones au sein de l'écosystème JAX facilite le couplage fluide des substituts de ML (comme QLKNN) en tant que modèles physiques, permettant de les développer dans le même langage que le simulateur.
Vitesse et Flexibilité : La compilation JIT assure des temps d'exécution rapides, tandis que l'architecture modulaire basée sur Python simplifie l'extensibilité du code, la maintenance et le couplage avec des flux de travail externes.

Les auteurs concluent que TORAX offre un "outil puissant et polyvalent" pour la communauté de la fusion, ciblant spécifiquement les applications qui nécessitent à la fois de la vitesse et de la différentiabilité, telles que la quantification d'incertitude, l'optimisation de scénarios et la conception de contrôleurs basés sur des modèles. Des travaux futurs sont prévus pour étendre la fidélité physique (ex: MHD, rayonnement, espèces multi-ions) et améliorer les fonctionnalités de différentiabilité (ex: caches de compilation persistants, calculs de sensibilité directe).