JCM v1.0: A Differentiable, Intermediate-Complexity Atmospheric Model

Davenport, Ellen H.; Madan, J. Varan; Gjini, Rebecca; Brzenski, Jared; Ho, Nick; Hsu, Tien-Yiao; Liang, Yueshan; Liu, Zhixing; Manivannan, Veeramakali; Pham, Eric; Vutukuru, Rohith; Williams, Andrew I. L.; Yang, Zhiqi; Yu, Rose; Lutsko, Nicholas J.; Hoyer, Stephan; Watson-Parris, Duncan

doi:10.5194/egusphere-2025-6266

Preprints

https://doi.org/10.5194/egusphere-2025-6266

Preprints

26 Jan 2026

| 26 Jan 2026

JCM v1.0: A Differentiable, Intermediate-Complexity Atmospheric Model

Ellen H. Davenport, J. Varan Madan, Rebecca Gjini, Jared Brzenski, Nick Ho, Tien-Yiao Hsu, Yueshan Liang, Zhixing Liu, Veeramakali Manivannan, Eric Pham, Rohith Vutukuru, Andrew I. L. Williams, Zhiqi Yang, Rose Yu, Nicholas J. Lutsko, Stephan Hoyer, and Duncan Watson-Parris

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.

Received: 16 Dec 2025 – Discussion started: 26 Jan 2026

Publisher's note: Copernicus Publications remains neutral with regard to jurisdictional claims made in the text, published maps, institutional affiliations, or any other geographical representation in this paper. While Copernicus Publications makes every effort to include appropriate place names, the final responsibility lies with the authors. Views expressed in the text are those of the authors and do not necessarily reflect the views of the publisher.

Download & links

Status: final response (author comments only)

RC1:
'Comment on egusphere-2025-6266', Anonymous Referee #1, 06 Feb 2026

The comment was uploaded in the form of a supplement: https://egusphere.copernicus.org/preprints/2026/egusphere-2025-6266/egusphere-2025-6266-RC1-supplement.pdf

Citation: https://doi.org/10.5194/egusphere-2025-6266-RC1
- AC2: 'Reply on RC1', Duncan Watson-Parris, 03 Apr 2026
  
  The comment was uploaded in the form of a supplement: https://egusphere.copernicus.org/preprints/2026/egusphere-2025-6266/egusphere-2025-6266-AC2-supplement.pdf
  
  Citation: https://doi.org/10.5194/egusphere-2025-6266-AC2
RC2:
'Comment on egusphere-2025-6266', Anonymous Referee #2, 19 Feb 2026
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
- AC1: 'Reply on RC2', Duncan Watson-Parris, 03 Apr 2026
  
  The comment was uploaded in the form of a supplement: https://egusphere.copernicus.org/preprints/2026/egusphere-2025-6266/egusphere-2025-6266-AC1-supplement.pdf
  
  Citation: https://doi.org/10.5194/egusphere-2025-6266-AC1
CC1:
'Comment on egusphere-2025-6266', Juan Antonio Añel, 11 Mar 2026

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
- AC5: 'Reply on CC1', Duncan Watson-Parris, 03 Apr 2026
  
  The comment was uploaded in the form of a supplement: https://egusphere.copernicus.org/preprints/2026/egusphere-2025-6266/egusphere-2025-6266-AC5-supplement.pdf
  
  Citation: https://doi.org/10.5194/egusphere-2025-6266-AC5
CC2:
'Comment on egusphere-2025-6266', Maximilian Gelbrecht, 20 Mar 2026

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, USA
We 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
- AC4: 'Reply on CC2', Duncan Watson-Parris, 03 Apr 2026
  
  The comment was uploaded in the form of a supplement: https://egusphere.copernicus.org/preprints/2026/egusphere-2025-6266/egusphere-2025-6266-AC4-supplement.pdf
  
  Citation: https://doi.org/10.5194/egusphere-2025-6266-AC4
CC3:
'Comment on egusphere-2025-6266', Milan Klöwer, 23 Mar 2026

The authors present JCM a JAX-based general circulation model whereby the authors contributions is to have translated existing parameterizations from Fortran SPEEDY into JAX so they can be used together with the existing dynamical core Dinosaur (which was developed for NeuralGCM). I very much agree with the overall narrative of the paper: Modern GCMs written in easily accessible and productive languages (such as Python) and for modern hardware (GPU, TPU) with modern features such as differentiability can lower the barrier for climate modelling and accelerate climate research. In that sense, the paper is timely, important and bridges existing model (Dinosaur and Fortran SPEEDY's parameterizations) for a useful and user-friendly model.
I do recommend this paper for publication after following points have been addressed
Major points:
- Similar efforts by other authors: The introduction currently misses many modelling efforts pursued by other teams that achieve similar goals: On the Python side there is the ocean model Veros (only mentioned in the conclusions) which is also JAX-based and which has been tested on multiple GPUs too, and I believe they also have been working with JAX's AD capabilities too. Please cite them accordingly, otherwise the introduction reads as if JCM is the first JAX-based Earth-system model component that as GPU/TPU and differentiability in mind. On the Julia side there are also many efforts, primarily Oceananigans (e.g. https://arxiv.org/abs/2502.14148) and SpeedyWeather (https://doi.org/10.21105/joss.06323) with the latter being the closest to your efforts. AD for both are documented in https://doi.org/10.22541/essoar.176314951.18114616/v1. Otherwise the reader would get an incomplete summary on ongoing modelling efforts and the literature published on them.
- Spectral ringing (Fig 2a, b, Fig. 4a-c). Evident from the figures is a strong spectral ringing present in JCM simulations. They do seem to be constant in time such that I would conclude they do not necessarily present a dynamic (e.g. CFL) stability issue in, e.g., strong winds. But they may originate as a consequence of orography. While you could use a super smooth orography this also may show deficiencies in your dynamical core. Some spectral ringing isn't strictly problematic as it's the equivalent of step-wise changes of smooth functions in grid-point models but given that this particularly projects onto precipitation your model clearly seems to react to "fake mountains" somehow. I suggest to include a discussion on this and test and report on the associated stability. A little bit of ringing is not a problem but a lot of ringing can easily make the model unusable through numerical instability.
Kudos to your efforts and many thanks for writing them up.
Minor points:
l30: Many/most operational NWP models use a hand-written adjoint? It's still tedious, error-prone and can massively slow down code development or require additional developers to maintain so your point still stands.
l48: I guess you argue that the gained stability in hybrid models over MLWP is because NeuralGCM can be integrated for decades while e.g. GraphCast/GenCast can't? I don't disagree with this but it's still an apple to pears comparison as these models use vastly different ML architectures too. A fair comparison was done in https://doi.org/10.1029/2025MS004969
l54: ICON won the Gordon Bell prize with a heterogeneous Fortran-CPU-GPU computing so not impossible, but given they required a whole team ended up with less than 50% of lines of code actually describing algorithms (most is OpenMP/MPI/CUDA boilerplate instructions for the compiler, wrapped around Fortran loops) and sold this positively as "separation of concerns" you have a very good argument for "extremely costly" and likely also unwieldy and killing a lot of user-friendliness.
l81: geopotential is not a prognostic variable in the hydrostatic equations? Rather diagnosed from surface geopotential (gravity*orography) and the locally vertical virtual temperature profile?
l84: T? are the *spectral* resolutions (or truncation) which can be paired with different *grid* resolutions, presumably the full Gaussian grid with N latitudes yielding linear, quadratic or cubic truncation. Fortran SPEEDY uses T31 with a 96x48 full Gaussian grid which is a quadratic truncation. Your sentences here could be clearer on this?
l95: Are these the same or different from Dinosaur or Fortran SPEEDY?
l100: The Fortran 90 version was written by Sam Hatfield, maybe give him credit?
l126: Fortran SPEEDY has no daily cycle, have you changed that?
l148: The time scales (or equivalently a normalization of the diffusion operator) should also depend on the horizontal resolution, is this currently done?
l152: These appear in the code but aren't documented in Fortran SPEEDY, I also don't know why there were introduced. SpeedyWeather left them out.
l172: Why is 'forcing' and option of model.run and not of the model constructor 'Model'?
l182: Most of the forcing (in terms of energy) comes from solar radiation?
l184: I thought a 30min time step are default?
l193: What is this wind drag? Only in the PBL? Why did you leave it out?
l207: Is this because you apply a much stronger diffusion? implicit 4th order Laplacian vs explicit 2nd order (or even 1st for divergence)?
l228: Why only sometimes? Otherwise you are violating CFL?
l229: Do you mean undersaturation of GPU/TPU threads at lower resolution?
l232: Give these timings in simulated year per day (SYPD) as it's the more common unit for performance?
Fig. 5: If you put simulation period on the xaxis the this graph actual measures the overhead of initializating/compiling/precomputing (depends on what happens when you hit model.run) rather than the actual performance after that? Given you use a log yaxis that information should be in the vertical offset of the lines in the limit? I believe this is the actually interesting performance metric. The initial overhead may depend a lot on what's done at construction via 'Model' and what's done at initialization 'model.run'? Which is still useful information to be added! E.g. can a user run many performant 10-day simulations starting from varying initial conditions? Atm there seems to be quite an overhead for doing that!
Fig. 6: Y axis in SYPD? Add CPU, GPU, TPU performance to inform about where low/high resolutions are faster on CPU vs GPU/TPU?
l241: See also comment on NWP adjoints above.
l266: I agree with you that such a simple example is the first evidence that the method works but maybe note that for your example only the sign of the gradient has to be correct for the optimization to work?
Fig 8/9: I'm missing a very brief discussion on the results of the sensitivity analysis. I agree that a proper analysis is beyond the scope of this paper but a) have you checked that the gradients are correct with finite differences? and b) it looks like there's some dynamically consistent response in the system that could be briefly explained, e.g. is the wind geostrophically adjusting to a higher geopotential due to higher SST? This would also provide more evidence on the correctness.

Citation: https://doi.org/10.5194/egusphere-2025-6266-CC3
- AC3: 'Reply on CC3', Duncan Watson-Parris, 03 Apr 2026
  
  The comment was uploaded in the form of a supplement: https://egusphere.copernicus.org/preprints/2026/egusphere-2025-6266/egusphere-2025-6266-AC3-supplement.pdf
  
  Citation: https://doi.org/10.5194/egusphere-2025-6266-AC3

Viewed

Total article views: 1,411 (including HTML, PDF, and XML)

HTML	PDF	XML	Total	BibTeX	EndNote
802	570	39	1,411	25	37

HTML: 802
PDF: 570
XML: 39
Total: 1,411
BibTeX: 25
EndNote: 37

Views and downloads (calculated since 26 Jan 2026)

Month	HTML	PDF	XML	Total
Jan 2026	334	178	7	519
Feb 2026	209	216	8	433
Mar 2026	199	149	13	361
Apr 2026	60	27	11	98

Cumulative views and downloads (calculated since 26 Jan 2026)

Month	HTML	PDF	XML	Total
Jan 2026	334	178	7	519
Feb 2026	209	216	8	433
Mar 2026	199	149	13	361
Apr 2026	60	27	11	98

Viewed (geographical distribution)

Total article views: 1,397 (including HTML, PDF, and XML) Thereof 1,397 with geography defined and 0 with unknown origin.

Country	#	Views	%

Latest update: 11 Apr 2026

Short summary

We introduce version 1.0 of the JAX Circulation Model (JCM), an open-source atmosphere model. JCM is written in JAX, a framework for high-performance Python code that supports automatic differentiation (automated calculation of how sensitive any program output is to any input). JCM's differentiability and modular design make it easier to train, test, and combine physical-theory-based and data-driven model components, thus providing a flexible and modern platform to facilitate climate research.


Total:	0
HTML:	0
PDF:	0
XML:	0