pyses 0.0.3

A highly readable, well documented, well tested set of atmospheric dynamical cores with support for variable resolution meshes, as well as automatic differentiation for machine learning and data assimilation.

Overview

The purpose of this project is to create a highly readable, well documented, well tested atmospheric dynamical core with support for variable resolution meshes, as well as automatic differentiation for machine learning and data assimilation. This project prioritizes code readability and maintainability, in the sense that code that runs 10% slower but is much easier for a second-year graduate student to understand and modify is better than its inaccessible optimized counterpart. We want to minimize external dependencies and, insofar as it is possible, create a codebase that is entirely written in python. Given the constraints of these design decisions, it is unlikely that the resulting dynamical core will scale to hundreds or thousands of nodes on an HPC computing system. This aligns with our stated goal of making atmospheric modeling accessible.

Quickstart with uv

Common steps

• Run git clone [e-mail]:OkayHughes/pyses.git and navigate to the pyses directory.
• Install uv, e.g. pip install uv in your python environment. uv is a modern environment manager for Python that we use for development and testing.

MPI-dependent steps

If you don't have MPI installed, run:

• Run uv sync --group mpich --group dev --group jax --group torch
• If you don't plan to use either jax or torch, you can omit one or both of those groups.

If you have a system MPI installed (e.g. on HPC systems), run:

• Run which mpicc
• Create a .env file containing export MPICC=${MPICC_LOCATION}
• Run uv sync --env-file .env --group dev --group jax --group torch
• If you don't plan to use either jax or torch, you can omit one or both of those groups.

Run tests

• Navigate to the tests directory and run bash run_test.sh. These should catch if there are problems with CPU configurations of pyses, and test whether there are issues with your MPI environment.

Accelerator support

Most scientific code can be easily written in a (nearly) purely functional programming style. Consequently, this means that the codebase can be written to satisfy the requirements of the Jax library's just-in-time compilation and automatic differentiation, while retaining the ability to run with array/tensor operations provided by PyTorch or Numpy. Due to Google's history of abruptly discontinuing widely used software products, we have chosen to future proof this code base by ensuring that GPU parallelism (and ideally automatic differentiation) can be sourced from any library that provides an array implementation that resembles Numpy.ndarray, along with a list of array operations that is enumerated in the documentation.

Jax support

The Jax configuration of the code is significantly more performant than either the Numpy or Torch configurations. This is because the allowable control flow constructs of jax.jit are significantly less constraining than, e.g., numba or torch.compile. Consequently, the programming style of this project treats the idiosyncracies of the Jax functional programming model as authorative. This means that Numpy and Torch performance suffers, as Jax disallows assignment of slices of arrays after initial assignment, so code like

x = np.eye((3, 3))
x[:, 0] = x[:, 1]

is entirely disallowed outside of initialization. Instead, this would be written

x = np.eye((3, 3))
x = np.stack((x[:, 1], x[:, 1:]), axis=1)

which appears almost deliberately tailored to force Numpy/Torch view semantics to copy the data in x.

PyTorch support

PyTorch can also provide automatic differentiation capabilities and GPU parallelism. This is used as a fallback, as the dominance of PyTorch at so-called "AI" companies makes it more likely to be supported in the future. Therefore, we are trying to ensure that code runs with PyTorch as a backend, but we don't typically optimize for PyTorch performance.

Policy on intellectual property

The view of the maintainer is that since the training data used to train LLMs was obtained without the consent of the people who made it, LLMs trained on this data are structurally incapable of determining when they are commiting plagiarism or theft. Contributing LLM generated code should be treated with the utmost care and consideration for other people.

Shoulders of giants disclaimer

The core functionality of this dynamical core is not primarily my work, and is mostly a framework that helps unify the Higher Order Methods Modeling Environment (HOMME) and the Community Atmosphere Model-Spectral Element dynamical cores. I've been so fortunate to work with Drs. Mark Taylor, Oksana Guba, and Peter Lauritzen, and I would not have been able to write this codebase without their mentorship. None of this mentorship would have happened at all if my advisor, Professor Christiane Jablonowski, hadn't helped me find a place for myself in the atmospheric modeling community, and helped develop my taste for science-driven dynamical core development.