The Ocean Model That Can Look Backward: How a Python Shadow Could Transform Climate Science
A Shadow Worth Having
In the basement of the Alfred Wegener Institute in Bremerhaven, Germany, a climate model is running. Not the production version—not the Fortran code that has been refined over decades and sits at the heart of Europe's next generation of climate simulations—but its shadow. A projection. A mirror written in a different language, on different hardware, with one crucial addition: it can look backward.
That might sound paradoxical for a model designed to forecast the future. But the ability to look backward—to trace, with mathematical precision, how every output depends on every input—is precisely what has been missing from state-of-the-art ocean modeling. It is the difference between knowing that your forecast is sensitive to cloud parameterizations and knowing exactly how sensitive, at every point in the ocean, down to every parameter in every equation. This is the gift of differentiation, and the shadow that now carries it is called FESOM2-JAX.
The numbers are striking on their own. A single GPU can now hold an entire coarse-resolution ocean, something that once required a cluster. A modest four-GPU node can simulate 113 years of Earth's ocean in a single wall-clock day. The gap between the new shadow and its Fortran original is two orders of magnitude smaller than the gap between either and reality. And for the first time, a global ocean-sea-ice model of CMIP-class complexity—one that countries will use to inform the next round of climate negotiations—runs natively in a differentiable framework, on an unstructured mesh, from a laptop to 256 GPUs without changing a line of code.
This is not a new model. It is a new way of working with the old one.
The Science
The ocean does not care about our computational convenience. It moves in eddies that range from tens to hundreds of kilometers, along boundary currents that pinch and release unpredictably, around coastlines that twist into fjords and archipelagos. A model that hopes to capture this behavior needs a mesh that can follow the geography—refining where the action is, coarsening where it is not. The Finite-volumE Sea ice-Ocean Model, known as FESOM2, does exactly this. Its triangular mesh can concentrate resolution in the Arctic, where sea ice demands it, or along the Gulf Stream, where eddies spin off into the deep Atlantic, without being forced into the regular grid that suits a textbook equation but not the real ocean.
FESOM2 has been under development for over a decade. It is the ocean component of the AWI climate model that Germany will submit to CMIP7, the seventh round of the Coupled Model Intercomparison Project that forms the scientific backbone of the IPCC reports. It has been assessed extensively against observations and against its predecessor. It scales to thousands of cores in its current Fortran-MPI implementation. It is, in other words, a serious tool with serious pedigree.
But Fortran is old. Not wrong—FESOM2's numerics are sound, its physics well-tested—but old in the way that matters: the generation of scientists now entering the field learned to code in Python. They are comfortable with PyTorch and TensorFlow, with the libraries that have transformed machine learning over the past decade. And they find Fortran's world opaque. Not because they cannot read it, but because the ecosystem around it—the interactivity, the rapid prototyping, the ability to inspect gradients and tweak parameters in real time—is missing.
This is the first pressure. The second is hardware. The computational power of modern supercomputers lives in graphics processing units, GPUs—specialized chips that excel at performing the same operation on millions of data points simultaneously. Fortran's toolchains for GPUs have improved, but they remain less mature than those for C++ or Python. The result is a gap: the science is in Fortran, but the hardware is speaking a different language.
The team at the Alfred Wegener Institute has been bridging this gap for several years. Their previous work produced a C++/Kokkos version of FESOM2, a performance-portable translation that allows the model to run on GPUs without rewriting the numerics. But Kokkos is still C++, still closer to the metal than to the Python world where so much modern scientific computing now happens.
FESOM2-JAX is a different kind of bridge. It is a re-implementation of FESOM2 in Python, using the JAX framework—a NumPy-style library developed by Google that has become the lingua franca of differentiable programming. The port was carried out kernel by kernel, each Fortran subroutine translated into a corresponding Python function, verified against the reference at every step. Large language models handled much of the mechanical translation; the researchers handled the verification.
The port covers the forced ocean-sea-ice configuration: the hydrostatic Boussinesq primitive equations on an unstructured triangular mesh, with a cell-vertex finite-volume discretization. Temperature and salinity live at mesh vertices; horizontal velocity lives at triangle centroids. The vertical coordinate is the arbitrary Lagrangian-Eulerian (ALE) scheme, used here in its rescaled form. Sea ice is a single-class dynamic-thermodynamic model with seven thickness classes. All of this matches the Fortran version exactly—what differs is the implementation.
Three choices shape the JAX implementation. First, one timestep is a pure function. The model state—every prognostic field, the Adams-Bashforth history, the warm-start vector of the conjugate gradient solver—is stored as a single immutable data structure, and one timestep takes that state and returns the next. This purity is what allows JAX to compile the entire timestep into a single optimized program, running it without returning to the Python interpreter until the integration is done. Second, differentiability is treated as a correctness property of every kernel, not a feature added later. Iterative solvers with data-dependent stopping are avoided in favor of fixed iteration counts; non-smooth operations are given finite gradients at their switching points; the memory of the reverse sweep is managed through checkpointing. Third, one source runs on any device count. The domain decomposition is FESOM2's own, but the exchange between devices is not a call into a message-passing library but an operation written in the same language as the rest of the model and compiled with it.
The verification procedure is rigorous. The port was built against an instrumented C translation of FESOM2 produced as part of the earlier Kokkos project. Each kernel was gated against this reference, with tolerances set by what the kernel does: point-wise operations agree to about 10⁻¹⁵ relative, accumulation operations to about 10⁻¹². Multi-step tests verify the state threading that single-step tests cannot see. The sharded code path reproduces the single-device path exactly on one device. Gradients are checked against finite differences kernel by kernel and through short full-model integrations. The suite comprises 78 test modules, roughly 640 tests in total.
Crucially, the chain closes on the original. The hindcast described in the next section compares the assembled JAX model against Fortran FESOM2 itself, not against the C reference the port was built from. An infidelity introduced at either translation step—Fortran to C, or C to JAX—would appear in this comparison. When a freshwater-budget defect was found over climate timescales, the kernel-by-kernel verification allowed the researchers to narrow the search to where two kernels were joined.
What They Found
The first test was a 62-year hindcast, from 1958 to 2019, on the CORE2 mesh with 1° equivalent resolution—126,858 surface vertices, 47 vertical layers, a timestep of 18,000 seconds. The configuration was identical to a Fortran run: same physics, same forcing, same initial conditions. The question was not whether the two models would give exactly the same answer—JAX does not reproduce the loop order of Fortran, so rounding differences accumulate—but whether the differences would be scientifically meaningful.
They were not.
The mean states of the JAX and Fortran versions differ from each other by two orders of magnitude less than either differs from observations. In the 1980-2009 climatology, the root-mean-square error between the JAX model and the PHC3.0 observational climatology is 0.87°C for sea surface temperature and 0.18°C for sea surface salinity. The error between Fortran and the same climatology is nearly identical: 0.85°C and 0.18°C. The error between JAX and Fortran is 0.01°C and 0.004°C—smaller than the observational uncertainty, and two orders of magnitude below the gap to reality.
This is visible in the bias maps. The JAX and Fortran versions show the same patterns of warm and cold patches, the same subtle shifts away from observations, the same eddy-rich regions where the 1° mesh struggles to resolve subgrid processes. The difference between them is an order of magnitude smaller than the difference between either and truth.
The agreement extends over six decades of integration. Global mean temperature, salinity, heat content, and sea ice area all track the Fortran run with barely perceptible divergence. The JAX model does not drift relative to its original; the shadow stays close to the body that casts it.
The performance numbers are equally striking. The complete 1° CORE2 configuration fits on a single GPU, with 6 million vertex-levels stored in memory and advanced each timestep. A node of four GH200 superchips integrates approximately 113 simulated years per wall-clock day. Meshes of up to 7.4 million surface vertices—roughly 5 km resolution over most of the globe—scale to 128 GPUs.
Simulation Throughput Across Mesh Sizes and Hardware
| Label | Value |
|---|---|
| CORE2 (4 A100) | 57.7 |
| CORE2 (8 A100) | 66.5 |
| DARS (16 A100) | 1.5 |
| DARS (16 GH200) | 3.6 |
| FORCA20 (16 A100) | 1.7 |
| NG5 (32 A100) | 0.9 |
| NG5 (64 A100) | 1.5 |
| NG5 (128 A100) | 2.8 |
The chart above shows throughput in simulated years per wall-clock day across the five meshes used in the study, comparing A100 GPUs (DKRZ Levante supercomputer) and GH200 GPUs (JUPITER booster). The CORE2 mesh saturates quickly, achieving its maximum throughput around four GPUs before communication overhead dominates. The larger meshes show sub-linear strong scaling: doubling the GPUs does not quite double the throughput, because the communication cost grows even as the per-device arithmetic shrinks. But the absolute numbers are impressive. At 128 GPUs on the NG5 mesh (7.4 million vertices, 511 million vertex-levels), the model advances 2.8 simulated years per wall-clock day. A full century-long integration would take 36 wall-clock days—a timeframe accessible to a researcher with GPU allocation on a European supercomputer.
GH200 is roughly twice as fast per GPU as A100 wherever the per-device partition is large enough to hide memory latency. This shifts each mesh's scaling knee to about half the device count. The CORE2 mesh, which saturates at four A100 GPUs, would likely saturate at two GH200 GPUs. The NG5 mesh, which peaks at 128 A100 GPUs at 2.8 SYPD, would reach the same throughput at roughly 64 GH200 GPUs.
Computational Cost per Simulated Year
| Label | Value |
|---|---|
| CORE2 (4 A100) | 2 |
| CORE2 (8 A100) | 3 |
| DARS (16 A100) | 263 |
| DARS (16 GH200) | 107 |
| FORCA20 (16 A100) | 227 |
| NG5 (32 A100) | 883 |
| NG5 (64 A100) | 1,024 |
| NG5 (128 A100) | 551 |
This chart tells a different story. It shows GPU-hours required to simulate one year of ocean behavior, plotted against the number of GPUs used. Lower is better. The CORE2 mesh is efficient at small device counts but becomes increasingly wasteful at large ones—the cost per simulated year rises because communication overhead grows even as the per-device work shrinks. The larger meshes are efficient across a wider range. The NG5 mesh at 128 A100 GPUs requires 551 GPU-hours per simulated year; at 64 GH200 GPUs, it requires 370 GPU-hours. In both cases, this is a quantity of compute that would have been unthinkable for a community ocean model a decade ago.
What limits the model is not arithmetic. It is communication. On the meshes tested here, a device has between 6 and 14 neighbors regardless of how many devices are used, because the spatial partition keeps nearby vertices together. This is the promise of unstructured meshes: bounded communication volume even as the problem size grows. But the implementation of that communication matters. The researchers implemented four halo exchange transports—ways of moving data between devices—and their performance differs significantly.
The all-gather broadcast moves the whole field, regardless of how little halo is actually needed. The ragged all-to-all moves exactly the halo entries but carries a slot for every device whether or not that device owns any halo data. The padded all-to-all adds padding so every slot is the same size, allowing vectorized communication but sending zeros that still cross the wire. The colored ppermute exchanges with one partner at a time in a sequence of rounds; its cost is bounded by the number of neighbors rather than by the device count, which keeps it flat as the partition spreads, but it pays startup costs for each round.
On CORE2 (small mesh, few GPUs), the padded and ragged all-to-alls are fastest. On NG5 at 64 GPUs, the colored ppermute wins at 0.236 seconds per step, versus 0.317 seconds for the padded all-to-all. The ragged all-to-all has a defective reverse-mode rule in JAX 10.1 and cannot be used in the backward pass—important, because what the shadow adds is the gradient.
Why This Changes Things
The Fortran production model remains the single source of truth. FESOM2-JAX is not a fork or a successor; it is a projection, kept faithful to the original so that developments made in the shadow can be transferred back. A parameter set tuned in JAX, a numerical variant validated in Python, a neural network trained through the dynamics—these can be carried into the Fortran codebase because the physics is identical. The shadow illuminates without replacing.
What the shadow adds is the gradient.
In the original Fortran model, obtaining an adjoint—the reverse-mode derivative that tells you how every output depends on every input—has been notoriously difficult. Source-to-source transformation tools require coding restrictions and hand-tuned checkpointing. Hand-coded adjoints must be regenerated and reverified whenever the forward model changes. Most ocean models, including FESOM2, never acquired one. Parameter estimation has relied on gradient-free methods: evolutionary algorithms, ensembleKalman filters, approximate Bayesian computation. These work, but they are slow, noisy, and limited in the questions they can answer.
In a differentiable framework, the adjoint is derived automatically from the same source that defines the forward model. It is exact for the discretization, remains consistent with the code as the code changes, and extends to whatever is added—including neural network components. One reverse-mode pass through the full time loop returns the sensitivity of a model diagnostic to a parameter at every mesh vertex.
The implications are practical. Gradient-based optimization can tune parameterizations more efficiently than gradient-free methods, with clear convergence criteria and the ability to use second-order information. Adjoint sensitivity studies can identify which parameters matter most at which locations—critical for understanding where model biases originate and how they might be reduced. Variational data assimilation can use the adjoint to constrain initial conditions or forcing fields to observations, the approach behind the ECCO ocean state estimates that have been so influential in quantifying ocean heat uptake.
But the most exciting applications involve machine learning.
Hybrid physics-machine-learning models have been a frontier in weather and climate modeling for several years. The idea is to retain the known physics—the dynamical core that solves the primitive equations—while replacing uncertain parameterizations with neural networks trained on high-resolution simulations or observations. The problem has always been training: how do you update the neural network weights to minimize error against reality when the loss function depends on the output of a physics model that itself depends on those weights?
The answer is backpropagation. You run the model forward, compute the loss against observations, then run the computation backward through the model to compute how each weight contributed to the error. This requires a differentiable model. In Fortran, it has required hand-coded adjoints or approximate methods. In FESOM2-JAX, it comes for free.
The researchers have not yet demonstrated online training of neural parameterizations—this is ongoing work. But they have verified the gradient. A single reverse-mode pass through the full time loop returns sensitivities that match finite-difference approximations to the same quantities. The gradient is correct. What remains is to use it.
This is not the first differentiable ocean model. Veros translated the pyOM2 primitive-equation model into Python with a JAX backend and runs global structured-grid configurations on GPUs. Veris recently added sea ice to the same ecosystem by re-implementing the MITgcm sea-ice model in JAX, matching hundreds of CPU cores on a single GPU. NeuralGCM couples a differentiable spectral dynamical core to a learned physics package and trains the hybrid model end to end against reanalysis.
What distinguishes FESOM2-JAX is complexity and geometry. It is the first global ocean-sea-ice model of CMIP-class complexity—the full parameterization suite of such a model, including mesoscale eddy parameterization (GM/Redi), three choices of vertical mixing, two sea-ice rheologies, and dynamic-thermodynamic sea ice—running in a single differentiable time loop. And it is the first built on an unstructured mesh, which opens the door to the regional refinement and variable resolution that are increasingly recognized as essential for next-generation climate modeling.
There is also a human dimension. The same properties that make JAX more programmable—interactivity, Python syntax, the ability to inspect and modify code on the fly—make it more teachable. A graduate student can read the model. They can run it on a laptop, see the output, change a parameter, see the output change. They can experiment in ways that are difficult with compiled Fortran production code. FESOM2-JAX is a vehicle for teaching as well as research.
The translation itself is a kind of proof of concept. The port was carried out with large language model assistance, using Claude Code as the agent and Opus 4.8 and Fable as the language models. The researchers describe this as an LLM-agentic workflow, and the process is documented in their earlier paper on the Fortran-to-Kokkos translation. What this suggests is that such translations are becoming systematic—that the bottleneck is no longer the mechanical work of porting but the verification of the result. The same tooling can likely be applied to other models, other codebases, other projections onto other frameworks.
What's Next
The paper is careful about what it claims. FESOM2-JAX is a model description, not an application paper. The authors write that the "applications enabled by this capability, including gradient-based calibration, hybrid physics-machine-learning modelling, and the transfer of their products back into the Fortran original, are the subject of ongoing work and lie beyond the scope of this model-description paper."
This is honest. The gradient exists, but no one has yet used it to calibrate a parameterization or train a neural network. The scaling results are impressive, but the largest mesh tested—7.4 million surface vertices, roughly 5 km resolution—is still far from the kilometer-scale resolution that some advocates argue is necessary for credible climate projections. The model does not yet include ice shelves, partial bottom cells, tides, icebergs, or ocean biogeochemistry—components that the Fortran version supports and that CMIP-class simulations require.
These are boundaries of a first paper, not limitations of the approach. The same can be said for the hardware: the measurements were made on A100 and GH200 GPUs; future architectures may behave differently. The four halo transports were implemented and tested; other approaches, including MPI-based communication for very large device counts, may be necessary as the model moves toward production use.
One technical limitation is worth noting: the ragged all-to-all, which ships exactly the halo entries each device needs and nothing more, has a defective reverse-mode rule in JAX 10.1. It cannot be used in the backward pass. The colored ppermute is an excellent alternative—it keeps communication bounded as the device count grows—but it pays round-by-round startup costs that a single all-to-all does not. This is a software limitation, not a mathematical one, and it may be fixed in future JAX releases.
There are scientific questions the paper does not address. How well does the gradient-based calibration perform compared to gradient-free methods on the same problem? Can neural parameterizations trained in FESOM2-JAX be transferred to the Fortran model and retain their skill? Does the unstructured mesh offer advantages for machine learning that structured grids do not—for example, the ability to refine resolution in data-rich regions without changing the global grid? These are the questions that the next papers will answer.
The shadow has been cast. What it illuminates is still being discovered.
The climate models that inform international policy are built on decades of work in Fortran. They are the product of thousands of person-years of development, tested and evaluated against every observation humanity has taken of the planetary interior. They are not going to be replaced overnight by Python re-implementations. But the boundary between the production model and the experimental model is softening. FESOM2-JAX brings the full complexity of a CMIP-class ocean model into the ecosystem where machine learning lives, where rapid iteration is possible, where the barrier to entry is lower.
It is, in the language of the paper, a code shadow. But shadows are not just reflections. In the right light, they illuminate what was previously in darkness.
For the generation of scientists who learned to code in Python, who think in terms of gradients and automatic differentiation, who want to train neural networks through the full complexity of a global ocean model—for them, FESOM2-JAX is not a shadow at all. It is the main event.
And for the first time, they can look backward.