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

Resumo Técnico: Tradução Sistemática de Código Científico Legado para Frameworks Diferenciáveis via LLM

Definição do Problema
Modelos de Sistemas Terrestres (ESMs) legados, particularmente modelos de superfície terrestre (LSMs), são críticos para políticas climáticas e previsões, mas permanecem amplamente inacessíveis às modernas técnicas de programação diferenciável. Esses modelos, que frequentemente compreendem centenas de milhares de linhas de código Fortran, sofrem com baixa modularidade, componentes obsoletos e uma força de trabalho de desenvolvedores em declínio. Embora a programação diferenciável permita a estimativa de parâmetros, análise de sensibilidade e assimilação de dados baseadas em gradientes de forma eficiente, migrar bases de código legadas para frameworks diferenciáveis (como o JAX) permanece um desafio não resolvido. As tentativas atuais de modernização de código envolvem ou reescritas completas que descartam a física legada ou falham em escalar além de funções isoladas devido a dependências entre arquivos e à falta de fidelidade numérica nas traduções feitas por Large Language Models (LLMs). Além disso, os benchmarks atuais de LLM não medem adequadamente a validade científica ou a paridade numérica exigida para modelos físicos.

Metodologia: Um Pipeline Agêntico de Cinco Fases
Os autores propõem um pipeline agêntico sistemático de cinco fases para traduzir o modelo de superfície terrestre CLM-ml-v2, composto por 19.000 linhas e 102 módulos em Fortran, para uma implementação JAX totalmente diferenciável chamada clm-ml-jax. O pipeline baseia-se em três princípios de design fundamentais: análise de dependência estática, testes de paridade numérica modulares e verificação abrangente de gradiente.

1. Fase 1: Escopo e Análise de Dependência: Para evitar alucinações de LLM em relação aos tipos de dados em bases de código acopladas, os autores realizam análise estática no grafo de dependência de módulos completo. Uma ordenação topológica determina uma ordem de tradução que respeita as dependências, garantindo que todos os módulos importados sejam traduzidos antes dos módulos que deles dependem.
2. Fase 2: Configuração da Documentação de Estado: O pipeline utiliza documentos de estado persistentes (CLAUDE.md, plan.md, CHANGELOG.md) para servir como memória externa do agente. Esses documentos codificam convenções do projeto, rastreiam flags de conclusão e registram tentativas falhas, permitindo a retomada coerente de sessões através de fronteiras de janela de contexto e evitando que os agentes retornem a caminhos sem saída.
3. Fase 3: Construção do Oráculo Fortran: Antes da tradução, uma equipe de agentes constrói um conjunto de testes funcionais. As sub-rotinas são classificadas em três níveis baseados na testabilidade. O sistema executa o código Fortran original sob entradas fixas para gerar pares de entrada-saída "golden", estabelecendo um oráculo de referência para validação numérica.
4. Fase 4: Tradução, Teste e Reparo Iterativo: Os módulos são processados na ordem topológica estabelecida na Fase 1. Um agente tradutor gera o código JAX, que é então validado por um agente de teste contra o oráculo Fortran (exigindo tolerância relativa dentro de $10^{-4}$). Se um teste falhar, um "Ralph Loop" (um padrão de orquestração) aciona um agente de reparo para diagnosticar discrepâncias, editar o código e retestar autonomamente até que o módulo passe ou seja sinalizado para revisão humana.
5. Fase 5: Integração e Validação de Diferenciabilidade: Uma vez traduzidos os módulos, um Agente de Integração monta o pipeline de coluna completa e valida o comportamento end-to-end. Um Agente de Diferenciabilidade então resolve iterativamente questões específicas do JAX, tais como:
   • Desenrolamento de Loop Estático (Static Loop Unrolling): Substituição de loops for de nível Python por jax.lax.scan para rastrear corpos de loop uma única vez no tempo de compilação, evitando grafos de computação incontroláveis para solvers de Runge-Kutta.
   • Tipagem Estática vs. Dinâmica: Conversão de constantes para escalares JAX para evitar a recompilação total causada por mudanças nas chaves de cache.
   • Salvaguardas de Segurança de Ramificação (Branch-Safe Guards): Aplicação de limites inferiores (ex: jnp.maximum(x, 1e-30)) para evitar gradientes NaN em operações como $0 \times \infty$ em expressões jnp.where.
   • Adjuntos do Teorema da Função Implícita (IFT Adjoins): Diferenciação de solvers iterativos (ex: bisseção da condutância estomática, comprimento de Monin–Obukhov) usando adjuntos IFT em vez de desenrolar o loop, o que resultaria em gradientes incorretos.

Principais Resultados
O modelo resultante, clm-ml-jax, foi avaliado no site de plantação de nogueiras CHATS7 AmeriFlux utilizando dados meteorológicos de maio de 2007.

• Equivalência Numérica: O modelo traduzido alcançou paridade numérica quase perfeita com a referência Fortran ao longo de 1488 passos de tempo, com coeficientes de correlação de 1.000 e viés negligenciável para todas as variáveis (ex: calor sensível, calor latente, GPP).
• Correção de Gradiente: A diferenciação de modo reverso foi validada contra diferenças finitas centrais. Todos os parâmetros ativos alcançaram erros relativos menores que $1 \times 10^{-4}$. O pipeline lidou com sucesso com estruturas complexas, incluindo solvers iterativos e sub-passos de Runge-Kutta.
• Eficiência de Calibração de Parâmetros: Em uma prova de conceito de calibração de três parâmetros, o otimizador baseado em gradiente L-BFGS-B recuperou os valores reais com precisão de máquina ($\sim 10^{-19}$) em menos de 50 avaliações. Isso foi 8 vezes mais rápido (em termos de avaliações) do que o otimizador livre de gradiente Nelder-Mead.
• Aceleração de Performance: Usando jax.vmap em uma GPU (Quadro RTX 8000), o modelo alcançou um ganho de velocidade de 24x no tempo de execução (wall-clock) sobre a execução sequencial em Fortran para um tamanho de ensemble de $N=2.048$. O custo amortizado por amostra caiu de 24,9 ms para 11,4 ms conforme o tamanho do ensemble aumentava, enquanto o Fortran permaneceu constante em $\approx 54$ ms/amostra.
• Eficiência Agêntica: Através de 73 tarefas de tradução de módulos e 46 sessões agênticas, 81,2% dos bugs foram resolvidos autonomamente na primeira tentativa. A supervisão humana total foi estimada em apenas 10–15 horas.

Significância e Alegações
O artigo afirma demonstrar que bases de código científicas legadas podem ser sistematicamente traduzidas para frameworks totalmente diferenciáveis sem reimplantação manual, preservando a verdade científica fundamental.

• Retorno Científico: O trabalho permite a otimização de parâmetros baseada em gradiente e análise de sensibilidade que eram intratáveis para o CLM-ml-v2. A análise Jacobiana revela que a temperatura do ar e a radiação de ondas curtas dominam as sensibilidades da Produção Primária Bruta (GPP) e dos fluxos de energia, enquanto $V_{c,max}$ exerce um controle distinto.
• Contribuição Metodológica: O artigo apresenta um fluxo de trabalho reprodutível de cinco etapas com loops de reparo autônomos e validação de oráculo de domínio específico. Ele quantifica a viabilidade da "coautoria" entre humanos e agentes de IA na modernização de código científico, mostrando que a supervisão humana de alto nível é suficiente para tarefas de tradução complexas de vários dias.
• Reutilização: Tanto o modelo traduzido quanto a infraestrutura do pipeline são lançados como um framework reutilizável, destinado a apoiar a comunidade na diferenciação de outros componentes de modelos de sistemas terrestres.

Os autores posicionam este trabalho não como uma solução final para todos os ESMs, mas como uma prova de conceito validada que preenche a lacuna entre a física legada em Fortran e a moderna programação diferenciável, pavimentando o caminho para arquiteturas híbridas de física-ML e uma incerteza quantificada mais eficiente na modelagem climática.