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

Sintesi Tecnica: Traduzione Sistematica di Codice Scientifico Legacy tramite LLM in Framework Differenziabili

Definizione del Problema
I modelli legacy dei sistemi terrestri (ESM), in particolare i modelli della superficie terrestre (LSM), sono critici per le politiche climatiche e le previsioni ma rimangono ampiamente inaccessibili alle moderne tecniche di programmazione differenziabile. Questi modelli, che spesso comprendono centinaia di migliavere di righe di codice Fortran, soffrono di scarsa modularità, componenti deprecate e una forza lavoro di sviluppatori in contrazione. Sebbene la programmazione differenziabile permetta una stima dei parametri e un'analisi di sensibilità basata su gradienti efficienti, nonché l'assimilazione dati, la migrazione di codebase legacy verso framework differenziabili (come JAX) rimane una sfida irrisolta. I tentativi attuali di modernizzazione del codice consistono o in riscritture complete che scartano la fisica legacy o falliscono nel scalare oltre funzioni isolate a causa delle dipendenze tra file e della mancanza di fedeltà numerica nelle traduzioni effettuate da Large Language Models (LLM). Inoltre, i benchmark attuali degli LLM non misurano adeguatamente la validità scientifica o la parità numerica richiesta dai modelli fisici.

Metodologia: Una Pipeline Agente a Cinque Fasi
Gli autori propongono una pipeline agente sistematica a cinque fasi per tradurre il modello della superficie terrestre CLM-ml-v2, composto da 19.000 righe e 102 moduli, in un'implementazione JAX completamente differenziabile denominata clm-ml-jax. La pipeline si basa su tre principi di progettazione fondamentali: analisi delle dipendenze statiche, test di parità numerica modulare e verifica completa del gradiente.

1. Fase 1: Scoping e Analisi delle Dipendenze: Per prevenire le allucinazioni degli LLM riguardo ai tipi di dati nei codebase accoppiati, gli autori eseguono un'analisi statica sull'intero grafo delle dipendenze dei moduli. Un ordinamento topologico determina un ordine di traduzione che rispetta le dipendenze, garantendo che tutti i moduli importati siano tradotti prima dei moduli che dipendono da essi.
2. Fase 2: Configurazione della Documentazione dello Stato: La pipeline utilizza documenti di stato persistenti (CLAUDE.md, plan.md, CHANGELOG.md) per fungere da memoria esterna dell'agente. Questi documenti codificano le convenzioni del progetto, tracciano i flag di completamento e registrano i tentativi falliti, consentendo una ripresa coerente delle sessioni attraverso i confini della finestra di contesto ed evitando che gli agenti tornino su strade senza uscita.
3. Fase 3: Costruzione dell'Oracolo Fortran: Prima della traduzione, un team di agenti costruisce una suite di test funzionali. Le subroutine sono classificate in tre livelli in base alla testabilità. Il sistema esegue il codice Fortran originale sotto input fissi per generare coppie input-output "golden", stabilendo un oracolo di riferimento per la validazione numerica.
4. Fase 4: Traduzione, Test e Riparazione Iterativa: I moduli vengono elaborati nell'ordine topologico stabilito nella Fase 1. Un agente traduttore genera codice JAX, che viene poi validato da un agente di testing rispetto all'oracolo Fortran (richiedendo una tolleranza relativa entro $10^{-4}$). Se un test fallisce, un "Ralph Loop" (un pattern di orchestrazione) attiva un agente di riparazione per diagnosticare le discrepanze, modificare il codice e testarlo nuovamente in modo autonomo finché il modulo non passa o non viene segnalato per la revisione umana.
5. Fase 5: Integrazione e Validazione della Differenziabilità: Una volta tradotti i moduli, un Agente di Integrazione assembla l'intera pipeline della colonna e valida il comportamento end-to-end. Un Agente di Differenziabilità risolve iterativamente problemi specifici di JAX, come:
   • Svolgimento Statico dei Cicli (Static Loop Unrolling): Sostituzione dei cicli for a livello Python con jax.lax.scan per tracciare i corpi dei cicli una sola volta al momento della compilazione, evitando grafi di computazione ingestibili per i solver Runge-Kutta.
   • Tipizzazione Statica vs. Dinamica: Conversione delle costanti in scalari JAX per evitare ricompilazioni complete causate dal cambiamento delle chiavi di cache.
   • Guardie per la Sicurezza dei Branch (Branch-Safe Guards): Applicazione di limiti inferiori (es. jnp.maximum(x, 1e-30)) per prevenire gradienti NaN da operazioni come $0 \times \infty$ in espressioni jnp.where.
   • Adjoint del Teorema della Funzione Implicita (IFT) per i Solvers: Differenziazione di solver iterativi (es. bisezione della conduttanza stomatica, lunghezza di Monin–Obukhov) utilizzando gli adjoint di IFT invece di svolgere il ciclo (unrolling), il che produrrebbe gradienti errati.

Risultati Chiave
Il modello clm-ml-jax risultante è stato valutato sul sito AmeriFlux CHATS7 walnut orchard utilizzando dati meteorologici di maggio 2007.

• Equivalenza Numerica: Il modello tradotto ha raggiunto una parità numerica quasi perfetta con il riferimento Fortran attraverso 1488 timestep, con coefficienti di correlazione di 1.000 e un bias trascurabile per tutte le variabili (es. calore sensibile, calore latente, GPP).
• Correttezza del Gradiente: La differenziazione in modalità inversa è stata validata rispetto alle differenze centrali finite. Tutti i parametri attivi hanno raggiunto errori relativi inferiori a $1 \times 10^{-4}$. La pipeline ha gestito con successo strutture complesse, inclusi solver iterativi e sottosteppi Runge-Kutta.
• Efficienza della Calibrazione dei Parametri: In una prova di concetto di calibrazione di tre parametri, l'ottimizzatore basato su gradiente L-BFGS-B ha recuperato i valori reali con precisione di macchina ($\sim 10^{-19}$) in meno di 50 valutazioni. Questo è stato 8 volte più veloce (in termini di valutazioni) rispetto all'ottimizzatore gradient-free Nelder-Mead.
• Accelerazione delle Prestazioni: Utilizzando jax.vmap su una GPU (Quadro RTX 8000), il modello ha ottenuto un'accelerazione del tempo di esecuzione (wall-clock) di 24x rispetto all'esecuzione sequenziale in Fortran con una dimensione dell'ensemble di $N=2.048$. Il costo ammortizzato per campione è sceso da 24,9 ms a 11,4 ms all'aumentare della dimensione dell'ensemble, mentre il Fortran è rimasto piatto a $\approx 54$ ms/campione.
• Efficienza Agente: Attraverso 73 task di traduzione di moduli e 46 sessioni agenti, l'81,2% dei bug è stato risolto autonomamente al primo tentativo. La supervisione umana totale è stata stimata in soli 10–15 ore.

Significato e Rivendicazioni
Il documento sostiene di aver dimostrato che le codebase scientifiche legacy possono essere sistematicamente tradotte in framework completamente differenziabili senza reimplementazione manuale, preservando la verità scientifica di base.

• Risultato Scientifico: Il lavoro abilita l'ottimizzazione dei parametri basata su gradienti e l'analisi di sensibilità che erano precedentemente intrattabili per CLM-ml-v2. L'analisi del Jacobiano rivela che la temperatura dell'aria e la radiazione a onde corte dominano le sensibilità della produzione primaria lorda (GPP) e dei flussi energetici, mentre $V_{c,max}$ esercita un controllo distinto.
• Contributo Metodologico: Il documento presenta un flusso di lavoro riproducibile in cinque fasi con loop di riparazione autonomi e validazione dell'oracolo specifica del dominio. Quantifica la fattibilità della "co-autoria" tra umani e agenti IA nella modernizzazione del codice scientifico, mostrando che una supervisione umana di alto livello è sufficiente per compiti di traduzione complessi che durano più giorni.
• Riutilizzabilità: Il modello tradotto e l'infrastruttura della pipeline sono rilasciati come framework riutilizzabile, destinato a supportare la comunità nella differenziazione di altri componenti dei modelli del sistema terrestre.

Gli autori posizionano questo lavoro non come una soluzione finale per tutti gli ESM, ma come una prova di concetto validata che colma il divario tra la fisica legacy in Fortran e la moderna programmazione differenziabile, aprendo la strada ad architetture ibride fisica-ML e a una quantificazione dell'incertezza più efficiente nella modellazione climatica.