TokaMark: A Comprehensive Benchmark for MAST Tokamak Plasma Models

Este artículo presenta TokaMark, un benchmark integral y de código abierto diseñado para estandarizar la evaluación de modelos de inteligencia artificial en la dinámica del plasma del tokamak MAST, abordando la falta de datos unificados y facilitando comparaciones justas para acelerar el desarrollo de la energía de fusión.

Lo esencial

Imagina que intentar predecir el comportamiento del Sol es como intentar adivinar qué hará un toro salvaje en medio de una plaza, pero sin poder tocarlo, solo observándolo a través de ventanas pequeñas, con gafas de sol y con los datos llegando a diferentes velocidades. Eso es, básicamente, lo que hacen los científicos con los tokamaks: máquinas gigantes que intentan recrear la energía del Sol en la Tierra para tener una fuente de energía limpia e infinita.

El problema es que estos "soles artificiales" son extremadamente complejos, caóticos y difíciles de entender. Los sensores que los monitorean a veces fallan, a veces dan datos ruidosos y a veces hablan idiomas diferentes (algunos miden temperatura, otros campos magnéticos, otros presión).

Aquí es donde entra TokaMark, el tema de este artículo.

¿Qué es TokaMark? (El "Simulador de Videojuegos" de la Fusión)

Piensa en TokaMark como un gran videojuego de entrenamiento o un "campo de pruebas" estandarizado para la Inteligencia Artificial (IA).

Antes de este trabajo, si un científico quería crear una IA para predecir el comportamiento del plasma (el gas supercaliente dentro de la máquina), tenía que:

1. Buscar sus propios datos, que estaban escondidos en diferentes instituciones.
2. Limpiar esos datos manualmente (como ordenar una habitación llena de juguetes mezclados).
3. Inventar sus propias reglas para ver si su IA funcionaba bien.

Esto hacía que fuera muy difícil comparar quién era el mejor. ¿Era la IA del laboratorio A mejor que la del laboratorio B? Nadie lo sabía con certeza porque usaban reglas diferentes.

TokaMark soluciona esto creando un "código de vestimenta" y un "campo de juego" único para todos.

¿Qué hace exactamente TokaMark?

TokaMark es una caja de herramientas que ofrece tres cosas principales:

1. Los Datos (La "Materia Prima"): Han recopilado y organizado miles de experimentos reales de una máquina llamada MAST (en el Reino Unido). Han limpiado los datos, puesto etiquetas en ellos y los han hecho públicos para que cualquiera pueda usarlos. Es como si alguien hubiera abierto una biblioteca gigante de recetas de cocina y las hubiera puesto en orden alfabético para que todos los chefs las usen.
2. Las Pruebas (Los "Niveles del Juego"): Han diseñado 14 misiones diferentes para poner a prueba a las IAs. Estas misiones van desde lo fácil a lo difícil:
   • Nivel 1 (Reconstrucción): "Mira estos sensores magnéticos y dime qué forma tiene el plasma ahora mismo".
   • Nivel 2 (Dinámica Rápida): "Si cambio la corriente en este cable, ¿cómo reaccionará el plasma en los próximos milisegundos?".
   • Nivel 3 (Evolución Lenta): "Basándote en lo que ha pasado en los últimos segundos, ¿cómo cambiará la temperatura del plasma en el futuro?".
   • Nivel 4 (El "Jefe Final"): "¿Puedes predecir si el plasma va a explotar o desestabilizarse antes de que ocurra?". Esto es vital para la seguridad.
3. Las Reglas de Puntuación (El "Marcador"): Han creado un sistema de puntuación justo. No importa si tu IA es rápida o lenta; lo importante es qué tan bien predice la realidad. Si tu IA acierta, gana puntos. Si falla, pierde. Así, todos compiten en igualdad de condiciones.

¿Por qué es importante esto?

Imagina que quieres construir un coche autónomo. No puedes simplemente ponerle un motor y esperar que funcione. Necesitas un circuito de pruebas donde puedas chocar contra cosas (de forma segura) y ver qué tan bien aprende el coche a evitar accidentes.

TokaMark es ese circuito de pruebas para la energía de fusión.

• Antes: Cada científico construía su propio coche y su propia pista. Era un caos.
• Ahora: Todos usan la misma pista (TokaMark). Si una IA aprende a conducir bien aquí, sabemos que es buena.

El Modelo de "Referencia" (El "Novato")

Los autores también crearon un modelo de IA básico (un "novato") para usarlo como punto de comparación. Es como poner un coche estándar en la pista. Si un nuevo investigador crea una IA y su coche va más rápido que el estándar, ¡sabemos que ha hecho un buen trabajo!

En resumen

TokaMark es el gran paso adelante para que la Inteligencia Artificial ayude a domar al "Sol en una caja". Al dar a todos los científicos los mismos datos, las mismas pruebas y las mismas reglas, aceleran el descubrimiento de cómo hacer que la fusión nuclear sea una realidad segura y comercial.

Es como si, en lugar de que cada inventor de aviones intentara volar en su propio garaje con sus propias reglas, todos volaran en el mismo aeropuerto con el mismo controlador de tráfico. Así, descubrimos más rápido cómo llegar a las estrellas (o en este caso, a la energía infinita).

Resumen técnico

1. El Problema

El desarrollo de reactores de fusión nuclear comercialmente viables (como los tokamaks) requiere predicciones precisas de la dinámica del plasma basadas en lecturas de sensores que a menudo son escasas, ruidosas e incompletas.

• Desafíos actuales: Los métodos numéricos tradicionales basados en primeros principios (ecuaciones diferenciales parciales) son computacionalmente costosos y difíciles de usar en tiempo real. Por otro lado, aunque los enfoques de IA prometedores han demostrado resultados en tareas específicas, carecen de conjuntos de datos estandarizados y abiertos.
• Barreras: Los datos de fusión existentes están fragmentados entre instituciones, tienen formatos propietarios, son inconsistentes en su anotación y son difíciles de acceder para investigadores de aprendizaje automático. Esto impide la comparación justa, la reproducibilidad y el desarrollo de modelos de IA generalistas capaces de entender el plasma de manera integral.

2. Metodología: TokaMark

Para abordar estas carencias, los autores presentan TokaMark, el primer benchmark grande y abierto diseñado para evaluar modelos de IA entrenados con datos experimentales reales del tokamak MAST (Mega Ampere Spherical Tokamak).

A. Datos y Taxonomía

• Fuente: Utilizan datos del conjunto FAIR-MAST, que contiene 11,573 disparos (shots) experimentales.
• Selección: Se seleccionaron 39 señales heterogéneas de múltiples modalidades (magnéticas, cinéticas, radiativas, corrientes, voltajes, etc.).
• Taxonomía de Señales: Las señales se organizan por:
  • Origen: Diagnósticos (mediciones directas), Actuadores (parámetros controlables) y Señales derivadas (reconstruidas, como perfiles de equilibrio).
  • Frecuencia: Desde 0.2 kHz hasta 500 kHz.
  • Modalidad: Series temporales (1D), perfiles (2D) y mapas/video (3D).
• Preprocesamiento: Se estableció una división estricta de datos (80% entrenamiento, 10% validación, 10% prueba) a nivel de disparo para evitar fugas de información. Se implementó un manejo robusto de datos faltantes (NaN) sin descartar muestras valiosas.

B. Definición de Tareas (14 Tareas en 4 Grupos)

El benchmark define 14 tareas de aprendizaje supervisado y auto-supervisado agrupadas en cuatro categorías que reflejan los desafíos físicos reales:

1. Reconstrucción de Equilibrio (Grupo 1): Inferir la forma, límites y parámetros del equilibrio del plasma a partir de mediciones magnéticas instantáneas (ej. reconstrucción de la superficie de flujo poloidal).
2. Dinámica Magnética (Grupo 2): Predicción a corto plazo de señales magnéticas y corrientes en respuesta a comandos de actuadores (control de posición y forma).
3. Dinámica de Perfiles (Grupo 3): Modelado de la evolución temporal de perfiles cinéticos (densidad y temperatura) y transiciones de modo de confinamiento.
4. Actividad MHD y Disrupciones (Grupo 4): Predicción a largo plazo de eventos críticos e inestabilidades (como quenches térmicos, desplazamientos verticales y modos bloqueados) para la prevención de disrupciones.

C. Protocolo de Evaluación Jerárquico

Se introduce un protocolo de evaluación que agrupa los errores en varios niveles para proporcionar una visión completa:

• Niveles: Muestras $\rightarrow$ Ventanas $\rightarrow$ Señales $\rightarrow$ Tareas $\rightarrow$ Disparos.
• Métrica Principal: Error Normalizado de Raíz Cuadrática Media (NRMSE). Esto permite comparar el error de predicción con la variabilidad natural de la señal, haciendo las métricas comparables entre diferentes señales físicas.
• Ventanas: Se utilizan ventanas deslizantes con longitudes de entrada y salida específicas para cada tarea (desde 5 ms para dinámicas rápidas hasta 100 ms para eventos MHD).

D. Modelo Baseline

Se proporciona un modelo de referencia de codificador-decodificador convolucional multi-rama:

• Arquitectura: Utiliza codificadores convolucionales independientes (1D, 2D o 3D) según la modalidad de entrada, fusionados en un espacio latente compartido, y decodificadores específicos para cada objetivo.
• Propósito: Establecer una línea base reproducible y accesible para que la comunidad pueda comparar sus nuevos modelos.

3. Contribuciones Clave

1. Diseño del Benchmark: Definición de 14 tareas estructuradas con un protocolo de evaluación jerárquico y métricas estandarizadas.
2. Empaquetado de Datos: Liberación de un subconjunto estable de datos FAIR-MAST con metadatos normalizados, unidades consistentes y solucionadas las inconsistencias de esquema.
3. Herramientas y API: Un paquete Python integrado con PyTorch para la carga de datos, procesamiento, enmascaramiento, alineación y lógica de evaluación.
4. Modelos Baseline: Liberación de configuraciones y scripts de entrenamiento para el modelo convolucional multi-rama, estableciendo un punto de referencia reproducible para todas las tareas.

4. Resultados

Los resultados del modelo baseline en TokaMark revelan:

• Rendimiento por Grupo: El modelo logra un buen rendimiento en tareas de reconstrucción de equilibrio (Grupo 1) y dinámica magnética (Grupo 2), con puntuaciones de grupo (NRMSE) por debajo de 0.17.
• Dificultad Creciente: El rendimiento disminuye significativamente en la dinámica de perfiles (Grupo 3) y, especialmente, en las actividades MHD (Grupo 4).
• Análisis de Errores:
  • La reconstrucción de la frontera del plasma (Tarea 1-2) y su predicción (Tarea 2-2) son las tareas mejor resueltas.
  • La tarea 4-5 (predicción de diagnósticos Mirnov para modos bloqueados) presenta el peor rendimiento (NRMSE > 1.0), indicando que las señales correspondientes están mal restringidas o son extremadamente difíciles de predecir con esta arquitectura genérica.
• Interpretación: Estos resultados no solo evalúan el modelo, sino que destacan la dificultad intrínseca del benchmark, especialmente para tareas que requieren integrar contexto temporal de largo alcance y física no lineal compleja.

5. Significado e Impacto

• Unificación de Comunidades: TokaMark actúa como un puente entre la comunidad de fusión nuclear y la de inteligencia artificial, proporcionando un terreno común para la investigación.
• Aceleración de la Investigación: Al estandarizar las tareas y los datos, permite comparaciones justas entre diferentes algoritmos, fomentando la reproducibilidad y evitando la duplicación de esfuerzos en el preprocesamiento de datos.
• Hacia Modelos Generalistas: El benchmark está diseñado para probar la capacidad de los modelos para aprender representaciones transferibles del plasma, un paso necesario hacia el desarrollo de "Modelos Fundamentales" (Foundation Models) para la fusión.
• Objetivo Final: Contribuir al desarrollo de modelos de control y predicción más rápidos y precisos, esenciales para la operación segura y estable de futuros reactores de fusión comercial.

El código, los datos y la documentación de TokaMark se han hecho de código abierto para fomentar la adopción y la contribución de la comunidad científica.