the Creative Commons Attribution 4.0 License.
the Creative Commons Attribution 4.0 License.
JCM v1.0: A Differentiable, Intermediate-Complexity Atmospheric Model
Abstract. In this paper we present version 1.0 of the JAX Circulation Model (JCM). JCM is an open-source, differentiable atmospheric model built in Python using the JAX numerical library. Earth system modeling is rapidly evolving, particularly through hybrid approaches that combine known dynamics with data-driven components. However, the training and validation of hybrid methods in traditional models remain difficult due to the absence of gradients and the complexity of legacy code. Differentiable models written in modern frameworks offer a path forward. JCM couples physics parameterizations to the Dinosaur dynamical core through a flexible interface that makes substitution of other schemes easy. The default parameterization scheme uses the SPEEDY (Simplified Parameterizations, primitivE-Equation DYnamics) intermediate-complexity physics scheme. This modularity supports benchmarking across physical and machine-learned schemes, with direct access to gradients for sensitivity analysis, calibration, and online learning. We show validation of JCM against the original Fortran SPEEDY code at T31 resolution. We also highlight JCM's differentiability and efficiency and outline plans for extending the framework to a differentiable Earth system model. JCM provides a lightweight yet expressive platform for accelerating research in climate modeling.
- Preprint
(3623 KB) - Metadata XML
- BibTeX
- EndNote
Status: open (until 23 Mar 2026)
-
RC1: 'Comment on egusphere-2025-6266', Anonymous Referee #1, 06 Feb 2026
reply
The comment was uploaded in the form of a supplement: https://egusphere.copernicus.org/preprints/2026/egusphere-2025-6266/egusphere-2025-6266-RC1-supplement.pdfReplyCitation: https://doi.org/
10.5194/egusphere-2025-6266-RC1 -
RC2: 'Comment on egusphere-2025-6266', Anonymous Referee #2, 19 Feb 2026
reply
This paper describes the JAX Circulation Model (JCM), a differentiable atmospheric model written in Python using JAX. The paper is well written. I downloaded the code, and it was easy to follow the documentation and run the code on a laptop. I would be happy to accept the paper subject to some minor suggestions.
- Line 152: Could you explain a bit more what the orographic corrections are?
- Line 186: Is the RMS difference using the monthly data, or the average of the three-year simulation?
- Lines 189 - 198 and Figures 2 and 3: The differences in zonal wind between JCM and speedy look pretty large to me. Could you explain more about how the differences in horizontal diffusion and orographic corrections result in these changes?
- Figures 5 and 6: I think it would be clearer to present performance metrics as SYPD (simulated years per day). It would also make it easier to compare with other models.
- Line 256: I don’t think this is accurate. For example, ensemble methods that are gradient-free are commonly used in data assimilation in weather forecasting. They can be used for tuning parameters, too.
- Section 6.2: I was hoping the authors could show a calibration of more parameters, especially as it is mentioned that “Calibration with gradient information allows for the assessment of model behaviour across more of the parameter space”. I understand if it is beyond the scope of this paper, though.
- Line 269: Could you describe what is in y (the observation)?
Citation: https://doi.org/10.5194/egusphere-2025-6266-RC2 -
CC1: 'Comment on egusphere-2025-6266', Juan Antonio Añel, 11 Mar 2026
reply
I would like to note that the term "open-source" used by the authors here is commercial terminology, which does not correspond to the academic nature of the scientific work. The more correct terminology to refer in the text to the model would be a free software or a FLOSS model. This argument is backed by the fact that the authors release it under the Apache v2.0 license, which is widely defined as a free software license by the Apache Foundation.
Citation: https://doi.org/10.5194/egusphere-2025-6266-CC1 -
CC2: 'Comment on egusphere-2025-6266', Maximilian Gelbrecht, 20 Mar 2026
reply
Comment on JCM v1.0: A Differentiable, Intermediate-Complexity Atmospheric Model
Maximillian Gelbrecht [1,2], Brian Groenke [1,2], and Gregory Wagner [3,4]
1 Potsdam Institute for Climate Impact Research, Member of the Leibniz Association, Potsdam, Germany,
2 Munich Climate Center and Earth System Modeling Group, Technical University of Munich, Munich, Germany
3 Aeolus Labs, San Francisco, CA, USA
4 Massachusetts Institute for Technology, Cambridge, MA, USAWe are pleased to see the timely publication of the authors’ work on developing a new General Circulation Model in python using the JAX automatic differentiation (AD) library. We emphatically agree with their view that differentiable models provide a promising way to bridge the gap between physics-based and data-driven modeling techniques and would constitute a major step forward for Earth system modeling. The model presented here, the JAX Circulation Model (JCM), is an intermediate-complexity atmospheric general circulation model based on the existing dynamical core of NeuralGCM (Kochkov et al., 2024), called Dinosaur, which solves the primitive equations based on a spectral discretization. Dinosaur is itself derived from the SPEEDY GCM (Molteni, 2003) which generally serves as a convenient basis for hybrid atmospheric modeling thanks to its reduced complexity and computational costs.
However, it is impossible to ignore that the authors of JCM seem to make no mention of other parallel efforts to develop modern, differentiable, and GPU-accelerated variants of SPEEDY, most notably the Julia package, SpeedyWeather.jl (Klöwer et al., 2024) which has recently been made differentiable using the Enzyme.jl automatic differentiation framework (Moses et al., 2025). Other similar efforts that might be mentioned include Oceananigans.jl (Wagner et al., 2025) for ocean modeling and the HOPE shallow water model in PyTorch (Zhou, Xue and Shen, 2025).
Another possible issue is visible in Figure 4, JCM seems to exhibit what could be spectral ringing artefacts in some of its parametrizations presented. Those are not present in the Fortran Speedy figures just below. It is hard to tell just from looking at these coarse figures, but it would be a good idea to investigate this further as they could point to problems with the parametrizations. Furthermore, since JCM and NeuralGCM share an identical dynamical core, it would be really interesting to see a comparison between the learned parametrizations of NeuralGCM and the SPEEDY parametrizations of JCM, either in the present or future work.
A third point where we feel the authors could be clearer is in their discussion regarding the benefits of JAX as the underlying numerical framework. The authors note that “Python is the most commonly used programming language in the world, and its libraries are especially powerful for scientific analysis (e.g., Xarray and Dask) and ML (e.g., PyTorch and TensorFlow)...”. However, we find this statement to be potentially misleading due to the fact that none of the aforementioned python libraries are directly compatible with the JAX domain-specific language. Due to its reliance on JAX, JCM would not be able to directly integrate existing python code based on Xarray/Dask, PyTorch, or TensorFlow; rather, this code would need to be partially or wholly ported to JAX, which may range of mildly inconvenient to enormously difficult depending on the complexity of the software. These tools could of course still be used for preprocessing and postprocessing of input and output data; however, this is already possible with models written in other programming languages. The authors should therefore clarify these limitations and briefly outline a potential strategy for the integration of JCM into the broader Python ecosystem.
We are nevertheless excited to see how JCM further develops and we congratulate the authors on their pioneering work. It is clear that the integration of data-driven methods into the model will be particularly promising given its basis on JAX. JCM undoubtedly represents an important milestone for differentiable Earth system modelling and a great basis for a lot of exciting future work.
References
Klöwer, M. et al. (2024) “SpeedyWeather.jl: Reinventing atmospheric general circulation models towards interactivity and extensibility,” Journal of Open Source Software, 9(98), p. 6323. Available at: https://doi.org/10.21105/joss.06323.
Kochkov, D. et al. (2024) “Neural general circulation models for weather and climate,” Nature, 632(8027), pp. 1060–1066. Available at: https://doi.org/10.1038/s41586-024-07744-y.
Molteni, F. (2003) “Atmospheric simulations using a GCM with simplified physical parametrizations. I: model climatology and variability in multi-decadal experiments,” Climate Dynamics, 20(2), pp. 175–191. Available at: https://doi.org/10.1007/s00382-002-0268-2.
Moses, W.S. et al. (2025) “DJ4Earth: Differentiable, and Performance-portable Earth System Modeling via Program Transformations.” ESS Open Archive. Available at: https://doi.org/10.22541/essoar.176314951.18114616/v1.
Wagner, G.L. et al. (2025) “High-level, high-resolution ocean modeling at all scales with Oceananigans.” arXiv. Available at: https://doi.org/10.48550/arXiv.2502.14148.
Zhou, L., Xue, W. and Shen, X. (2025) “HOPE: an arbitrary-order non-oscillatory finite-volume shallow water dynamical core with automatic differentiation,” Geoscientific Model Development, 18(21), pp. 8175–8201. Available at: https://doi.org/10.5194/gmd-18-8175-2025.
Citation: https://doi.org/10.5194/egusphere-2025-6266-CC2
Viewed
| HTML | XML | Total | BibTeX | EndNote | |
|---|---|---|---|---|---|
| 677 | 516 | 22 | 1,215 | 24 | 29 |
- HTML: 677
- PDF: 516
- XML: 22
- Total: 1,215
- BibTeX: 24
- EndNote: 29
Viewed (geographical distribution)
| Country | # | Views | % |
|---|
| Total: | 0 |
| HTML: | 0 |
| PDF: | 0 |
| XML: | 0 |
- 1