Sintesi Tecnica: TORAX – Un simulatore di trasporto per Tokamak veloce e differenziabile in JAX

Definizione del Problema

La simulazione accurata del trasporto integrato è fondamentale per la progettazione e l'operatività dei dispositivi di confinamento magnetico toroidale come i tokamak, servendo applicazioni che spaziano dall'interpretazione sperimentale al design di controllori. Tuttavia, gli esistenti codici di trasporto del core (es. JETTO, ASTRA, PTRANSP) affrontano un compromesso tra fedeltà fisica e velocità di simulazione; i modelli ad alta fedeltà richiedono spesso ore o giorni per simulare pochi tempi di confinamento dell'energia. Inoltre, i codici tradizionali, tipicamente scritti in Fortran o MATLAB, mancano di differenziabilità nativa, ostacolando l'ottimizzazione basata su gradienti, l'analisi di sensibilità e l'integrazione fluida con i surrogati di machine learning (ML). Sebbene alcuni codici dedicati come RAPTOR offrano velocità, essi si affidano a Jacobiani codificati manualmente, limitando l'estensibilità verso modelli ML complessi. Esiste una chiara necessità di un simulatore che combini tempi di esecuzione rapidi, differenziabilità automatica e supporto nativo per l'integrazione con il ML.

Metodologia

Gli autori presentano TORAX, un simulatore open-source del trasporto del core del tokamak implementato in Python utilizzando il framework JAX.

Equazioni Governative e Discretizzazione

TORAX risolve un insieme accoppiato di PDE (equazioni alle derivate parziali) di trasporto 1D in coordinate di flusso toroidale normalizzate ($\hat{\rho}$). Il sistema governa:

1. Trasporto del calore ionico ($T_i$)
2. Trasporto del calore elettronico ($T_e$)
3. Trasporto delle particelle elettroniche ($n_e$)
4. Diffusione della corrente (flusso poloidale $\psi$)

Le equazioni sono derivate dai momenti mediati sulle superfici di flusso delle equazioni cinetiche e dalla legge di Faraday. La discretizzazione spaziale è eseguita tramite un Metodo dei Volumi Finiti (FVM) su una griglia uniforme. Lo schema impiega una ponderazione a legge di potenza per i flussi di convezione-diffusione per passare fluidamente tra differenze centrali e upwinding in base al numero di Péclet.

Implementazione del Solutore

L'integrazione temporale utilizza il metodo theta, una media pesata tra gli schemi espliciti ed espliciti di Eulero. TORAX offre tre opzioni di solutore per gestire i sistemi algebrici non lineari risultanti:

• Solutore Lineare (Predictor-Corrector): Utilizza l'iterazione a punto fisso. Per migliorare la stabilità con coefficienti rigidi (es. da modelli ML), implementa opzionalmente il metodo Pereverzev-Corrigan, aggiungendo diffusione artificiale bilanciata dalla convezione.
• Solutore Newton-Raphson: Un metodo di ricerca delle radici basato sul gradiente. Fondamentalmente, sfrutta la differenziazione automatica di JAX per calcolare automaticamente il Jacobiano del vettore residuo, eliminando la necessità di codificare manualmente il Jacobiano.
• Solutore Optimizer: Riconfigura il residuo come una funzione di perdita minimizzata tramite algoritmi di ottimizzazione iterativa (utilizzando jaxopt).

Modelli Fisici

TORAX presenta un framework modulare per i componenti fisici:

• Geometria: Supporta i dati di equilibrio CHEASE e una geometria circolare semplificata.
• Trasporto: Include modelli costanti, il Modello del Gradiente Critico (CGM) e QLKNN, un surrogato di machine learning addestrato sul codice gyrocinetico QuaLiKiz. Il modello QLKNN è implementato nativamente in JAX utilizzando la libreria Flax.
• Sorgenti: Gestisce lo scambio di calore ione-elettrone, la potenza di fusione (parametrizzazione Bosch-Hale), il riscaldamento Ohmic e le sorgenti di particelle (gas puff, iniezione di pellet, NBI). Le sorgenti possono essere trattate come implicite o esplicite.
• Fisica Neoclassica: Utilizza il modello Sauter per la corrente bootstrap e la conducibilità neoclassica.

Architettura del Software

Il codice è strutturato in moduli distinti (Geometria, Modelli di Trasporto, Sorgenti, Stepper, Calcolatori del Passo Temporale, Stato). Utilizza la compilazione Just-In-Time (JIT) di JAX per accelerare le componenti computazionalmente intensive (modelli di sorgente, coefficienti di trasporto, algebra lineare) mantenendo il controllo del flusso in standard Python per maggiore flessibilità. Questo design permette l'esecuzione su CPU, GPU e TPU.

Risultati Chiave

Gli autori hanno verificato TORAX rispetto al codice stabilito RAPTOR, il quale è stato a sua volta validato rispetto ad altri solver di trasporto.

1. Benchmark di Stato Stazionario (ITER L-mode):

   • È stato simulato un scenario simile a ITER di 10 secondi con corrente e riscaldamento costanti.
   • I risultati hanno mostrato un eccellente accordo tra TORAX e RAPTOR per i profili di temperatura ionica/elettronica, densità, flusso poloidale e fattore di sicurezza.
   • I valori di Deviazione Radice Quadrata Media Normalizzata (NRMSD) sono stati approssimativamente dello 0,7% - 1,1% per i profili stazionari.
   • Le deviazioni transitorie della temperatura sono state limitate a circa il 2,5%.

2. Benchmark di Scenario Dinamico:

   • È stato simulato uno scenario di 300 secondi con corrente di plasma variabile nel tempo (ramp-up, flattop, ramp-down) e potenza di riscaldamento modulata.
   • TORAX e RAPTOR hanno seguito da vicino la dinamica, con valori NRMSD di circa il 3,5% - 4,0% per i profili di temperatura su scale temporali sia di confinamento dell'energia che di diffusione della corrente.

3. Performance:

   • Su un nodo CPU a 48 core, una simulazione di 80 secondi di ramp-up di ITER (usando QLKNN e Newton-Raphson) è stata completata in 22 secondi di runtime, con un tempo di compilazione di 15,6 secondi.
   • Ciò dimostra una capacità faster-than-realtime (più veloce del tempo reale) per configurazioni specifiche, sebbene gli autori notino che questa non sia una conclusione generale per tutti i fenomeni transitori.

Significato e Rivendicazioni

Il paper posiziona TORAX come uno strumento versatile progettato per accelerare la ricerca nella modellazione di scenari per tokamak, nel design di impulsi e nel controllo. La sua importanza primaria risiede in tre aree:

1. Differenziabilità: Sfruttando JAX, TORAX fornisce un residuo di PDE differenziabile. Ciò consente l'ottimizzazione basata su gradienti e l'analisi di sensibilità (tramite Analisi di Sensibilità Forward) senza la codifica manuale del Jacobiano richiesta dai precedenti codici differenziabili come RAPTOR.
2. Integrazione ML: Il supporto nativo per le reti neurali all'interno dell'ecosistema JAX facilita il perfetto accoppiamento di surrogati ML (come QLKNN) come modelli fisici, permettendo loro di essere sviluppati nello stesso linguaggio del simulatore.
3. Velocità e Flessibilità: La compilazione JIT garantisce tempi di esecuzione rapidi, mentre l'architettura modulare basata su Python semplifica l'estensibilità del codice, la manutenzione e l'accoppiamento con workflow esterni.

Gli autori concludono che TORAX offre uno "strumento potente e versatile" per la comunità della fusione, mirando specificamente ad applicazioni che richiedono sia velocità che differenziabilità, come l'incertezza quantitativa, l'ottimizzazione degli scenari e il design di controllori basati su modelli. Il lavoro futuro prevede l'espansione della fedeltà fisica (es. MHD, radiazione, specie multi-ioniche) e il potenziamento delle caratteristiche di differenziabilità (es. cache di compilazione persistenti, calcoli di sensibilità forward).