legoESM: a modular and differentiable Earth system model, rewritten in JAX by AI programming agents

Its components fit together like building blocks, from metre-scale simulations to global climate projections, and can be calibrated using a gradient on a GPU.

Researchers present legoESM, a model of the Earth system designed for accelerator hardware and machine learning. It takes decades of parameterizations and numerical methods developed by the community, but brings them together in a unified framework written in JAX. Particularity: this rewriting work was entrusted to AI programming agents, supervised by a "scientific contract" defined by humans and verified by reference tests.

Interchangeable bricks

Dynamic cores, physics schemes, grids, complexity levels and components can be swapped at will, and each brick can be based on classical physics or a learned emulator. The same code covers both meter-scale simulation of large vortices and global simulations from weather to climate. Because the model is end-to-end differentiable, it lends itself to gradient calibration, variational data assimilation and online training.

The authors indicate that legoESM reproduces realistic simulations at several scales, reduces the temperature bias at the land surface by calibration and efficiently switches to kilometer resolution on GPU. They propose it as an open infrastructure to test hypotheses, especially on the radiative feedbacks of clouds, the main source of uncertainty in projections.