Development of the Differentiable Ice Sheet Model JAX-MALI
Abstract
Ice sheet models that accurately resolve physics are critical for predicting Sea Level Rise (SLR) and its sensitivity to changing climate conditions. Here we present JAX-MALI, a differentiable ice sheet model built around the established MPAS-Albany Land Ice (MALI) model. Ice sheet physics and dynamic changes are simulated within a JAX differentiable framework, while the momentum balance is solved by a dynamical core: either the accurate and well-tested MALI/Albany package, or an equivalent solver of our own. Because the time loop is differentiable in JAX, the model provides a transient adjoint with no bespoke code to write and maintain, together with Jacobian and Hessian products supporting hybrid physics-ML emulators, differentiable physics-based residuals and losses, data assimilation, uncertainty quantification and optimization. The model was developed with coding assistants using a test harness translated from MALI's own physics test suite, gating each module on MALI's per-cell fields to ensure accurate physics reproduction. On a reference simulation of the full Antarctic ice sheet we reproduce MALI's 15-year projection to 2e-6 mm sea-level equivalent, at 1.6x MALI's speed. The adjoint allows for data assimilation and end-to-end sensitivity of projected SLR to input data, parameters and forcings. Conversely, the model can inform new observational campaigns, closing the loop between observation and modeling.