Skip to content

docs: update examples #285

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
romanc merged 8 commits into from
Oct 27, 2025
Merged

docs: update examples #285

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
fa5ec64
Updated NDSL notebooks with clarifications
gmao-ckung Jul 9, 2024
98187cf
Merge remote-tracking branch 'origin/develop' into feature/NDSL_tutorial
romanc Oct 27, 2025
aec4059
fix typos
romanc Oct 27, 2025
8d0b010
cleanup
romanc Oct 27, 2025
78e0f71
whitespace cleanup and typos
romanc Oct 27, 2025
a67e363
Update headers to be pure markdown
romanc Oct 27, 2025
9293f12
update error message
romanc Oct 27, 2025
4227ca2
update error message
romanc Oct 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 68 additions & 42 deletions examples/NDSL/01_gt4py_basics.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,34 +4,32 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"# **GT4Py Tutorial : Stencil Basics**\n",
"# GT4Py Tutorial: Stencil Basics\n",
"\n",
"## **Introduction**\n",
"## Introduction\n",
"\n",
"This notebook will show how to create a simple GT4Py stencil that copies data from one variable to another.\n",
"\n",
"### **Notebook Requirements**\n",
"### Notebook Requirements\n",
"\n",
"- Python v3.11.x to v3.12.x\n",
"- Python v3.11.x\n",
"- [NOAA/NASA Domain Specific Language Middleware](https://github.com/NOAA-GFDL/NDSL)\n",
"- `ipykernel==6.1.0`\n",
"- [`ipython_genutils`](https://pypi.org/project/ipython_genutils/)\n",
"\n",
"### **Quick GT4Py (Cartesian version) Overview**\n",
"### Quick GT4Py (Cartesian version) Overview\n",
"\n",
"GT4Py is a Domain Specific Language (DSL) in Python that enables a developer to write stencil computations. Compared to simply running under Python, GT4Py achieves performance when the Python code is translated and compiled into a lower level language such as C++ and CUDA, which enables the codebase to execute on a multitude of architectures. In this notebook, we will cover the basics of creating GT4Py stencils and demonstrate several intracies of the DSL. Additional information about GT4Py can be found at the [GT4Py site](https://gridtools.github.io/gt4py/latest/index.html). One small note is that this tutorial covers and uses the Cartesian version of GT4Py and not the unstructured version.\n",
"GT4Py is a Domain Specific Language (DSL) in Python that enables a developer to write stencil computations. Compared to simply running under Python, GT4Py achieves performance when the Python code is transformed and compiled into a lower level language such as C++ and CUDA. This code transformation capability also enables GT4Py-based code to execute on multiple architectures. In this notebook, we will cover the basics of creating GT4Py stencils and demonstrate several intricacies of the DSL. Additional information about GT4Py can be found in the [GT4Py documentation](https://gridtools.github.io/gt4py/latest/index.html). One small note is that this tutorial covers and uses the Cartesian version of GT4Py and not the unstructured (\"next\") version.\n",
"\n",
"### **GT4Py Parallel/Execution Model**\n",
"### GT4Py Parallel/Execution Model\n",
"\n",
"Within a 3-dimensional domain, GT4Py considers computations in two parts. If we assume an `(I,J,K)` coordinate system as a reference, GT4Py separates computations in the Horizontal (`IJ`) spatial plane and Vertical (`K`) spatial interval. In the Horizontal spatial plane, computations are implicitly executed in parallel, which also means that there is no assumed calculation order within the plane. In the Vertical spatial interval, comptuations are specified by an iteration policy that will be discussed later through examples.\n",
"Within a 3-dimensional domain, GT4Py considers computations in two parts. If we assume an `(I,J,K)` coordinate system as the reference, GT4Py separates computations in the horizontal (`IJ`) spatial plane and vertical (`K`) spatial dimension. In the horizontal spatial plane, computations are implicitly executed in parallel, which also means that there is no assumed calculation order within the plane. In the vertical spatial dimension, computations are specified by an iteration policy that will be discussed later through examples.\n",
"\n",
"Another quick note is that the computations are executed sequentially in the order they appear in code.\n",
"Another quick note is that the stencil computations are executed sequentially in the order they appear in code.\n",
"\n",
"## **Tutorial**\n",
"## Tutorial\n",
"\n",
"### **Copy Stencil example**\n",
"### Copy Stencil example\n",
"\n",
"To demonstrate how to implement a GT4Py stencil, we'll step through an example that copies the values of one array into another array. First, we import several packages. "
"To demonstrate how to implement a GT4Py stencil, we'll step through a code example that copies the values of one array into another array. First, we import several packages."
]
},
{
Expand All @@ -50,7 +48,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"As we walk through the example, we'll highlight different terms and such from the imported packages. Let's first define, in GT4Py terms, two arrays of size 5 by 5 by 2 (dimensionally `I` by `J` by `K`). These arrays are defined using a `Quantity` object, an NDSL data container for physical quantities. More detailed information about the `Quantity` object and its arguments can be found from the [`Quantity` docstring](https://github.com/NOAA-GFDL/NDSL/blob/develop/ndsl/quantity.py#L270). To make debugging easier, the `numpy` backend will be used."
"As we walk through the example, we'll highlight different terms and such from the imported packages. Let's first define, in GT4Py terms, two arrays of size 5 by 5 by 2 (dimensionally `I` by `J` by `K`). These arrays are defined using a `Quantity` object, an NDSL data container for physical quantities. More detailed information about the `Quantity` object and its arguments can be found from the [`Quantity` documentation](https://noaa-gfdl.github.io/NDSL/docstrings/quantity/quantity/#quantity.quantity.Quantity.__init__). To make debugging easier, we'll use the `numpy` backend."
]
},
{
Expand Down Expand Up @@ -82,7 +80,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"We will next create a simple GT4Py stencil that copies values from one input to another. A stencil will look like a Python subroutine or function except that it uses specific GT4Py functionalities."
"Next, we will next create a simple GT4Py stencil that copies values from one argument to another. A stencil will look like a Python function except that it uses specific GT4Py functionalities."
]
},
{
Expand All @@ -92,7 +90,7 @@
"outputs": [],
"source": [
"@stencil(backend=backend)\n",
"def copy_stencil(input_field: FloatField, output_field: FloatField):\n",
"def copy_stencil(input_field: FloatField, output_field: FloatField) -> None:\n",
" with computation(PARALLEL), interval(...):\n",
" output_field = input_field"
]
Expand All @@ -101,18 +99,44 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"As mentioned before, GT4Py (cartesian version) was designed for stencil-based computation. Since stencil calculations generally are localized computations, GT4Py stencils are written using variables and the variable's relative location if it's an array. If there are no indices in brackets next to a GT4Py type (such as `FloatField`), it's implied to be at the [0] (for 1-dimension), [0,0] (for 2-dimension), or [0,0,0] (for 3-dimension) location. For the simple example `copy_stencil`, the value of `input_field` simply gets copied to `output_field` at every point in the domain of interest.\n",
"As mentioned before, GT4Py (the cartesian version) was designed for stencil-based computation. Since stencil calculations generally are localized computations, GT4Py stencils are written using variables and the variable's **relative location** if the variable is an array. For example, a stencil-based [Laplacian](https://en.wikipedia.org/wiki/Discrete_Laplace_operator#Finite_differences) calculation can be implemented as follows."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def laplacian(input_field: FloatField, output_field: FloatField):\n",
" with computation(PARALLEL), interval(...):\n",
" output_field = (\n",
" -4.0 * input_field[0, 0, 0]\n",
" + input_field[1, 0, 0]\n",
" + input_field[-1, 0, 0]\n",
" + input_field[0, 1, 0]\n",
" + input_field[0, -1, 0]\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In the `laplacian` stencil, every value of `output_field` at a particular location is a function of `input_field` at that same location (`input_field[0,0,0]`), the adjacent `input_field` values in the `I`-direction (`input_field[1,0,0]` and `input_field[-1,0,0]`), and the adjacent `input_field` values in the `J`-direction (`input_field[0,1,0]` and `input_field[0,-1,0]`)\n",
"\n",
"We see that this stencil does not contain any explicit loops. As mentioned above in the notebook, GT4Py has a particular computation policy that implicitly executes in parallel within an `IJ` plane and is user defined in the `K` interval. This execution policy in the `K` interval is dictated by the `computation` and `interval` keywords. \n",
"If there are no indices in brackets next to a GT4Py type (such as `FloatField`), it's implied to be at the `[0]` (for 1-dimension), `[0,0]` (for 2-dimension), or `[0,0,0]` (for 3-dimension) location. Note again that these locations are **relative**. In the `laplacian` stencil, `output_field` has an implied `[0,0,0]` location, and `input_field[0,0,0]` could have been written simply as `input_field`.\n",
"\n",
"- `with computation(PARALLEL)` means that there's no order preference to executing the `K` interval. This also means that the `K` interval can be computed in parallel to potentially gain performance if computational resources are available.\n",
"We see that both, `copy_stencil` and the `laplacian` stencil, do not contain any explicit loops. As mentioned above, GT4Py has a particular computation policy that implicitly executes in parallel within the `IJ` plane and is user defined in the `K` dimension. This execution policy in the `K` dimension is dictated by the `computation` and `interval` keywords.\n",
"- `with computation(PARALLEL)` means that there's no order preference to executing the `K` dimension. This effectively means that the `K` dimension can be computed in parallel to potentially gain performance if computational resources are available.\n",
"\n",
"- `interval(...)` means that the entire `K` interval is executed. Instead of `(...)`, more specific intervals can be specified using a tuple of two integers. For example... \n",
"- `interval(...)` means that the \"entire\" `K` dimension range is executed (as dictated by the keywords `origin` and `domain`, which will be covered [later in the notebook](#setting-domain-subsets-in-a-stencil-call)). Instead of `(...)`, more specific intervals can be specified using a tuple of two integers. For example...\n",
"\n",
" - `interval(0,2)` : The interval `K` = 0 to 1 is executed.\n",
" - `interval(0,-1)` : The interval `K` = 0 to N-2 (where N is the size of `K`) is executed.\n",
" - `interval(0, 2)` : The interval `K` = 0 to 1 is executed.\n",
" - `interval(0, -1)` : The interval `K` = 0 to N-2 (where N is the size of `K`) is executed.\n",
" - Note: The actual intervals for the above two `interval` are also dictated by [`origin` and `domain`](#setting-domain-subsets-in-a-stencil-call)\n",
"\n",
"The decorator `@stencil(backend=backend)` (Note: `stencil` comes from the package `ndsl.dsl.gt4py`) converts `copy_stencil` to use the specified `backend` to \"compile\" the stencil. `stencil` can also be a function call to create a stencil object."
"The decorator `@stencil(backend=backend)` (Note: `stencil` comes from the package `ndsl.dsl.gt4py`) converts `copy_stencil` to use the specified `backend` to \"compile\" the stencil. `stencil` can also be a function call to create a stencil object."
]
},
{
Expand All @@ -128,9 +152,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Note that the input and output parameters to `copy_stencil` are of type `FloatField`, which can essentially be thought of as a 3-dimensional NumPy array of `float` types.\n",
"Note that the input and output parameters to `copy_stencil` and `laplacian` are of type `FloatField`, which can essentially be thought of as a 3-dimensional NumPy array of `float` types.\n",
"\n",
"`plot_field_at_kN` plots the values within the `IJ` plane at `K = 0` if no integer is specified or at `K` equal to the integer that is specified as an argument. As we can see in the plots below, `copy_stencil` copies the values from `qty_in` into `qty_out`."
"`plot_field_at_kN` plots the values within the `IJ` plane at `K = 0` if no integer is specified or at `K` equal to the integer that is specified as an argument. As we can see in the plots below, `copy_stencil` copies the values from `qty_in` into `qty_out`."
]
},
{
Expand All @@ -155,15 +179,17 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Setting domain subsets in a stencil call**\n",
"### Setting domain subsets in a stencil call\n",
"\n",
"GT4Py also allows a subset to be specified from a stencil call and executed in a fashion similar to using `interval(...)` in the `K` dimension. This is done by setting the stencil call's `origin` and `domain` argument.\n",
"\n",
"GT4Py also allows a subset to be specified from a stencil call and executed in a fashion similar to using `interval(...)` in the K interval. This is done by setting the stencil call's `origin` and `domain` argument.\n",
"- `origin`: This specifies the \"starting\" coordinate to perform computations.\n",
"\n",
"- `origin` : This specifies the \"starting\" coordinate to perform computations. \n",
"- `domain`: This specifies the range of the stencil computation based on `origin` as the \"starting\" coordinate.\n",
"\n",
"- `domain` : This specifies the range of the stencil computation based on `origin` as the \"starting\" coordinate (Note: May need to check whether this affects `interval()`)\n",
"- Note: May need to check whether `domain` and `origin` affect `interval()`\n",
"\n",
"If these two parameters are not set, the stencil call by default will iterate over the entire input domain. The following demonstrates the effect of specifying different `origin` and `domain`."
"If these two parameters are not set, the stencil call by default will iterate over the entire input domain. The following demonstrates the effect of specifying different `origin` and `domain`."
]
},
{
Expand Down Expand Up @@ -242,13 +268,13 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **`FORWARD` and `BACKWARD` `computation` keywords and Offset Indexing within a stencil call**\n",
"### `FORWARD` and `BACKWARD` `computation` keywords and Offset Indexing within a stencil call\n",
"\n",
"Besides `PARALLEL`, the developer can specify `FORWARD` or `BACKWARD` as the iteration policy in `K` for a stencil. Essentially, the `FORWARD` policy has `K` iterating consecutively starting from the lowest vertical index to the highest, while the `BACKWARD` policy performs the reverse.\n",
"Besides `PARALLEL`, the developer can specify `FORWARD` or `BACKWARD` as the iteration policy in `K` for a stencil. Essentially, the `FORWARD` policy has `K` iterating consecutively starting from the lowest vertical index to the highest, while the `BACKWARD` policy performs the reverse.\n",
"\n",
"An array-based stencil variable can also have an integer dimensional offset if the array variable is on the right hand side of the `=` for the computation. When a computation is performed at a particular point, an offset variable's coordinate is based on that particular point plus (or minus) the offset in the offset dimension.\n",
"An array-based stencil variable can also have an integer dimensional offset if the array variable is on the right hand side of the `=` for the computation. When a computation is performed at a particular point, an offset variable's coordinate is based on that particular point plus (or minus) the offset in the offset dimension.\n",
"\n",
"The following examples demonstrate the use of these two iteration policies and also offset indexing in the `K` dimension. Note that offsets can also be applied to the `I` or `J` dimension."
"The following examples demonstrate the use of these two iteration policies and also offset indexing in the `K` dimension. Note that offsets can also be applied to the `I` or `J` dimension."
]
},
{
Expand Down Expand Up @@ -336,7 +362,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Regarding offsets, GT4Py does not allow offsets to variables in the left hand side of the `=`. Uncomment and execute the below code to see the error `Assignment to non-zero offsets is not supported.`."
"Regarding offsets, GT4Py does only allow offsets to variables in the left hand side of the `=` in the `K` dimension. Uncomment and execute the below code to see the error message \"Assignment to non-zero offsets is not supported in IJ\"."
]
},
{
Expand All @@ -355,9 +381,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Limits to offset : Cannot set offset outside of usable domain**\n",
"### Limits to offset : Cannot set offset outside of usable domain\n",
"\n",
"Note that there are limits to the offsets that can be applied in the stencil. An error will result if the specified shift results attemps to read data that is not available or allocated. In the example below, a shift of -2 in the `J` axis will shift `field_in` out of its possible range in `J`."
"Note that there are limits to the offsets that can be applied in the stencil. An error will result if the specified shift results attempts to read data that is not available or allocated. In the example below, a shift of -2 in the `J` axis will shift `field_in` out of its possible range in `J`."
]
},
{
Expand Down Expand Up @@ -406,9 +432,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **`if/else` statements**\n",
"### `if/else` statements\n",
"\n",
"GT4Py allows for `if/else` statements to exist within a stencil. The following simple example shows a stencil `stencil_if_zero` modifing values of `in_out_field` depending on its initial value."
"GT4Py allows for `if/else` statements to exist within a stencil. The following simple example shows a stencil `stencil_if_zero` modifying values of `in_out_field` depending on its initial value."
]
},
{
Expand Down Expand Up @@ -463,9 +489,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Function calls**\n",
"### Function calls\n",
"\n",
"GT4Py also has the capability to create functions in order to better organize code. The main difference between a GT4Py function call and a GT4Py stencil is that a function does not (and cannot) contain the keywords `computation` and `interval`. However, array index referencing within a GT4py function is the same as in a GT4Py stencil.\n",
"GT4Py also has the capability to create functions in order to better organize code. The main difference between a GT4Py function call and a GT4Py stencil is that a function does not (and cannot) contain the keywords `computation` and `interval`. However, array index referencing within a GT4py function is the same as in a GT4Py stencil.\n",
"\n",
"GT4Py functions can be created by using the decorator `function` (Note: `function` originates from the package `ndsl.dsl.gt4py`)."
]
Expand Down
Loading