🔢 mathematics

JAX-FVM: A differentiable, entropy-stable finite volume solver on unstructured meshes for compressible flows

Este artículo presenta JAX-FVM, un resolvedor de volúmenes finitos de código abierto y totalmente diferenciable escrito íntegramente en JAX que permite la diferenciación automática de extremo a extremo para flujos compresibles en mallas triangulares no estructuradas, cerrando así una brecha crítica dejada por los marcos existentes restringidos a mallas estructuradas o métodos espectrales.

Autores originales: Guillaume de Romémont

Publicado 2026-07-09

📖 5 min de lectura🧠 Análisis profundo

CC0 1.0

Autores originales: Guillaume de Romémont

Artículo original dedicado al dominio público bajo CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando predecir cómo el aire fluye alrededor de una forma compleja, como un coche de carreras o el ala de un pájaro. Durante décadas, los científicos han utilizado potentes programas informáticos llamados solucionadores CFD para simular esto. Normalmente, estos programas funcionan como una calle de un solo sentido: les das una forma y algunas reglas, y ellos te dicen qué hace el viento.

Pero, ¿y si quisieras trabajar hacia atrás? ¿Qué tal si quisieras preguntar: "¿Cómo cambio la forma para que el viento sea más suave?" o "¿Qué fuerzas invisibles están causando esta turbulencia?". Para responder a estas preguntas, el ordenador necesita ser capaz de rastrear sus pasos hacia atrás, calculando exactamente cómo pequeños cambios en la entrada afectan al resultado final. Esto se llama diferenciabilidad.

Hasta ahora, la mayoría de los solucionadores "inteligentes" que podían hacer este rastreo hacia atrás solo funcionaban en patrones simples y con cuadrícula (como un papel milimetrado). Sin embargo, las formas del mundo real son desordenadas e irregulares. Necesitan mallas no estructuradas, que son como un mosaico de triángulos irregulares que encajan perfectamente alrededor de curvas complejas. ¿El problema? Estas formas irregulares son muy difíciles de manejar para los chips informáticos modernos y superrápidos (GPU), y aún más difíciles de hacer "rastreables hacia atrás".

Presentamos JAX-FVM.

Piensa en JAX-FVM como un traductor nuevo y superinteligente que habla dos lenguajes con fluidez:

El lenguaje de las formas desordenadas del mundo real (usando esos mosaicos triangulares irregulares).
El lenguaje de la IA moderna y el aprendizaje automático (usando una herramienta llamada JAX que permite el "rastreo hacia atrás" automático e instantáneo).

Así es como funciona, utilizando algunas analogías de la vida cotidiana:

1. La "Malla Inteligente" (Triángulos no estructurados)

Imagina que intentas cubrir el suelo de un baño con baldosas. Si el suelo es un cuadrado perfecto, usas baldosas cuadradas. Pero si el baño tiene curvas y esquinas extrañas, necesitas cortar las baldosas en todo tipo de formas irregulares para que encajen.
JAX-FVM utiliza estas "baldosas" irregulares (triángulos) para mapear formas complejas. La mayoría de las otras herramientas de "rastreo hacia atrás" se niegan a usar estas baldosas irregulares porque son demasiado desordenadas para que el ordenador las procese rápidamente. JAX-FVM, sin embargo, está construido específicamente para manejar este desorden de manera eficiente.

2. La "Red de Seguridad de Entropía Estable"

Cuando el aire se mueve muy rápido (como en un motor de un avión de combate), puede crear saltos repentinos y violentos llamados choques. Si tu simulación informática no es cuidadosa, estos choques pueden hacer que los números se vuelvan locos y la simulación colapse, como un coche perdiendo el control en una carretera helada.
JAX-FVM utiliza una red de seguridad matemática especial llamada estabilidad de entropía. Piensa en esto como un sistema de dirección autocorrectivo. No importa cuán salvaje se vuelva el viento, las matemáticas aseguran que la simulación se mantenga estable y no explote, tal como un coche con control de estabilidad avanzado que se mantiene en la carretera incluso durante un derrape.

3. El "Rastreo hacia Atrás Instantáneo" (Diferenciabilidad)

Este es el truco de magia. En el pasado, si querías saber cómo un pequeño cambio en la forma afectaría al viento, tenías que escribir manualmente fórmulas matemáticas complejas para averiguarlo. Era como intentar calcular manualmente la trayectoria exacta de una pelota lanzada en un huracán.
JAX-FVM está construido de modo que cada uno de los pasos del cálculo es "diferenciable". Esto significa que el ordenador puede calcular instantáneamente el camino hacia atrás de forma automática. Si preguntas: "¿Qué pasa si muevo la velocidad del viento un 1%?", el solucionador te dice la respuesta al instante sin necesidad de que un humano reescriba el código.

4. El "Motor Universal" (CPU y GPU)

Normalmente, el software se construye para ordenadores estándar (CPU) o chips de juego superrápidos (GPU). JAX-FVM es como un coche híbrido que funciona perfectamente en ambos. Puede procesar números en un portátil estándar o acelerarse masivamente en una potente tarjeta gráfica, sin necesidad de cambiar el código.

¿Qué puede hacer?

El artículo muestra que JAX-FVM puede simular con éxito pruebas famosas y difíciles de túneles de viento, tales como:

El Escalón Frontal: Imagina el viento golpeando un escalón repentino en el suelo. El solucionador predice correctamente las complejas ondas de choque que rebotan en él.
La Reflexión de Doble Mach: Un escenario donde las ondas de choque golpean una pared y rebotan entre sí en un patrón complejo.

¿Por qué es esto importante?

Los autores explican que esta herramienta llena un vacío específico. Es la primera herramienta de acceso abierto que combina formas desordenadas del mundo real con rastreo hacia atrás compatible con la IA para aire de alta velocidad.

Está diseñado como una plataforma de investigación. Permite a los científicos:

Probar nuevas ideas sobre cómo fluye el aire.
Entrenar modelos de aprendizaje automático para comprender mejor la física.
Resolver "problemas inversos" (figurar la causa a partir del efecto) para cuestiones de aerodinámica.

En resumen, JAX-FVM es un nuevo tipo de simulador de túnel de viento construido para el futuro de la IA, capaz de manejar las formas irregulares y desordenadas del mundo real, manteniendo al mismo tiempo la estabilidad matemática y la reversibilidad de los cálculos.

Resumen Técnico: JAX-FVM

Planteamiento del Problema

La Dinámica de Fluidos Computacional (CFD) se ha utilizado tradicionalmente para simulaciones directas, pero muchas aplicaciones críticas —como la optimación de formas, la identificación de parámetros y el modelado de la turbulencia— requieren resolver problemas inversos. Estas tareas necesitan computar los gradientes de las salidas del solver con respecto a las entradas (por ejemplo, condiciones de contorno, parámetros físicos o estados iniciales). Aunque los métodos adjuntos han abordado esto históricamente, son intrusivos al código y a menudo están desacoplados de los ecosistemas modernos de aprendizaje automático (ML).

Los desarrollos recientes en programación diferenciable (por ejemplo, JAX, PyTorch) han permitido la diferenciabilidad de extremo a extremo, permitiendo que los gradientes se computen mediante diferenciación automática (AD) sin necesidad de adjuntos derivados a mano. Sin embargo, los marcos de trabajo de CFD diferenciables existentes (como JAX-CFD, JAX-Fluids y ΦFlow) están predominantemente restringidos a mallas cartesianas estructuradas o discretizaciones espectrales. Esto crea una brecha significativa: las mallas no estructuradas, que son esenciales para manejar dominios geométricamente complejos en la CFD industrial, permanecen mayoritariamente ausentes del ecosistema de simulación diferenciable debido a su conectividad irregular, la cual entra en conflicto con los modelos de ejecución vectorizada de los aceleradores. Además, simular de forma robusta flujos compresibles en mallas no estructuradas requiere esquemas de estabilidad de entropía para prevenir la inestabilidad no lineal y las oscilaciones espurias cerca de los choques, una característica que suele ser difícil de implementar dentro de un marco de trabajo totalmente diferenciable.

Metodología

El artículo presenta JAX-FVM, un solver de método de volúmenes finitos (FVM) totalmente diferenciable y de código abierto para las ecuaciones de Euler y Navier-Stokes compresibles en dos dimensiones sobre mallas triangulares no estructuradas. El solver está escrito íntegramente en JAX, lo que garantiza que cada operación —desde la conectividad de la malla y la evaluación de flujos hasta el limitador de pendiente y la integración temporal— sea compilada justo a tiempo (JIT), vectorizada y diferenciable.

Discretización Numérica

Ecuaciones Gobernantes: El solver aborda la forma conservativa de las ecuaciones de Navier-Stokes compresibles 2D, cerradas por la ley de los gases ideales.
Representación de la Malla: El dominio se triangula utilizando el generador MeshPy/Triangle. La estructura de datos de la malla almacena vértices, baricentros de celdas, áreas y conectividad de caras, soportando condiciones de contorno periódicas y diversos marcadores de condiciones de contorno.
Formulación Semi-discreta: La actualización centrada en la celda se deriva mediante el teorema de la divergencia. El residuo se computa como la suma de los flujos numéricos sobre las caras de las celdas.
Reconstrucción de Gradientes: Se logra una precisión de segundo orden mediante la reconstrucción MUSCL de variables primitivas. Los gradientes de celda se computan mediante un ajuste de mínimos cuadrados ponderados (LSQ), resuelto mediante un solver lineal vectorizado (jax.vmap). El limitador de pendiente se aplica utilizando ya sea el limitador de Venkatakrishnan o el clásico limitador minmod.
Flujos Numéricos: El núcleo del esquema es una formulación de estabilidad de entropía.
- El flujo convectivo base es un flujo de dos puntos conservativo de la entropía (Tadmor/Ismail-Roe).
- La disipación se añade mediante un término de disipación de tipo Rusanov (Lax-Friedrichs local) o Roe de variables de entropía.
- Crucialmente, la acción del Jacobiano requerida para mapear los saltos de variables de entropía a variables conservativas se computa exactamente usando AD de modo directo (jax.jvp) en lugar de una derivación analítica manual.
- Los flujos viscosos se evalúan utilizando estados primitivos promediados en la cara y gradientes LSQ.
Integración Temporal: El solver soporta integradores explícitos (RK2, RK3, RK4) e integradores implícitos libres de matriz (Backward Euler, SDIRK2). Las etapas implícitas resuelven sistemas no lineales mediante iteraciones de Newton, donde los productos Jacobiano-vector se computan de forma libre de matriz usando AD de modo directo y se pasan a un solver GMRES.

Arquitectura de Software

Diferenciabilidad: Todo el residuo y la lógica de paso de tiempo están envueltos en jax.jit. Se utiliza AD de modo directo internamente para la disipación de entropía y los Jacobianos del solver implícito, mientras que el AD de modo inverso permite la propagación de gradientes para cualquier función escalar hacia las entradas.
Agnosticismo de Hardware: El código se ejecuta de forma transparente en CPUs o GPUs (vía JAX con soporte CUDA).
Integración: La librería depende de flax y optax, lo que facilita la construcción de bucles de optimización y diseño inverso directamente sobre el solver.

Contribuciones Clave

Primer FVM No Estructurado Diferenciable de Flujo Compresible de Acceso Abierto: JAX-FVM llena el vacío dejado por las herramientas de CFD diferenciables existentes que están restringidas a mallas estructuradas. Es, según los autores, uno de los primeros solvers diferenciables para las ecuaciones de Euler/Navver-Stokes compresibles que opera nativamente sobre mallas triangulares no estructuradas.
Estabilidad de Entropía vía AD: La implementación integra la teoría de flujo conservativo de la entropía con la disipación de variables de entropía. Al usar AD para computar las acciones de Jacobiano necesarias para los términos de disipación, el solver logra una estabilidad de entropía demostrable sin la derivación manual de Jacobianos complejos.
Diferenciabilidad de Extremo a Extremo: El solver permite la computación directa de gradientes de funcionales de flujo con respecto a condiciones iniciales, datos de contorno, parámetros físicos (viscosidad, conductividad, $\gamma$ ) o modelos neuronales embebidos.
Plataforma de Investigación: El código proporciona una implementación compacta orientada a la investigación que combina geometría no estructurada, numeración consistente con la entropía y AD basado en JAX, sirviendo como base para problemas inversos basados en gradientes y modelado híbrido de física/ML.

Resultados y Verificación

El artículo valida JAX-FVM a través de una suite de casos de verificación estándar:

Escalón Frontal Mach 3: Un problema de choque transitorio canónico que involucra un choque de proa desprendido, un vástago de Mach, choques reflejados y una línea de deslizamiento. El solver captura exitosamente estas características en una triangulación totalmente no estructurada utilizando el flujo conservativo de la entropía, reconstrucción MUSCL con limitación de Venkatakrishnan y paso de tiempo explícito RK2.
Benchmarks Adicionales: La distribución incluye 19 configuraciones de Riemann bidimensionales (Lax y Liu), reflexión de Mach doble, inestabilidad de Kelvin-Helmholtz, vórtice de Taylor-Green y varios problemas de vórtice de bajo número de Mach.
Diagnósticos: El solver monitorea la entropía total, la energía cinética, la enstrofia, la palinstrofia, la vorticidad y el número de Mach para verificar las propiedades de preservación de estructura.

Significación y Reivindicaciones

El artículo posiciona a JAX-FVM como un puente entre la CFD de alta fidelidad en geometrías complejas y la programación diferenciable moderna. Su significación radica en permitir el diseño inverso basado en gradientes y la optimización con restricciones de PDE para flujos compresibles en mallas no estructuradas, una capacidad previamente no disponible en marcos de trabajo diferenciables de código abierto.

Los autores afirman que, al combinar la capacidad de mallas no estructuradas con numeración estable de la entropía y la diferenciabilidad de extremo a extremo, JAX-FVM crea una plataforma natural para:

Recuperar parámetros a partir de datos dispersos (por ejemplo, resistencia vascular).
Optimación de forma basada en presiones superficiales.
Aprender cierres de turbulencia o discretizaciones basadas en datos de una manera consistente con el solver.

El trabajo es modesto en su alcance, reconociendo limitaciones actuales: está restringido a mallas triangulares 2D (con una extensión a 3D planeada), depende del generador MeshPy (con un lector VTK planeado) y actualmente carece de modelos de turbulencia (los cierres RANS o LES son trabajo futuro). El solver se presenta principalmente como una herramienta de investigación para facilitar el desarrollo de modelos híbridos de física/ML y flujos de trabajo de diseño inverso.