Veris: Fast & Efficient Sea-Ice Modeling in Python with GPU Acceleration

Abstract. Climate models are typically developed in Fortran, due to its long-standing use in scientific computing and its excellent computational performance. While Python offers significant advantages in terms of code readability, maintainability, and the availability of libraries and tools, the performance gap between Python and Fortran has limited Python's use in large scale climate simulations. The JAX library for array based computing addresses this gap by enabling Just-In-Time compilation, which significantly enhances execution speed. As climate models iterate over numerous time steps, using compiled functions substantially increases the performance of the model. Furthermore, JAX supports both CPU and GPU execution, allowing models to leverage the high parallelism of GPUs. Given that climate models are highly parallelizable, GPUs offer a more efficient alternative to CPUs, both in terms of performance and energy consumption, thereby reducing the computational carbon footprint. This work presents Veris, a sea ice model that is a Fortran to Python translation of the sea ice component of the general circulation model MITgcm. Benchmark tests on a grid with 10⁶ grid cells show that Veris with JAX as the backend is only 1.7 times slower than the Fortran reference. When running on a high-end GPU, Veris matches the performance of the parallelized Fortran reference running on 45 CPUs on 45 nodes or 224 CPUs on 2 nodes. Additionally, Veris can be coupled with the Python-based ocean model Veros to form a fully Python-based coupled sea ice-ocean model, that can be used with large grids in HPC-based simulations.