GT4Py: GridTools for Python

GT4Py is a Python library for generating high performance implementations of stencil kernels from a high-level definition using regular Python functions. GT4Py is part of the GridTools framework, a set of libraries and utilities to develop performance portable applications in the area of weather and climate modeling.

NOTE: The gt4py.next subpackage contains a new version of GT4Py which is not compatible with the current stable version defined in gt4py.cartesian. The new version is still experimental.

📃 Description

GT4Py is a Python library for expressing computational motifs as found in weather and climate applications. These computations are expressed in a domain specific language (GTScript) which is translated to high-performance implementations for CPUs and GPUs.

The DSL expresses computations on a 3-dimensional Cartesian grid. The horizontal axes (I, J) are always computed in parallel, while the vertical (K) can be iterated in sequential, forward or backward, order. Cartesian offsets are expressed relative to a center index.

In addition, GT4Py provides functions to allocate arrays with memory layout suited for a particular backend.

The following backends are supported:

• numpy: A vectorized python backend
• gt:cpu_ifirst: GridTools C++ CPU backend using I-first data ordering
• gt:cpu_kfirst: GridTools C++ CPU backend using K-first data ordering
• gt:gpu: GridTools backend for CUDA
• dace:cpu: Dace code-generated CPU backend
• dace:cpu_kfirst: Dace code-generated CPU backend using K-first data ordering
• dace:gpu: Dace code-generated GPU backend
• debug: A pure python backend used for prototyping new features

🚜 Installation

GT4Py can be installed as a regular Python package using uv, pip or any other PEP-517 compatible frontend. We strongly recommended to useuv to create and manage virtual environments for your own projects.

⚙ Configuration

Other useful available environment variables are:

• CUDA_ARCH: Set the compute capability of the NVIDIA GPU if it is not detected automatically by cupy.
• CXX: Set the C++ compiler.
• GT_CACHE_DIR_NAME: Name of the compiler's cache directory (defaults to .gt_cache)
• GT_CACHE_ROOT: Path to the compiler cache (defaults to ./)

More options and details are available in config.py.

🛠 Development Instructions

Follow the installation instructions below to initialize a development virtual environment containing an editable installation of the GT4Py package. Make sure you read the CONTRIBUTING.md and CODING_GUIDELINES.md documents before you start working on the project.

Development Environment Installation using uv

GT4Py uses the uv project manager for the development workflow. uv is a versatile tool that consolidates functionality usually distributed across different applications into subcommands.

• The uv pip subcommand provides a fast Python package manager, emulating pip.
• The uv export | lock | sync subcommands manage dependency versions in a manner similar to the pip-tools command suite.
• The uv init | add | remove | build | publish | ... subcommands facilitate project development workflows, akin to hatch.
• The uv tool subcommand serves as a runner for Python applications in isolation, similar to pipx.
• The uv python subcommands manage different Python installations and versions, much like pyenv.

We require a reasonably recent version of uv, which can be installed in various ways (see its installation instructions), with the recommended method being the standalone installer:

$ curl -LsSf https://astral.sh/uv/install.sh | sh

Once uv is installed in your system, it is enough to clone this repository and let uv handling the installation of the development environment.

# Clone the repository
git clone https://github.com/gridtools/gt4py.git
cd gt4py

# Let uv create the development environment at `.venv`.
uv sync --extra cartesian  # for gt4py.cartesian
uv sync --extra next       # for gt4py.next

# The `--extra cartesian` (or `--extra next`) option tells uv to install
# the optional run-time dependencies of gt4py.cartesian (or gt4py.next),
# and thus it is not strictly necessary.
# The 'dev' group with the core development dependencies is installed by default.

# Finally, activate the virtual environment and start writing code!
source .venv/bin/activate

The newly created venv is a standard Python virtual environment preconfigured with all necessary runtime and development dependencies. Additionally, the gt4py package is installed in editable mode, allowing for seamless development and testing. To install new packages in this environment, use the uv pip subcommand which emulates the pip interface and is generally much faster than the original pip tool (which is also available within the venv although its use is discouraged).

The pyproject.toml file contains both the definition of the gt4py Python distribution package and the settings of the development tools used in this project, most notably uv, ruff, and mypy. It also contains dependency groups (see PEP 735 for further reference) with the development requirements listed in different groups (build, docs, lint, test, typing, ...) and collected together in the general dev group, which gets installed by default by uv as mentioned above.

Development Tasks (/scripts)

Recurrent development tasks like bumping versions of used development tools or required third party dependencies have been collected as separate modules in the /scripts folder, and collected as subcommands of the ./scripts/run tool (see scripts/README.md for the conventions). Read the tool help for a brief description of every task and always use this tool to update the versions and sync the version configuration across different files (e.g. pyproject.toml and .pre-commit-config.yaml).

📖 Documentation

GT4Py uses the Sphinx tool for the documentation. To build browseable HTML documentation, install the required tools provided in the docs dependency group:

uv sync --group docs --extra standard  # or --group dev

(Note that most likely these tools are already installed in your development environment, since the docs group is included in the dev group, which installed by default by uv sync if no dependency groups are specified.)

Once the requirements are already installed, then build the docs using:

cd gt4py/docs/user/cartesian
make html  # run 'make help' for a list of targets

⚖️ License

GT4Py is licensed under the terms of the BSD-3-Clause.