The software productivity gap problem can be described as an issue with the diverse set of skills required by climate modelers to implement a production-grade climate and weather model, from discretizing the underlying physical equations all the way to the optimization details of a given hardware architecture. This process can be roughly split in two:

climate and weather modeling – scientific motivation, algorithmic design, and numerical implementation;

There is overlap between these; some algorithmic designs are better suited to certain hardware architectures, and implementation details can be changed to improve model speed, but this breakdown is a useful heuristic for the development of weather and climate models.

The design of a modern domain-specific language needs to respond to both classes of users. For climate modelers a DSL should

be easy to use, complete with debugging tools and simple methods to extracting scientific results;

For the performance developers it should

leverage a proven host language to strengthen basic development of a compiler tool chain and work on solid ground;

Python stands as a good candidate language to host a DSL given its robustness, ease of use, wide adoption across a range of research and industry groups, large ecosystem of pre-existing tools and packages, and capacity for introspection as an interpreted language. Furthermore, in the weather and climate community Python has established itself as the lingua franca for analysis, visualization, and post-processing. But Python has serious limitations in terms of execution performance: designed primarily as a scripting language, Python alone cannot achieve the performance required by workloads running on large HPC systems.

The solution is to escape Python at runtime and leverage compiled code. This is a technique already heavily used by frameworks such as Cython and Numba see, e.g.,. The DSL compiler is responsible for translating and compiling (transpiling) code written in Python into another, more performance-oriented language. The generated code can be tailored and optimized to a specified hardware target via hardware-specific compiler back ends.

For this work we use the GT4Py DSL compiler because the DSL provides a single-source solution for portable performance with the correct abstractions for weather and climate models. There are, however, other approaches using Python at different levels of abstraction that are worth discussing.

Cython is an optimizing static compiler based on Python and the Pyrex language. It is a powerful tool for accelerating Python code, but it lacks the portability and high-level abstractions we seek. Other packages provide portability but still lack the domain-level abstractions to simplify the language for climate scientists. An example of such a package is Numba , which is a just-in-time (JIT) compiler for Python and NumPy code. JAX similarly lacks the weather and climate abstractions we desire in the front-end DSL language but brings with it convenient features such as adjoint capability. Other application-building frameworks such as Exasim , FEniCS , and Dedalus mostly operate on the partial-differential-equation (PDE) level, which allows for more automated model development at the expense of flexibility in the discretization. In order to faithfully reproduce all aspects of the FV3 dynamical core and GFDL cloud microphysics, we required greater flexibility than these frameworks provided. Thus, the DSL we are targeting works at a lower level on the mathematical representation after the PDE has been discretized.

Our implementation of the domain-specific language described above is called GridTools for Python, or GT4Py. Developed in partnership with the Swiss National Supercomputing Centre (CSCS), it defines a DSL on top of Python. The code is analyzed and compiled into a C++ or CUDA executable that is bound back to the original Python, creating a seamless experience for the modelers but enabling fast and optimized execution (Fig. ). On top of those performance back ends, GT4Py also provides a back end to generate NumPy code, which is useful for debugging and quick development as no compilation is required. One GT4Py back end leverages the DaCe framework to be able to also include regular Python code in between DSL code in the translation.

GT4Py operates through the use of stencils: inside a GT4Py stencil each grid cell in an array is modified according to neighboring cells based on a fixed pattern. In numerical modeling of weather and climate two major computational patterns emerge due to the reliance on three-dimensional structured or unstructured grids: computations with dependencies on the local horizontal neighborhood of a grid cell and vertical solvers with column dependencies as shown in Fig. . These two stencil computational patterns therefore form the basis of GT4Py's design.

GT4Py stencils treat the horizontal dimensions differently from the vertical: they always execute in parallel over the entire horizontal domain. In the vertical dimension the order of execution may be specified with the computation keyword, and the vertical range is set with interval. A horizontal stencil can be written with computation(PARALLEL) to parallelize over the vertical domain as well if there are no vertical loop-carried dependencies. A column-based stencil that calculates up or down the k dimension, on the other hand, can be implemented sequentially in k with computation(FORWARD) or BACKWARD, respectively. Each stencil can also have multiple computation and interval blocks for flow control. Because stencils are applied to each point in a 3D grid, all indexation within a stencil is relative to the current computed grid point; i.e., Array[0,0,0] - Array[1,0,0] takes each array element and subtracts the array element immediately to its right along the x axis. This method of indexing three-dimensional arrays allows the modelers to use the same indexing convention (I, J, K), irrespective of the actual storage layout in memory.

GT4Py also allows zero-cost function calls, enabling more readable and reusable code within models. Extents of the computation are automatically determined by the DSL, including through these function calls. Figure  shows an example stencil function that executes in parallel over the entire vertical dimension. The domain of dependence for intermediate variables (advected_u and advected_v) is automatically determined by GT4Py. In this example, the advection helper functions take a horizontal difference of the u and v contravariant velocity components ub_contra and vb_contra, respectively. As a result, the interpolation function which computes these values is applied on a larger domain than the final operation in the stencil.

In order to implement the FV3 dynamical core we had to extend the stencil concept multiple times, enabling four- and two-dimensional arrays, loops over the k axis, and specialized handling of horizontal subdomains. Because GT4Py is written in Python, those extensions were easy to develop and they will be discussed further in Sects.  and .

GT4Py is able to optimize the performance of stencils; however, some Python overhead remains in the code linking the stencil operations together. In order to resolve this and optimize the entire Pace model we have incorporated DaCe as a back end for Pace. Using this framework both GT4Py stencils and raw Python code are exposed to optimization and transpilation to high-performance C++ or CUDA executables, removing all Python overhead seamlessly for the modelers. Section  contains more information on code optimizations, while details the performance optimizations more thoroughly.

Pace is designed to be modular; each model component of Pace (e.g., dynamics, microphysics, utilities, DSL integration) exists as a separate package. Computationally focused packages like the dynamics contain hierarchies of component modules (Fig. ). These components provide clear boundaries to document and change model behavior. For example, the horizontal transport scheme used by FV3 takes any one-dimensional finite-volume subgrid reconstruction scheme satisfying certain numerical conditions and can extend it to two dimensions. Within Pace, one-dimensional subgrid reconstruction code is contained in the XPPM and YPPM modules (X and Y piecewise parabolic methods, respectively). These modules take in scalar grid-cell-mean values and Courant numbers (speed as a fraction of grid-cell width) defined on transport interfaces and return the average value of the scalar within the section of grid cell to be advected through the cell interface. The FiniteVolumeTransport class extends these one-dimensional subgrid reconstruction schemes to produce two-dimensional horizontal fluxes. This allows a scientist to modify the behavior of the dynamics by replacing only the XPPM and YPPM components, for example with a cubic reconstruction scheme or a subgrid reconstruction scheme based on machine learning.

To maintain simplicity of the code and facilitate the separation of compile time and runtime, Pace uses a simple object-oriented framework that expresses nearly all of the internal computational components as initializable functions, as shown in Listing . Each of these functions is defined as a Python class with an initialization method and a call method. Initialization performs the necessary setup including memory allocation and compilation of DSL code. Each component's initialization code defines the component's temporary variables, compiles or loads its stencils, and recursively initializes subcomponents. Once initialized, the component may be called the same as any other Python function. This is similar to Fortran model structures where computational code is paired with initialize and finalize subroutines.

Stencils are compiled when components are initialized, using only explicitly passed configuration data. Pace uses the factory pattern (through the stencil_factory argument) to reduce compilation-specific logic as much as possible within computational components. The stencil factory class implements the code responsible for allocating and compiling stencils, allowing model code to instead focus on computational motifs. In lines 8–18 of Listing  we can see this factory at work compiling the compute_kinetic_energy stencil shown in Fig. . The from_dims_halo method takes a stencil function and a set of dimensions (either cell centers or cell interfaces) to execute over and returns a compiled stencil that writes its outputs over the compute domain, in this case on cell corners (x and y interfaces). Output can be extended into the computational halos with an optional compute_halos argument indicating how many halo points to write (not shown). We also see how compile-time constants are passed as externals to the compilation method. These constants can be set by the model configuration, extracted from the domain decomposition, or passed as arguments to the initialization method, and they have the benefit of being treated as compile-time constants. Because configuration settings such as domain-decomposition and namelist settings are now known at compile time, they also provide further avenues for the DSL back end to optimize performance.

The factory pattern used here is incredibly powerful when debugging the model. For example, if a model developer finds that at some point a variable air_pressure used by many routines has gone negative in the model, that developer can temporarily insert code to immediately raise an exception if any stencil writes a negative value to a variable named air_pressure. Or that model developer could alter every stencil to write its inputs and outputs to a netCDF file when executed, while only having to modify the code used for the stencil factory.

The introspection power of Python is used to great effect in the testing code for Pace. For example, we would like to ensure that Python array allocation only happens when initializing our model and not at all when it is called. To do this, we write a test which initializes the model and then replaces the storage allocation routines of GT4Py with routines which raise an exception if called before finally calling the model. If any arrays were allocated at call time, an exception would be raised and the test would fail.

We also want to ensure that the Pace model components are not stateful. The dynamical core, for example, has many temporary storage arrays assigned to it whose initial value should not matter when calling the model. However, a user could easily fail to initialize an array with zeros when they mean to, causing a bug in the model. We have a test which calls the dynamical core with the same state either once or twice and compares the value of all temporary data in the model between these two cases. It does this by dynamically crawling the Python object structure and comparing all array data. If any data differ between those two cases, it not only tells us that there is a bug but also exactly which temporary field has a bug – the one which is accessed first out of the ones which differ. Without such a test, it could take days, weeks, or months for a scientist to find the source of such an error, assuming they notice the presence of the bug.

In the process of writing Pace we needed to extend GT4Py in order to express the Fortran code in the DSL. FV3 discretizes the Earth into a gnomonic cubed sphere . GT4Py could only apply the same operations uniformly across the horizontal domain of a stencil. We added the ability to execute code on stencil subdomains to the GT4Py DSL language in order to perform special handling required near the corners and edges of the cubed-sphere tiles. Specifically, we included horizontal and region as stencil control keywords (similar to computation and interval) as shown in Fig. . with horizontal specifies a horizontally restricted block of code, and region specifies the extent of that subdomain through NumPy-like array slicing in the I and J dimensions. Though it does violate the GridTools concept of a stencil, this allows us to fully port the dynamical core and combine stencils to reproduce a natural amount of Fortran code. This also has the benefit of increasing the readability of the code: in Fortran the corner and edge computations are handled by the common pattern of if-conditionals inside of nested “for loops”, while in Pace the horizontal regions are used almost exclusively for this purpose.

Another example is in the vertical remapping component of the model. The FV3 dynamical core uses a Lagrangian vertical coordinate which is regularly remapped to its original Eulerian grid using the piecewise parabolic method . The remapping process requires a double k loop over the vertical dimension: an inner loop over the deformed vertical levels to sum their contributions to an Eulerian level and an outer loop over those Eulerian levels. When we initially ported the remapping code GT4Py did not support “for loops” or “while loops” inside of stencils, so we wrote the inner loop as a stencil over the deformed k levels and the outer loop in plain Python. While this implementation worked algorithmically, calling a stencil inside of a “for loop” has a large penalty in model speed due to the repeated kernel launches and the removal of the loop structure from the DSL optimization path. To resolve this we added “while loops” to GT4Py stencils. This allows us to consolidate the remapping code into one stencil over the Eulerian k levels, remove the Python “for loop”, and expose the entire remapping scheme to the DSL compiler. This reduced the runtime of the remapping step by over an order of magnitude.

GT4Py is able to cover a large number of algorithmic motifs used in weather and climate models. This section illustrated two examples where the DSL has been extended in order to express motifs which were present in FV3GFS. There is also an ongoing effort to extend GT4Py to unstructured grid computations. Nevertheless, any DSL will naturally be restricted from covering some algorithmic motifs. For GT4Py these include reductions (e.g., computing horizonal integrals over a field), interpolation (e.g., for multi-grid), and search patterns (e.g., looking for a neighboring grid cell with certain properties). Whether these patterns should fall into the scope of the GT4Py DSL or another framework is a design decision, which can be weighed for example against compiling these patterns with DaCe.

Each unit of code is validated by running it with the input data serialized from runs of the Fortran model, and comparing the outputs of our code to the serialized Fortran outputs. If a ported unit is particularly large (e.g., vertical remapping, grid generation) we serialize data from components of the larger unit and test the equivalent Pace components to reduce the complexity of each test. We use two sets of initial atmospheric conditions as our test cases, each with a corresponding configuration namelist. Our “standard” test case is generated from NCEP reanalysis data from 00:00:00Z on 1 August 2016 as in, and the other is the baroclinic instability test case described in . Initial versions of the DSL code were tested on the standard case run on six MPI ranks (one per tile) and a 12 by 12 horizontal grid on each tile face with 79 vertical levels. We also test with a 54-rank domain decomposition, as this gives each tile face a rank for each corner and edge as well as a central rank with no edges. Our tests thus cover the range of edge and corner handling required of a given rank from no corners or edges to all eight corners and edges.

In many cases the Pace outputs can be brought near roundoff error to the Fortran outputs. However, there are some instances where changes to the order of operations or the use of transcendental functions makes such reproducibility impossible, so we must choose a validation threshold for our code. We do this by perturbing the inputs to the Fortran code by small, floating-point differences and comparing the outputs. This leads us to adopt relative differences of 10-14 as our default validation threshold. If we find validation errors greater than our default, we test the components of the offending code to see what operations introduce discrepancies between Pace and Fortran. Often these differences are due to errors introduced during porting and we can reconcile Pace with the Fortran code. If these differences are instead due to algorithmic changes (reordering, etc.) we set the error tolerance for that piece of code based on the underlying Fortran–Python difference. For example, in the vertical remapping scheme FV3 makes use of multiple goto statements for its control flow. These are not available in GT4Py, and so some code has been rearranged to replicate the original algorithm, introducing small deviations from the Fortran outputs.

When we are confident that each component of the model accurately reproduces the Fortran version, we test larger combinations of these units to ensure the implementation of these larger modules also matches Fortran, as illustrated in Fig. . For example, after validating the cubic-spline vertical interpolation code and the code that calculates the contributions from deformed, Lagrangian pressure levels to the remapped Eulerian levels, we then validate the code to vertically remap a single variable from Lagrangian to Eulerian pressure levels. When this code validates, along with the tracer remapping, saturation adjustment, and moist potential temperature adjustment code, we can then validate the entire vertical remapping scheme. In this way we hierarchically assemble and validate Pace against the Fortran code. These validation tests are incorporated into our suite of continuous integration tests to ensure that future developments and code changes do not affect model validity.

Once the full Pace dynamical core passes these unit tests, we run the baroclinic instability test case for 9 model days and compare the results against the same test run in Fortran. In this test case, the dynamical core is initialized with zonally symmetric steady-state winds, on top of which a Gaussian perturbation of the zonal wind is superimposed in the northern mid-latitudes, triggering the evolution of a baroclinic wave. In Fig.  we show the region of interest in a 10 km resolution simulation (960 × 960 × 80 grid cells per tile, run on 150 MPI ranks) of the same instability from Pace and the reference Fortran model. We see that the model behavior is consistent with the reference Fortran model and replicates the results from the highest horizontal resolution in well. Figure  shows the difference in 850 mbar temperature between Pace and Fortran running the test case at 200 km (48 × 48 × 80 grid cells per tile) resolution. Both models were run with six MPI ranks. The Pace dycore matches the Fortran code closely early on, with random errors on the order of 10-11 after 3 model days of integration. Due to the nature of the baroclinic instability these small errors do eventually grow to the order of 10-1 after 9 simulation days (3888 time steps), but even on day 6 the errors are only on the order of 10-5. Based on these results we are confident that Pace accurately reproduces the Fortran model code.

Pace has only been validated using double-precision floating-point arithmetic and values. All performance results shown in the next section have been measured using double precision.

Because GT4Py has multiple back ends for various target architectures there is no one performance number that captures the entirety of our approach. Nonetheless, the aggregate performance across back ends indicates the power of the DSL paradigm through performance portability. With the capability of code generating for specific hardware targets, we are not limited to single, catch-all solutions in our user-facing code. Instead we can generate optimized code for each type of hardware individually through back-end logic, allowing user code to focus on numerical details. As an example, on CPU architectures it usually pays off to use coarse-grained parallelism such that threads are assigned to different subdomains, while on GPUs the parallelization strategy involves having blocks of threads execute subdomains.

Weather and climate models are written with large configuration files, namelists, to allow for flexible use. In order to support all the possible configurations, standard models use many conditionals that can not be resolved at compile time, such as which subgrid reconstruction to use for tracer advection. This limits how aggressively the low-level compiler can optimize the model. Code generation from a DSL allows us to generate code and compile the model for a specific namelist configuration to circumvent this problem.

Furthermore the DSL compiler is able to apply domain-specific optimizations under certain conditions on a per-stencil level before generating code at all. These include the following.

Reduce main memory accesses. Replace 3D temporary fields used for intermediate results that are stored in global memory with smaller buffers that allow for more reuse and faster access.

In order to achieve good performance in a Python-driven environment it is crucial to minimize the overhead moving from the driver language to the compiled language executable. Since our implementation of the model is a series of compiled stencils called from Python, the performance overhead linked to calling the compiled stencils from Python scales with the number of stencils called. A secondary issue is that with the fragmentation into individual stencils we limit the optimization potential of the DSL, as it is limited to the scope of a single stencil for the optimizations described above.

We leverage DaCe to address those shortcomings, providing a full-program optimization framework. With our DaCe back end in GT4Py we are able to compile the entire loop over time steps into a single executable called from Python, thus completely removing all Python to C call overhead during our simulation. With DaCe as our back end we are able to leverage custom optimization for our code, including improving the computational layout, improving where and how memory allocation happens, scheduling computation to improve parallelism as well as increasing data locality, and improving the pressure on global memory. A detailed explanation of our approach and how it affected performance can be found in .

As discussed in Sect. , the cubed-sphere discretization of FV3 requires special finite-difference stencils applied near edges and corners to account for the grid geometry. We added the ability to execute stencils on horizontal subdomains to enable these motifs in GT4Py, which unfortunately has negative performance implications for Pace.

FV3 parallelized the vertical plane K instead of the horizontal IJ ones. Flipping this order of parallelization increases the parallelism manyfold, which is necessary for high throughput on GPUs and other massively parallel architectures. FV3 parallelizes in this way because it is advantageous for their target architectures, which are CPUs with branch prediction and large caches. On CPU architectures instruction divergence inside the tight loops due to FV3's corner and edge handling is not an issue, but there is a limit to what can be done in a single-instruction, multiple-data (SIMD) GPU kernel. Because of their data dependencies these corner and edge operations cannot be executed in parallel with the rest of the code, requiring a separate kernel launch for each of them. This breaks what could be one large and optimized stencil into separate stencils separated by such corner calculations, though there are certain conditions under which we can fuse these operations into the same kernel.

This example shows how algorithmic choices have large impacts on performance. The choice of specialized corner and edge handling made for FV3 fits a CPU architecture well but is suboptimal for GPU architectures. Changes to the algorithm can be made to alleviate this problem, such as a duo-grid approach demonstrated in . Since Pace is intended to be a one-to-one port of FV3 in its current state, this is not a development we are focused on at the moment, though it could be a fruitful direction for other DSL models.

The driver code in Listing  showcases the power of Python as an API language for configuring, compiling, and running computational codes. The command_line entry function reads all information needed to run the model from disk and directly defines the behavior of the driver class. Within the driver class (not shown), the dynamical core and physics components only have access to the configuration and variables they use, which makes it much easier for a new developer of the model to understand what the code does. This is true all the way down through its components, as described in Sect. .

This example also showcases the advantage of working in a widely used language with a vibrant open-source package ecosystem. Python has powerful tools available for defining command-line interfaces. Taking a yaml configuration file and mapping it onto a nested configuration class is as simple as using pyyaml and dacite, which provide this functionality. Python's built-in datetime and timedelta types make it easy to manage model execution time, and external packages such as cftime provide support for a wide range of calendars. Diagnostics storage makes use of the xarray and zarr packages to greatly simplify the code we need to write. This well-established ecosystem of tools maintained by scientists and engineers from a range of disciplines and sectors is incredibly helpful when developing an atmospheric model.

And in turn, having a model written in Python means the tools we have written can be used by anyone who uses the language. This is particularly important given the popularity of Python as a language for scientific analysis. When processing the output of Pace, a scientist has direct access to the tools and numerical code used by the model itself.

One of the advantages of Python is the ecosystem surrounding it. For example, packages such as NumPy and SciPy make it easy to perform common mathematical and numerical operations and manipulate data; Matplotlib enables interactive visualization and creates static or animated images; a Jupyter Notebook allows users to create and share their computational documents; TensorFlow provides end-to-end libraries for machine learning and artificial intelligence applications. With Pace, we can leverage all these tools to make running, processing, and visualizing climate model output all in one Python script.

Pace also enables novel workflows through the use of Jupyter Notebooks. The Pace repository contains an example notebook running one component of Pace. We first initialize a cubed-sphere grid and an idealized atmospheric state, and then we run the tracer advection operation and visualize the results. In this case the analytic zonal wind profile should advect the tracer mass once around the Earth in 12 model days, and any deviation from this indicates a problem with the advection code. This workflow is meant to mimic how a model developer could develop, implement, and debug a new advection scheme. We take advantage of the fact that each component in Pace is modular and skip over running the entire model, which enables rapid prototyping without leaving the Python ecosystem. Beyond development, this capability is useful for teaching, allowing students to inspect elements of the model individually.

Another benefit of Pace is that it facilitates incorporating novel applications and approaches into the model, such as machine learning emulation. Oftentimes machine learning in climate models requires complicated workflows such as calling Fortran from Python or calling Python from Fortran. Because Pace is written in Python, we can train and execute an machine learning (ML) model directly, bypassing difficult infrastructure to pass variables between languages. In Listing  we show Pace's ability to replace a stencil-based microphysics scheme with a pre-trained TensorFlow-based microphysics emulator. Taking advantage of the modularity and separation of initialization and call time of the individual module, we simply load the TensorFlow model during the initialization of the Physics class and replace the call signature of microphysics. Figure  shows an example of the emulator applied to the baroclinic instability test case.