Articles | Volume 19, issue 3
https://doi.org/10.5194/gmd-19-1103-2026
https://doi.org/10.5194/gmd-19-1103-2026
Model description paper
 | 
03 Feb 2026
Model description paper |  | 03 Feb 2026

LEX v1.6.0: a new large-eddy simulation model in JAX with GPU acceleration and automatic differentiation

Xingyu Zhu, Yongquan Qu, and Xiaoming Shi

Related authors

Exploring the aerosol activation properties in coastal shallow convection using cloud- and particle-resolving models
Ge Yu, Yueya Wang, Zhe Wang, and Xiaoming Shi
Atmos. Chem. Phys., 25, 7527–7542, https://doi.org/10.5194/acp-25-7527-2025,https://doi.org/10.5194/acp-25-7527-2025, 2025
Short summary

Cited articles

Bezgin, D. A., Buhendwa, A. B., and Adams, N. A.: JAX-Fluids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows, Comput. Phys. Commun., 282, 108527, https://doi.org/10.1016/j.cpc.2022.108527, 2023. a
Bezgin, D. A., Buhendwa, A. B., and Adams, N. A.: JAX-Fluids 2.0: Towards HPC for differentiable CFD of compressible two-phase flows, Comput. Phys. Commun., 308, 109433, https://doi.org/10.1016/j.cpc.2024.109433, 2025a. a
Bezgin, D. A., Buhendwa, A. B., Schmidt, S. J., and Adams, N. A.: ML-ILES: End-to-end optimization of data-driven high-order Godunov-type finite-volume schemes for compressible homogeneous isotropic turbulence, J. Comput. Phys., 522, 113560, https://doi.org/10.1016/j.jcp.2024.113560, 2025b. a
Blossey, P. N., Bretherton, C. S., Cheng, A., Endo, S., Heus, T., Lock, A. P., and van der Dussen, J. J.: CGILS Phase 2 LES intercomparison of response of subtropical marine low cloud regimes to CO2 quadrupling and a CMIP3 composite forcing change, J. Adv. Model. Earth Sy., 8, 1714–1726, https://doi.org/10.1002/2016MS000765, 2016. a
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q.: JAX: composable transformations of Python+NumPy programs, GitHub, http://github.com/jax-ml/jax (last access: 30 January 2026), 2018. a, b
Download
Short summary
By using the newly developed Python library JAX, we developed a fast and differentiable large-eddy simulation model, named LEX. Evaluated with a warm bubble case, LEX maintains high accuracy as the Cloud Model 1. With the hardware acceleration and better numerical stability, LEX can be quite faster. To report its differentiability, we further trained a deep learning-based parameterization scheme. The newly trained model can surpass the conventional scheme and get proper forecast results.
Share