Systematic LLM Translation of Legacy Scientific Code to Differentiable Frameworks: Application to a Land Surface Model

Resumen Técnico: Traducción Sistemática de Código Científico Heredado a Marcos Diferenciables mediante LLM

Declaración del Problema
Los Modelos del Sistema Tierra (ESM) heredados, particularmente los modelos de superficie terrestre (LSM), son críticos para la política y el pronóstico climático, pero permanecen en gran medida inaccesibles para las técnicas modernas de programación diferenciable. Estos modelos, que a menudo constan de cientos de miles de líneas de código Fortran, sufren de una pobre modularidad, componentes obsoletos y una fuerza laboral de desarrolladores en disminución. Si bien la programación diferenciable permite una estimación de parámetros, análisis de sensibilidad y asimilación de datos basados en gradientes de manera eficiente, la migración de bases de código heredadas a marcos diferenciables (como JAX) sigue siendo un desafío sin resolver. Los intentos actuales de modernización de código o bien implican reescrituras limpias que descartan la física heredada, o fallan al no poder escalar más allá de funciones aisladas debido a las dependencias entre archivos y la falta de fidelidad numérica en las traducciones de Modelos de Lenguaje Extensos (LLM). Además, los benchmarks actuales de LLM no miden adecuadamente la validez científica o la paridad numérica requerida para modelos físicos.

Metodología: Un Pipeline Agéntico de Cinco Fases
Los autores proponen un pipeline agéntico sistemático de cinco fases para traducir el modelo de superficie terrestre en Fortran CLM-ml-v2, de 19,000 líneas y 102 módulos, a una implementación en JAX totalmente diferenciable llamada clm-ml-jax. El pipeline se basa en tres principios de diseño fundamentales: análisis de dependencias estáticas, pruebas de paridad numérica modulares y verificación exhaustiva de gradientes.

1. Fase 1: Alcance y Análisis de Dependencias: Para prevenir las alucinaciones de los LLM respecto a los tipos de datos en bases de código acopladas, los autores realizan un análisis estático del grafo de dependencias de módulos completo. Un ordenamiento topológico determina un orden de traducción que respeta las dependencias, asegurando que todos los módulos importados se traduzcan antes de los módulos que dependen de ellos.
2. Fase 2: Configuración de Documentación de Estado: El pipeline utiliza documentos de estado persistentes (CLAUDE.md, plan.md, CHANGELOG.md) para servir como memoria externa del agente. Estos documentos codifican las convenciones del proyecto, rastrean las banderas de finalización y registran intentos fallidos, lo que permite la reanudación coherente de sesiones a través de los límites de la ventana de contexto y evita que los agentes vuelvan a caer en callejones sin salida.
3. Fase 3: Construcción del Oráculo de Fortran: Antes de la traducción, un equipo de agentes construye una suite de pruebas funcionales. Las subrutinas se clasifican en tres niveles según su capacidad de prueba. El sistema ejecuta el código Fortran original bajo entradas fijas para generar pares de entrada-salida "dorados", estableciendo un oráculo de referencia para la validación numérica.
4. Fase 4: Traducción, Prueba y Reparación Iterativa: Los módulos se procesan en el orden topológico establecido en la Fase 1. Un agente traductor genera código JAX, que luego es validado por un agente de pruebas contra el oráculo de Fortran (requiriendo una tolerancia relativa dentro de $10^{-4}$). Si una prueba falla, se activa un "Bucle Ralph" (un patrón de orquestación), que emplea un agente de reparación para diagnosticar discrepancias, editar el código y volver a probar de forma autónoma hasta que el módulo pase o sea marcado para revisión humana.
5. Fase 5: Integración y Validación de Diferenciabilidad: Una vez traducidos los módulos, un Agente de Integración ensambla todo el pipeline de la columna y valida el comportamiento de extremo a extremo. Un Agente de Diferenciabilidad luego resuelve iterativamente problemas específicos de JAX, tales como:
   • Desenrollado de Bucles Estáticos: Reemplazar los bucles for a nivel de Python con jax.lax.scan para trazar los cuerpos de los bucles una sola vez en tiempo de compilación, evitando grafos de computación inmanejables para los resolvedores de Runge-Kutta.
   • Tipado Estático vs. Dinámico: Convertir constantes en escalares de JAX para evitar la recompilación completa causada por cambios en las claves de caché.
   • Salvaguardas de Ramificación Segura: Aplicar límites inferiores (por ejemplo, jnp.maximum(x, 1e-30)) para prevenir gradientes NaN en operaciones como $0 \times \infty$ en expresiones jnp.where.
   • Adjuntos del Teorema de la Función Implícita (IFT): Diferenciar los resolvedores iterativos (por ejemplo, la bisección de la conductancia estomática, la longitud de Monin–Obukhov) utilizando adjuntos de IFT en lugar de desenrollar el bucle, lo cual produce gradientes incorrectos.

Resultados Clave
El modelo resultante clm-ml-jax fue evaluado en el sitio de AmeriFlux CHATS7 (huerto de nogales) utilizando datos meteorológicos de mayo de 2007.

• Equivalencia Numérica: El modelo traducido logró una paridad numérica casi perfecta con la referencia de Fortran a través de 1488 pasos de tiempo, con coeficientes de correlación de 1.000 y un sesgo negligible para todas las variables (por ejemplo, calor sensible, calor latente, GPP).
• Corrección de Gradientes: La diferenciación de modo inverso fue validada contra diferencias finitas centrales. Todos los parámetros activos lograron errores relativos menores a $1 \times 10^{-4}$. El pipeline manejó con éxito estructuras complejas, incluyendo resolvedores iterativos y subpasos de Runge-Kutta.
• Eficiencia de Calibración de Parámetros: En una prueba de concepto de calibración de tres parámetros, el optimizador basado en gradientes L-BFGS-B recuperó los valores reales con precisión de máquina ($\sim 10^{-19}$) en menos de 50 evaluaciones. Esto fue 8 veces más rápido (en términos de evaluaciones) que el optimizador libre de gradientes Nelder-Mead.
• Mejora de Rendimiento: Usando jax.vmap en una GPU (Quadro RTX 8000), el modelo logró una aceleración de 24x en tiempo de ejecución (wall-clock) sobre la ejecución secuencial de Fortran para un tamaño de ensamble de $N=2,048$. El costo amortizado por muestra cayó de 24.9 ms a 11.4 ms a medida que el tamaño del ensamble aumentó, mientras que Fortran se mantuvo plano en $\approx 54$ ms/muestra.
• Eficiencia Agéntica: A través de 73 tareas de traducción de módulos y 46 sesiones agénticas, el 81.2% de los errores fueron resueltos de forma autónoma en el primer intento. La supervisión humana total se estimó en solo 10–15 horas.

Significancia y Reivindicaciones
El artículo afirma demostrar que las bases de código científicas heredadas pueden traducirse sistemáticamente a marcos totalmente diferenciables sin una reimplementación manual, preservando la verdad científica fundamental.

• Beneficio Científico: Este trabajo permite la optimización de parámetros basada en gradientes y el análisis de sensibilidad que antes eran intratables para CLM-ml-v2. El análisis de Jacobiano revela que la temperatura del aire y la radiación de onda corta dominan las sensibilidades de la Producción Primaria Bruta (GPP) y los flujos de energía, mientras que $V_{c,max}$ ejerce un control distinto.
• Contribución Metodológica: El documento presenta un flujo de trabajo reproducible de cinco pasos con bucles de reparación autónomos y validación de oráculo de dominio específico. Cuantifica la viabilidad de la "coautoría" entre humanos y agentes de IA en la modernización de código científico, mostrando que la supervisión humana de alto nivel es suficiente para tareas de traducción complejas de varios días.
• Reutilización: Tanto el modelo traducido como la infraestructura del pipeline se liberan como un marco reutilizable, destinado a apoyar a la comunidad en la diferenciación de otros componentes de los modelos del sistema tierra.

Los autores posicionan este trabajo no como una solución final para todos los ESM, sino como una prueba de concepto validada que cierra la brecha entre la física heredada de Fortran y la programación diferenciable moderna, allanando el camino para arquitecturas híbridas de física-ML y una cuantificación de la incertidumbre más eficiente en la modelación climática.