diff --git a/docs/user/next/.gitignore b/docs/user/next/.gitignore deleted file mode 100644 index 10cc078baa..0000000000 --- a/docs/user/next/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -*.py -*.ipynb diff --git a/docs/user/next/QuickstartGuide.md b/docs/user/next/QuickstartGuide.md index c994d0fd57..f8ff64a980 100644 --- a/docs/user/next/QuickstartGuide.md +++ b/docs/user/next/QuickstartGuide.md @@ -5,7 +5,7 @@ jupytext: extension: .md format_name: myst format_version: 0.13 - jupytext_version: 1.14.0 + jupytext_version: 1.16.1 kernelspec: display_name: Python 3 (ipykernel) language: python @@ -70,19 +70,17 @@ b = gtx.as_field([CellDim, KDim], np.full(shape=grid_shape, fill_value=b_value, Additional numpy-equivalent constructors are available, namely `ones`, `zeros`, `empty`, `full`. These require domain, dtype, and allocator (e.g. a backend) specifications. ```{code-cell} ipython3 -from gt4py._core import definitions as core_defs +I = gtx.Dimension("I") +J = gtx.Dimension("J") + array_of_ones_numpy = np.ones((grid_shape[0], grid_shape[1])) -field_of_ones = gtx.constructors.ones( +field_of_ones = gtx.ones( domain={I: range(grid_shape[0]), J: range(grid_shape[0])}, - dtype=core_defs.dtype(np.float64), + dtype=np.float64, allocator=gtx.program_processors.runners.roundtrip.backend ) ``` -_Note: The interface to construct fields is provisional only and will change soon._ - -+++ - #### Field operators Field operators perform operations on a set of fields, for example, elementwise addition or reduction along a dimension. You can write field operators as Python functions by using the `@field_operator` decorator. Field operators cannot have side effects, therefore you cannot modify its arguments within their body. Only a subset of the Python syntax is allowed inside field operators—the library checks for correctness. @@ -91,8 +89,8 @@ Let's see an example for a field operator that adds two fields elementwise: ```{code-cell} ipython3 @gtx.field_operator -def add(a: gtx.Field[gtx.Dims[CellDim, KDim], float64], - b: gtx.Field[gtx.Dims[CellDim, KDim], float64]) -> gtx.Field[gtx.Dims[CellDim, KDim], float64]: +def add(a: gtx.Field[Dims[CellDim, KDim], float64], + b: gtx.Field[Dims[CellDim, KDim], float64]) -> gtx.Field[Dims[CellDim, KDim], float64]: return a + b ``` @@ -116,9 +114,9 @@ This example program below calls the above elementwise addition field operator t ```{code-cell} ipython3 @gtx.program -def run_add(a : gtx.Field[gtx.Dims[CellDim, KDim], float64], - b : gtx.Field[gtx.Dims[CellDim, KDim], float64], - result : gtx.Field[gtx.Dims[CellDim, KDim], float64]): +def run_add(a : gtx.Field[Dims[CellDim, KDim], float64], + b : gtx.Field[Dims[CellDim, KDim], float64], + result : gtx.Field[Dims[CellDim, KDim], float64]): add(a, b, out=result) add(b, result, out=result) ``` @@ -163,9 +161,9 @@ This section approaches the pseudo-laplacian by introducing the required APIs pr The examples related to unstructured meshes use the mesh below. The edges (in blue) and the cells (in red) are numbered with zero-based indices. -| ![grid_topo](connectivity_numbered_grid.svg) | -| :------------------------------------------: | -| _The mesh with the indices_ | +| ![grid_topo](workshop/images/connectivity_numbered_grid.svg) | +| :----------------------------------------------------------: | +| _The mesh with the indices_ | +++ @@ -217,9 +215,9 @@ cell_values = gtx.as_field([CellDim], np.array([1.0, 1.0, 2.0, 3.0, 5.0, 8.0])) edge_values = gtx.as_field([EdgeDim], np.zeros((12,))) ``` -| ![cell_values](connectivity_cell_field.svg) | -| :-----------------------------------------: | -| _Cell values_ | +| ![cell_values](workshop/images/connectivity_cell_field.svg) | +| :---------------------------------------------------------: | +| _Cell values_ | +++ @@ -248,11 +246,11 @@ Pay attention to the syntax where the field offset `E2C` can be freely accessed ```{code-cell} ipython3 @gtx.field_operator -def nearest_cell_to_edge(cell_values: gtx.Field[gtx.Dims[CellDim], float64]) -> gtx.Field[gtx.Dims[EdgeDim], float64]: +def nearest_cell_to_edge(cell_values: gtx.Field[Dims[CellDim], float64]) -> gtx.Field[Dims[EdgeDim], float64]: return cell_values(E2C[0]) @gtx.program -def run_nearest_cell_to_edge(cell_values: gtx.Field[gtx.Dims[CellDim], float64], out : gtx.Field[gtx.Dims[EdgeDim], float64]): +def run_nearest_cell_to_edge(cell_values: gtx.Field[Dims[CellDim], float64], out : gtx.Field[Dims[EdgeDim], float64]): nearest_cell_to_edge(cell_values, out=out) run_nearest_cell_to_edge(cell_values, edge_values, offset_provider={"E2C": E2C_offset_provider}) @@ -262,9 +260,9 @@ print("0th adjacent cell's value: {}".format(edge_values.asnumpy())) Running the above snippet results in the following edge field: -| ![nearest_cell_values](connectivity_numbered_grid.svg) | $\mapsto$ | ![grid_topo](connectivity_edge_0th_cell.svg) | -| :----------------------------------------------------: | :-------: | :------------------------------------------: | -| _Domain (edges)_ | | _Edge values_ | +| ![nearest_cell_values](workshop/images/connectivity_numbered_grid.svg) | $\mapsto$ | ![grid_topo](workshop/images/connectivity_edge_0th_cell.svg) | +| :--------------------------------------------------------------------: | :-------: | :----------------------------------------------------------: | +| _Domain (edges)_ | | _Edge values_ | +++ @@ -274,12 +272,12 @@ Similarly to the previous example, the output is once again a field on edges. Th ```{code-cell} ipython3 @gtx.field_operator -def sum_adjacent_cells(cells : gtx.Field[gtx.Dims[CellDim], float64]) -> gtx.Field[gtx.Dims[EdgeDim], float64]: - # type of cells(E2C) is gtx.Field[gtx.Dims[CellDim, E2CDim], float64] +def sum_adjacent_cells(cells : gtx.Field[Dims[CellDim], float64]) -> gtx.Field[Dims[EdgeDim], float64]: + # type of cells(E2C) is gtx.Field[Dims[EdgeDim, E2CDim], float64] return neighbor_sum(cells(E2C), axis=E2CDim) @gtx.program -def run_sum_adjacent_cells(cells : gtx.Field[gtx.Dims[CellDim], float64], out : gtx.Field[gtx.Dims[EdgeDim], float64]): +def run_sum_adjacent_cells(cells : gtx.Field[Dims[CellDim], float64], out : gtx.Field[Dims[EdgeDim], float64]): sum_adjacent_cells(cells, out=out) run_sum_adjacent_cells(cell_values, edge_values, offset_provider={"E2C": E2C_offset_provider}) @@ -289,9 +287,9 @@ print("sum of adjacent cells: {}".format(edge_values.asnumpy())) For the border edges, the results are unchanged compared to the previous example, but the inner edges now contain the sum of the two adjacent cells: -| ![nearest_cell_values](connectivity_numbered_grid.svg) | $\mapsto$ | ![cell_values](connectivity_edge_cell_sum.svg) | -| :----------------------------------------------------: | :-------: | :--------------------------------------------: | -| _Domain (edges)_ | | _Edge values_ | +| ![nearest_cell_values](workshop/images/connectivity_numbered_grid.svg) | $\mapsto$ | ![cell_values](workshop/images/connectivity_edge_cell_sum.svg) | +| :--------------------------------------------------------------------: | :-------: | :------------------------------------------------------------: | +| _Domain (edges)_ | | _Edge values_ | +++ @@ -303,7 +301,7 @@ This function takes 3 input arguments: - mask: a field with dtype boolean - true branch: a tuple, a field, or a scalar - false branch: a tuple, a field, of a scalar - The mask can be directly a field of booleans (e.g. `gtx.Field[gtx.Dims[CellDim], bool]`) or an expression evaluating to this type (e.g. `gtx.Field[[CellDim], float64] > 3`). + The mask can be directly a field of booleans (e.g. `gtx.Field[Dims[CellDim], bool]`) or an expression evaluating to this type (e.g. `gtx.Field[[CellDim], float64] > 3`). The `where` builtin loops over each entry of the mask and returns values corresponding to the same indexes of either the true or the false branch. In the case where the true and false branches are either fields or scalars, the resulting output will be a field including all dimensions from all inputs. For example: @@ -313,8 +311,8 @@ result_where = gtx.as_field([CellDim, KDim], np.zeros(shape=grid_shape)) b = 6.0 @gtx.field_operator -def conditional(mask: gtx.Field[gtx.Dims[CellDim, KDim], bool], a: gtx.Field[gtx.Dims[CellDim, KDim], float64], b: float -) -> gtx.Field[gtx.Dims[CellDim, KDim], float64]: +def conditional(mask: gtx.Field[Dims[CellDim, KDim], bool], a: gtx.Field[Dims[CellDim, KDim], float64], b: float +) -> gtx.Field[Dims[CellDim, KDim], float64]: return where(mask, a, b) conditional(mask, a, b, out=result_where, offset_provider={}) @@ -330,13 +328,13 @@ result_1 = gtx.as_field([CellDim, KDim], np.zeros(shape=grid_shape)) result_2 = gtx.as_field([CellDim, KDim], np.zeros(shape=grid_shape)) @gtx.field_operator -def _conditional_tuple(mask: gtx.Field[gtx.Dims[CellDim, KDim], bool], a: gtx.Field[gtx.Dims[CellDim, KDim], float64], b: float -) -> tuple[gtx.Field[gtx.Dims[CellDim, KDim], float64], gtx.Field[gtx.Dims[CellDim, KDim], float64]]: +def _conditional_tuple(mask: gtx.Field[Dims[CellDim, KDim], bool], a: gtx.Field[Dims[CellDim, KDim], float64], b: float +) -> tuple[gtx.Field[Dims[CellDim, KDim], float64], gtx.Field[Dims[CellDim, KDim], float64]]: return where(mask, (a, b), (b, a)) @gtx.program -def conditional_tuple(mask: gtx.Field[gtx.Dims[CellDim, KDim], bool], a: gtx.Field[gtx.Dims[CellDim, KDim], float64], b: float, -result_1: gtx.Field[gtx.Dims[CellDim, KDim], float64], result_2: gtx.Field[gtx.Dims[CellDim, KDim], float64] +def conditional_tuple(mask: gtx.Field[Dims[CellDim, KDim], bool], a: gtx.Field[Dims[CellDim, KDim], float64], b: float, +result_1: gtx.Field[Dims[CellDim, KDim], float64], result_2: gtx.Field[Dims[CellDim, KDim], float64] ): _conditional_tuple(mask, a, b, out=(result_1, result_2)) @@ -361,17 +359,17 @@ result_2 = gtx.as_field([CellDim, KDim], np.zeros(shape=grid_shape)) @gtx.field_operator def _conditional_tuple_nested( - mask: gtx.Field[gtx.Dims[CellDim, KDim], bool], a: gtx.Field[gtx.Dims[CellDim, KDim], float64], b: gtx.Field[gtx.Dims[CellDim, KDim], float64], c: gtx.Field[gtx.Dims[CellDim, KDim], float64], d: gtx.Field[gtx.Dims[CellDim, KDim], float64] + mask: gtx.Field[Dims[CellDim, KDim], bool], a: gtx.Field[Dims[CellDim, KDim], float64], b: gtx.Field[Dims[CellDim, KDim], float64], c: gtx.Field[Dims[CellDim, KDim], float64], d: gtx.Field[Dims[CellDim, KDim], float64] ) -> tuple[ - tuple[gtx.Field[gtx.Dims[CellDim, KDim], float64], gtx.Field[gtx.Dims[CellDim, KDim], float64]], - tuple[gtx.Field[gtx.Dims[CellDim, KDim], float64], gtx.Field[gtx.Dims[CellDim, KDim], float64]], + tuple[gtx.Field[Dims[CellDim, KDim], float64], gtx.Field[Dims[CellDim, KDim], float64]], + tuple[gtx.Field[Dims[CellDim, KDim], float64], gtx.Field[Dims[CellDim, KDim], float64]], ]: return where(mask, ((a, b), (b, a)), ((c, d), (d, c))) @gtx.program def conditional_tuple_nested( - mask: gtx.Field[gtx.Dims[CellDim, KDim], bool], a: gtx.Field[gtx.Dims[CellDim, KDim], float64], b: gtx.Field[gtx.Dims[CellDim, KDim], float64], c: gtx.Field[gtx.Dims[CellDim, KDim], float64], d: gtx.Field[gtx.Dims[CellDim, KDim], float64], - result_1: gtx.Field[gtx.Dims[CellDim, KDim], float64], result_2: gtx.Field[gtx.Dims[CellDim, KDim], float64] + mask: gtx.Field[Dims[CellDim, KDim], bool], a: gtx.Field[Dims[CellDim, KDim], float64], b: gtx.Field[Dims[CellDim, KDim], float64], c: gtx.Field[Dims[CellDim, KDim], float64], d: gtx.Field[Dims[CellDim, KDim], float64], + result_1: gtx.Field[Dims[CellDim, KDim], float64], result_2: gtx.Field[Dims[CellDim, KDim], float64] ): _conditional_tuple_nested(mask, a, b, c, d, out=((result_1, result_2), (result_2, result_1))) @@ -426,19 +424,19 @@ The second lines first creates a temporary field using `edge_differences(C2E)`, ```{code-cell} ipython3 @gtx.field_operator -def pseudo_lap(cells : gtx.Field[gtx.Dims[CellDim], float64], - edge_weights : gtx.Field[gtx.Dims[CellDim, C2EDim], float64]) -> gtx.Field[gtx.Dims[CellDim], float64]: - edge_differences = cells(E2C[0]) - cells(E2C[1]) # type: gtx.Field[gtx.Dims[EdgeDim], float64] - return neighbor_sum(edge_differences(C2E) * edge_weights, axis=C2EDim) +def pseudo_lap(cells : gtx.Field[Dims[CellDim], float64], + edge_weights : gtx.Field[Dims[CellDim, C2EDim], float64]) -> gtx.Field[Dims[CellDim], float64]: + edges = cells(E2C[0]) # type: gtx.Field[Dims[EdgeDim], float64] + return neighbor_sum(edges(C2E) * edge_weights, axis=C2EDim) ``` The program itself is just a shallow wrapper over the `pseudo_lap` field operator. The significant part is how offset providers for both the edge-to-cell and cell-to-edge connectivities are supplied when the program is called: ```{code-cell} ipython3 @gtx.program -def run_pseudo_laplacian(cells : gtx.Field[gtx.Dims[CellDim], float64], - edge_weights : gtx.Field[gtx.Dims[CellDim, C2EDim], float64], - out : gtx.Field[gtx.Dims[CellDim], float64]): +def run_pseudo_laplacian(cells : gtx.Field[Dims[CellDim], float64], + edge_weights : gtx.Field[Dims[CellDim, C2EDim], float64], + out : gtx.Field[Dims[CellDim], float64]): pseudo_lap(cells, edge_weights, out=out) result_pseudo_lap = gtx.as_field([CellDim], np.zeros(shape=(6,))) @@ -455,7 +453,11 @@ As a closure, here is an example of chaining field operators, which is very simp ```{code-cell} ipython3 @gtx.field_operator -def pseudo_laplap(cells : gtx.Field[gtx.Dims[CellDim], float64], - edge_weights : gtx.Field[gtx.Dims[CellDim, C2EDim], float64]) -> gtx.Field[gtx.Dims[CellDim], float64]: +def pseudo_laplap(cells : gtx.Field[Dims[CellDim], float64], + edge_weights : gtx.Field[Dims[CellDim, C2EDim], float64]) -> gtx.Field[Dims[CellDim], float64]: return pseudo_lap(pseudo_lap(cells, edge_weights), edge_weights) ``` + +```{code-cell} ipython3 + +``` diff --git a/docs/user/next/workshop/exercises/1_simple_addition.ipynb b/docs/user/next/workshop/exercises/1_simple_addition.ipynb new file mode 100644 index 0000000000..918e72b084 --- /dev/null +++ b/docs/user/next/workshop/exercises/1_simple_addition.ipynb @@ -0,0 +1,139 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2ead8b70", + "metadata": {}, + "source": [ + "# 1. Simple Addition" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e7b501e0", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "markdown", + "id": "06113a1f", + "metadata": {}, + "source": [ + "Next we implement the stencil and a numpy reference version, in order to verify them against each other." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4c9f3427", + "metadata": {}, + "outputs": [], + "source": [ + "I = gtx.Dimension(\"I\")\n", + "J = gtx.Dimension(\"J\")\n", + "size = 10" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8272b87c", + "metadata": {}, + "outputs": [], + "source": [ + "def addition_numpy(a: np.array, b: np.array) -> np.array:\n", + " c = a + b\n", + " return c" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "05f2d09d", + "metadata": {}, + "outputs": [], + "source": [ + "def addition ... # TODO fix this cell" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db95ed00", + "metadata": {}, + "outputs": [], + "source": [ + "def test_addition(backend=None):\n", + " domain = gtx.domain({I: size, J: size})\n", + "\n", + " a_data = np.fromfunction(lambda xx, yy: xx, domain.shape, dtype=float)\n", + " a = gtx.as_field(domain, a_data, allocator=backend)\n", + " b_data = np.fromfunction(lambda xx, yy: yy, domain.shape, dtype=float)\n", + " b = gtx.as_field(domain, b_data, allocator=backend)\n", + "\n", + " c_numpy = addition_numpy(a.asnumpy(), b.asnumpy())\n", + "\n", + " c = gtx.zeros(domain, allocator=backend)\n", + "\n", + " addition(a, b, out=c, offset_provider={})\n", + "\n", + " assert np.allclose(c.asnumpy(), c_numpy)\n", + "\n", + " print(\"Result:\")\n", + " print(c)\n", + " print(c.asnumpy())\n", + " \n", + " # Plots\n", + " fig, ax = plt.subplot_mosaic([\n", + " ['a', 'b', 'c']\n", + " ])\n", + " ax['a'].imshow(a.asnumpy())\n", + " ax['b'].imshow(b.asnumpy())\n", + " ax['c'].imshow(c.asnumpy())\n", + "\n", + " print(\"\\nTest successful!\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "08502c34", + "metadata": {}, + "outputs": [], + "source": [ + "test_addition()" + ] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/1_simple_addition_solution.ipynb b/docs/user/next/workshop/exercises/1_simple_addition_solution.ipynb new file mode 100644 index 0000000000..59505f142a --- /dev/null +++ b/docs/user/next/workshop/exercises/1_simple_addition_solution.ipynb @@ -0,0 +1,209 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2ead8b70", + "metadata": {}, + "source": [ + "# 1. Simple Addition" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "e7b501e0", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import gt4py.next as gtx\n", + "from gt4py.next import Dims" + ] + }, + { + "cell_type": "markdown", + "id": "06113a1f", + "metadata": {}, + "source": [ + "Next we implement the stencil and a numpy reference version, in order to verify them against each other." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "4c9f3427", + "metadata": {}, + "outputs": [], + "source": [ + "I = gtx.Dimension(\"I\")\n", + "J = gtx.Dimension(\"J\")\n", + "size = 10" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "8272b87c", + "metadata": {}, + "outputs": [], + "source": [ + "def addition_numpy(a: np.array, b: np.array) -> np.array:\n", + " c = a + b\n", + " return c" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "05f2d09d", + "metadata": {}, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'Dims' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[4], line 3\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;129m@gtx\u001b[39m\u001b[38;5;241m.\u001b[39mfield_operator\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21maddition\u001b[39m(\n\u001b[0;32m----> 3\u001b[0m a: gtx\u001b[38;5;241m.\u001b[39mField[\u001b[43mDims\u001b[49m[I, J], \u001b[38;5;28mfloat\u001b[39m], b: gtx\u001b[38;5;241m.\u001b[39mField[Dims[I, J], \u001b[38;5;28mfloat\u001b[39m]\n\u001b[1;32m 4\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m gtx\u001b[38;5;241m.\u001b[39mField[Dims[I, J], \u001b[38;5;28mfloat\u001b[39m]:\n\u001b[1;32m 5\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m a \u001b[38;5;241m+\u001b[39m b\n", + "\u001b[0;31mNameError\u001b[0m: name 'Dims' is not defined" + ] + } + ], + "source": [ + "@gtx.field_operator\n", + "def addition(\n", + " a: gtx.Field[Dims[I, J], float], b: gtx.Field[Dims[I, J], float]\n", + ") -> gtx.Field[Dims[I, J], float]:\n", + " return a + b" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "db95ed00", + "metadata": {}, + "outputs": [], + "source": [ + "def test_addition(backend=None):\n", + " domain = gtx.domain({I: size, J: size})\n", + "\n", + " a_data = np.fromfunction(lambda xx, yy: xx, domain.shape, dtype=float)\n", + " a = gtx.as_field(domain, a_data, allocator=backend)\n", + " b_data = np.fromfunction(lambda xx, yy: yy, domain.shape, dtype=float)\n", + " b = gtx.as_field(domain, b_data, allocator=backend)\n", + "\n", + " c_numpy = addition_numpy(a.asnumpy(), b.asnumpy())\n", + "\n", + " c = gtx.zeros(domain, allocator=backend)\n", + "\n", + " addition(a, b, out=c, offset_provider={})\n", + "\n", + " assert np.allclose(c.asnumpy(), c_numpy)\n", + "\n", + " print(\"Result:\")\n", + " print(c)\n", + " print(c.asnumpy())\n", + " \n", + " # Plots\n", + " fig, ax = plt.subplot_mosaic([\n", + " ['a', 'b', 'c']\n", + " ])\n", + " ax['a'].imshow(a.asnumpy())\n", + " ax['b'].imshow(b.asnumpy())\n", + " ax['c'].imshow(c.asnumpy())\n", + "\n", + " print(\"\\nTest successful!\")" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "08502c34", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Result:\n", + "⟨Domain(I[horizontal]=(0:10), J[horizontal]=(0:10)) → DType(scalar_type=, tensor_shape=())⟩\n", + "[[ 0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]\n", + " [ 1. 2. 3. 4. 5. 6. 7. 8. 9. 10.]\n", + " [ 2. 3. 4. 5. 6. 7. 8. 9. 10. 11.]\n", + " [ 3. 4. 5. 6. 7. 8. 9. 10. 11. 12.]\n", + " [ 4. 5. 6. 7. 8. 9. 10. 11. 12. 13.]\n", + " [ 5. 6. 7. 8. 9. 10. 11. 12. 13. 14.]\n", + " [ 6. 7. 8. 9. 10. 11. 12. 13. 14. 15.]\n", + " [ 7. 8. 9. 10. 11. 12. 13. 14. 15. 16.]\n", + " [ 8. 9. 10. 11. 12. 13. 14. 15. 16. 17.]\n", + " [ 9. 10. 11. 12. 13. 14. 15. 16. 17. 18.]]\n", + "\n", + "Test successful!\n" + ] + }, + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAAhYAAAC9CAYAAADvAzTXAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuNSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/xnp5ZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAS1UlEQVR4nO3da2zU1b7G8aedtlMK06FcWqiM3LYEKTcFZCNqNDayOWD0xTFoakKqQaNFRIyRvgA0bKy84ZAoATcJwguu2QnRmCPGNEGCQLip0W3CRVGKWArG3WkLzJSZOS/c9KQR6PxXV+c/i34/ySR2Msu1nD6Fx39n5peTSqVSAgAAsCDX7wMAAIDbB8UCAABYQ7EAAADWUCwAAIA1FAsAAGANxQIAAFhDsQAAANbkZXrDZDKp8+fPKxQKKScnJ9Pb4zaQSqXU0tKi8vJy5eZmrhuTXdjgR37JLmxIN7sZLxbnz59XJBLJ9La4DTU0NGjYsGEZ24/swqZM5pfswqauspvxYhEKhSRJD+i/lKf8TG+P28A1tWu//rcjS5nS3ewG/jLSaN8rI0qM1klSS8TsR7xtqPn/1cbL48ZrS4c0G62bVnrWeM/K4n+ZrSsy+++MtiY1/N6fMprfbmd34ACzjUvCZuskJUqKjNbFw0HjPWP9A8Zr4yGznxnTdZLU3s9s3bXipNG65NWrOvfW37vMbsaLxfXLcHnKV14OxQIG/vMh9Jm+pNvd7AYCZn/g5eUXGq2TpECB2Y94oND8uc3tY355P9D3qtG6gn7mf5YUhcz+Miku6t6vMTKZ325nN7fAbGPDzEtSTp5Z7pP55nsm8s2LRaLA7PsZCJrnIGn4R0Oy0KxYXNdVdnnxJgAAsMaoWKxbt04jRoxQYWGhpk+frsOHD9s+F9AjyC5cRXbhCs/FYufOnVqyZIlWrFih48ePa9KkSZo1a5aampp64nyANWQXriK7cInnYrFmzRotWLBA1dXVGjdunDZs2KCioiJt2rSpJ84HWEN24SqyC5d4KhbxeFzHjh1TZWXl//8LcnNVWVmpgwcP3nBNLBZTNBrtdAMyjezCVWQXrvFULC5duqREIqGysrJO95eVlamxsfGGa+rq6hQOhztuvJcafiC7cBXZhWt6/F0htbW1am5u7rg1NDT09JaAFWQXriK78JOnN7kPGjRIgUBAFy5c6HT/hQsXNGTIkBuuCQaDCgbN31cM2EB24SqyC9d4umJRUFCgKVOmqL6+vuO+ZDKp+vp6zZgxw/rhAFvILlxFduEazx/Lt2TJEs2fP19Tp07Vfffdp7Vr16qtrU3V1dU9cT7AGrILV5FduMRzsZg3b54uXryo5cuXq7GxUZMnT9aePXv+9MIiINuQXbiK7MIlRoMEFi5cqIULF9o+C9DjyC5cRXbhiowPIbsu8JeRxkOZ0LulEjHptN+nANwTGDjAbKDYgP5G+yUG9DVaJ0nx/mZ/P3RnQmmsuBtTfQ3XtndjyG276ZTS0DWzdXnprWMIGQAAsIZiAQAArKFYAAAAaygWAADAGooFAACwhmIBAACsoVgAAABrKBYAAMAaigUAALCGYgEAAKyhWAAAAGsoFgAAwBqKBQAAsMa36aZXRpQoL7/Qr+3hsGvtV5luCpgoCUsGU6VNp5SaTiiVzKeU+jGhVDKfUmo6oVQyn1JaEIqb7RdIbx1XLAAAgDUUCwAAYA3FAgAAWOOpWNTV1WnatGkKhUIqLS3Vk08+qRMnTvTU2QBryC5cRXbhGk/F4osvvlBNTY0OHTqkzz//XO3t7XrsscfU1tbWU+cDrCC7cBXZhWs8vStkz549nb7evHmzSktLdezYMT300ENWDwbYRHbhKrIL13Tr7abNzc2SpAEDBtz0MbFYTLFYrOPraDTanS0BK8guXEV2ke2MX7yZTCa1ePFizZw5U+PHj7/p4+rq6hQOhztukUjEdEvACrILV5FduMC4WNTU1Oi7777Tjh07bvm42tpaNTc3d9waGhpMtwSsILtwFdmFC4x+FbJw4UJ98skn2rdvn4YNG3bLxwaDQQWD5p++BthEduEqsgtXeCoWqVRKr7zyinbv3q29e/dq5MiRPXUuwCqyC1eRXbjGU7GoqanRtm3b9NFHHykUCqmxsVGSFA6H1adPnx45IGAD2YWryC5c4+k1FuvXr1dzc7MefvhhDR06tOO2c+fOnjofYAXZhavILlzj+VchgIvILlxFduEa38amt0TyFCjwbXs4LBEnN4CJREmRcvIKPa8zHX9uOvpcMh9/7sfoc8l8/Lnp6HPJfPx5Seiy0bpEbqzrB4khZAAAwCKKBQAAsIZiAQAArKFYAAAAaygWAADAGooFAACwhmIBAACsoVgAAABrKBYAAMAaigUAALCGYgEAAKyhWAAAAGsoFgAAwBrfxkS2Dc1RoNB8Ch16r8RVcgOYiIeDSuZ7n1RqOqXUdEKpZD6l1I8JpZL5lFLTCaWS+ZTSIX1bjNa1K72zcsUCAABYQ7EAAADWUCwAAIA13SoW7777rnJycrR48WJLxwEyg+zCVWQX2c64WBw5ckQffPCBJk6caPM8QI8ju3AV2YULjIpFa2urqqqqtHHjRpWUlNg+E9BjyC5cRXbhCqNiUVNTozlz5qiysrLLx8ZiMUWj0U43wC9kF64iu3CF58+x2LFjh44fP64jR46k9fi6ujq9/fbbng8G2EZ24SqyC5d4umLR0NCgV199VVu3blVhYWFaa2pra9Xc3Nxxa2hoMDoo0B1kF64iu3CNpysWx44dU1NTk+69996O+xKJhPbt26f3339fsVhMgUDnT2gLBoMKBr1/0htgE9mFq8guXOOpWDz66KP69ttvO91XXV2tsWPH6s033/xTuIFsQXbhKrIL13gqFqFQSOPHj+90X9++fTVw4MA/3Q9kE7ILV5FduIZP3gQAANZ0e7rp3r17LRwDyDyyC1eRXWQz38amx8vjyu3DBRN4l7xiPmYY6M1i/QNK5Ht/TYbp+HPT0eeS+fhzP0afS+bjz01Hn0vm489H9PvNaF1c7Wk9jr/ZAQCANRQLAABgDcUCAABYQ7EAAADWUCwAAIA1FAsAAGANxQIAAFhDsQAAANZQLAAAgDUUCwAAYA3FAgAAWEOxAAAA1lAsAACANRQLAABgjW9j00uHNCvQ96pf28NhibaYzvl9CMBB8VCOEgXeR5mbjj83HX0umY8/92P0uWQ+/tx09LlkPv58dOFFo3VXr6X33HLFAgAAWEOxAAAA1nguFr/88oueffZZDRw4UH369NGECRN09OjRnjgbYBXZhavILlzi6TUWv//+u2bOnKlHHnlEn376qQYPHqxTp06ppKSkp84HWEF24SqyC9d4KharV69WJBLRhx9+2HHfyJEjrR8KsI3swlVkF67x9KuQjz/+WFOnTtVTTz2l0tJS3XPPPdq4ceMt18RiMUWj0U43INPILlxFduEaT8Xixx9/1Pr163XXXXfps88+00svvaRFixZpy5YtN11TV1encDjccYtEIt0+NOAV2YWryC5ck5NKpVLpPrigoEBTp07VgQMHOu5btGiRjhw5ooMHD95wTSwWUywW6/g6Go0qEono3n++pkDfYDeOjt4q0RbT8f/+HzU3N6u4uDitNTaz+7CeUF5OvudzB8aM9rxGkq6MGmC0TpKid5p9VE3bHWafWyBJsTvMPwugrPzfRuv+WvaT8Z5/C39rtq4o1vWDbiDaklTJmB/Tzq/N7E6oXqVAQaHnM/M5Fl3rFZ9j0XpNS+/7osvserpiMXToUI0bN67TfXfffbfOnj170zXBYFDFxcWdbkCmkV24iuzCNZ6KxcyZM3XixIlO9508eVLDhw+3eijANrILV5FduMZTsXjttdd06NAhvfPOOzp9+rS2bdumf/zjH6qpqemp8wFWkF24iuzCNZ6KxbRp07R7925t375d48eP18qVK7V27VpVVVX11PkAK8guXEV24RrPr+yaO3eu5s6d2xNnAXoU2YWryC5c4tt002mlZ1XQz/sr64F4a7uO+30IwEHxUI4CQe/v8DB9d4fpOzsk83d3+PHODsn83R2m7+yQzN/dcVew0Wjd5XgirccxhAwAAFhDsQAAANZQLAAAgDUUCwAAYA3FAgAAWEOxAAAA1lAsAACANRQLAABgDcUCAABYQ7EAAADWUCwAAIA1FAsAAGANxQIAAFhDsQAAANb4Nja9svhfKgoF/NoeDrucm9AHfh8CcFB7PylZaLDOcPy56ehzyXz8uR+jzyXz8eemo88l8/Hnd+WbnbU1P70ccMUCAABYQ7EAAADWUCwAAIA1nopFIpHQsmXLNHLkSPXp00ejR4/WypUrlUqleup8gBVkF64iu3CNpxdvrl69WuvXr9eWLVtUUVGho0ePqrq6WuFwWIsWLeqpMwLdRnbhKrIL13gqFgcOHNATTzyhOXPmSJJGjBih7du36/DhwzddE4vFFIvFOr6ORqOGRwXMkV24iuzCNZ5+FXL//fervr5eJ0+elCR988032r9/v2bPnn3TNXV1dQqHwx23SCTSvRMDBsguXEV24RpPVyyWLl2qaDSqsWPHKhAIKJFIaNWqVaqqqrrpmtraWi1ZsqTj62g0SsiRcWQXriK7cI2nYrFr1y5t3bpV27ZtU0VFhb7++mstXrxY5eXlmj9//g3XBINBBYNBK4cFTJFduIrswjWeisUbb7yhpUuX6umnn5YkTZgwQT///LPq6upuGnAgG5BduIrswjWeXmNx+fJl5eZ2XhIIBJRMmn3cK5ApZBeuIrtwjacrFo8//rhWrVqlO++8UxUVFfrqq6+0Zs0aPffccz11PsAKsgtXkV24xlOxeO+997Rs2TK9/PLLampqUnl5uV588UUtX768p84HWEF24SqyC9fkpDL88W3RaFThcFi/nxyl4hCfKA7voi1JlYz5Uc3NzSouLs7cvv/J7sN6Qnk5+Z7XB8aMNtr3yqgBRuskKXqn2QDjtjtyjPeM3WE2lVKSysr/bbTur2U/Ge/5t/C3ZuuKYl0/6Ab8yO/17N757t+VW+h9vKnplFLTCaWS+ZRSPyaUSuZTSk0nlErmU0pH5/czWpdudvmbHQAAWEOxAAAA1lAsAACANRQLAABgDcUCAABYQ7EAAADWUCwAAIA1FAsAAGANxQIAAFhDsQAAANZQLAAAgDUUCwAAYA3FAgAAWGM2+rAbrg9TjbYmM701bhPXs5Phwbwd+11Tu2SwdSphNg3zWvtVo3WSlIib/YgnrppPN01eMZ9omWgze47ire3Ge17OTRitiybM/gzzI7/X90peNctSMs9sumky0I0s5JploV3me8ZlnqOr18yeo8txs/xJUmu+YQZN16WZ3YyPTT937pwikUgmt8RtqqGhQcOGDcvYfmQXNmUyv2QXNnWV3YwXi2QyqfPnzysUCiknp/P/FUWjUUUiETU0NNxy1ntvxnP0R1tuaWlReXm5cnMz99s8sts9PEd/8CO/ZLd7eI7+kG52M/6rkNzc3C5benFxca/+5qWjtz9H4XA443uSXTt4jjKfX7JrB89RetnlxZsAAMAaigUAALAmq4pFMBjUihUrFAwG/T5K1uI5yk58X7rGc5Sd+L50jefIm4y/eBMAANy+suqKBQAAcBvFAgAAWEOxAAAA1lAsAACANRQLAABgTdYUi3Xr1mnEiBEqLCzU9OnTdfjwYb+PlFXeeust5eTkdLqNHTvW72NBZLcrZDe7kd+bI7tmsqJY7Ny5U0uWLNGKFSt0/PhxTZo0SbNmzVJTU5PfR8sqFRUV+vXXXztu+/fv9/tIvR7ZTQ/ZzU7kt2tk17usKBZr1qzRggULVF1drXHjxmnDhg0qKirSpk2b/D5aVsnLy9OQIUM6boMGDfL7SL0e2U0P2c1O5LdrZNc734tFPB7XsWPHVFlZ2XFfbm6uKisrdfDgQR9Pln1OnTql8vJyjRo1SlVVVTp79qzfR+rVyG76yG72Ib/pIbve+V4sLl26pEQiobKysk73l5WVqbGx0adTZZ/p06dr8+bN2rNnj9avX68zZ87owQcfVEtLi99H67XIbnrIbnYiv10ju2YyPjYdZmbPnt3xzxMnTtT06dM1fPhw7dq1S88//7yPJwNujezCVWTXjO9XLAYNGqRAIKALFy50uv/ChQsaMmSIT6fKfv3799eYMWN0+vRpv4/Sa5FdM2Q3O5Bf78huenwvFgUFBZoyZYrq6+s77ksmk6qvr9eMGTN8PFl2a21t1Q8//KChQ4f6fZRei+yaIbvZgfx6R3bTlMoCO3bsSAWDwdTmzZtT33//feqFF15I9e/fP9XY2Oj30bLG66+/ntq7d2/qzJkzqS+//DJVWVmZGjRoUKqpqcnvo/VqZLdrZDd7kd9bI7tmsuI1FvPmzdPFixe1fPlyNTY2avLkydqzZ8+fXlTUm507d07PPPOMfvvtNw0ePFgPPPCADh06pMGDB/t9tF6N7HaN7GYv8ntrZNdMTiqVSvl9CAAAcHvw/TUWAADg9kGxAAAA1lAsAACANRQLAABgDcUCAABYQ7EAAADWUCwAAIA1FAsAAGANxQIAAFhDsQAAANZQLAAAgDX/B2SI2ivYywPQAAAAAElFTkSuQmCC", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "test_addition()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "91115f3c-a294-4489-adec-1f4788ce71bd", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/2_divergence_exercise.ipynb b/docs/user/next/workshop/exercises/2_divergence_exercise.ipynb new file mode 100644 index 0000000000..86baf90901 --- /dev/null +++ b/docs/user/next/workshop/exercises/2_divergence_exercise.ipynb @@ -0,0 +1,258 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c841c53b", + "metadata": {}, + "source": [ + "# 3. Divergence" + ] + }, + { + "cell_type": "markdown", + "id": "bcac0f9b", + "metadata": {}, + "source": [ + "Next we will translate a divergence stencil. We approximate the divergence of a vector field $\\mathbf{v}$ at the middle point of a cell $\\mathbf{P}$ in the following way: We take the dot product of the normal velocity $\\mathbf{n}_e$ of each direct neighbor edge of $\\mathbf{P}$ with $\\mathbf{v}_e$ which is multipled with the edge length $L_e$. The contributions from all three edges of a cell are summed up and then divided by the area of the cell $A_P$. In the next pictures we can see a graphical representation of all of the quantities involved:\n", + "\n", + "![](../images/divergence_picture.png \"Divergence\")\n", + "\n", + "And the equation:\n", + "\n", + "![](../images/divergence_formula.png \"Divergence\")\n", + "\n", + "The orientation of the edge has to factor in, since we do not know, in general, if the normal of an edge is pointed inwards or outwards of any cell we are looking at. We cannot have only outwards pointing edge normals, because if we look at two neighboring cells, the normal of their shared edge has to point outwards for one of the cells, but inwards for the other.\n", + "\n", + "![](../images/edge_orientation.png \"Edge Orientation\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "4eba62c1", + "metadata": {}, + "outputs": [], + "source": [ + "from helpers import *" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "0cb870eb", + "metadata": {}, + "outputs": [], + "source": [ + "def divergence_numpy(\n", + " c2e: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " L: np.array,\n", + " A: np.array,\n", + " edge_orientation: np.array,\n", + ") -> np.array:\n", + " uv_div = (\n", + " np.sum(\n", + " (u[c2e] * nx[c2e] + v[c2e] * ny[c2e]) * L[c2e] * edge_orientation, axis=1\n", + " )\n", + " / A\n", + " )\n", + " return uv_div" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "8fc6416f", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def divergence(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " L: gtx.Field[Dims[E], float],\n", + " A: gtx.Field[Dims[C], float],\n", + " edge_orientation: gtx.Field[Dims[C, C2EDim], float],\n", + ") -> gtx.Field[Dims[C], float]:\n", + " uv_div = A\n", + "\n", + " return uv_div" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "5dbd2f62", + "metadata": {}, + "outputs": [], + "source": [ + "def test_divergence():\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + "\n", + " cell_domain = gtx.domain({C: n_cells})\n", + " edge_domain = gtx.domain({E: n_edges})\n", + "\n", + " u = random_field(edge_domain, allocator=backend)\n", + " v = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " L = random_field(edge_domain, allocator=backend)\n", + " A = random_field(cell_domain, allocator=backend)\n", + " edge_orientation = random_sign(\n", + " gtx.domain({C: n_cells, C2EDim: 3}), allocator=backend\n", + " )\n", + "\n", + " divergence_ref = divergence_numpy(\n", + " c2e_table,\n", + " u.asnumpy(),\n", + " v.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " L.asnumpy(),\n", + " A.asnumpy(),\n", + " edge_orientation.asnumpy(),\n", + " )\n", + "\n", + " c2e_connectivity = gtx.NeighborTableOffsetProvider(\n", + " c2e_table, C, E, 3, has_skip_values=False\n", + " )\n", + "\n", + " divergence_gt4py = gtx.zeros(cell_domain, allocator=backend)\n", + "\n", + " divergence(\n", + " u,\n", + " v,\n", + " nx,\n", + " ny,\n", + " L,\n", + " A,\n", + " edge_orientation,\n", + " out=divergence_gt4py,\n", + " offset_provider={C2E.value: c2e_connectivity},\n", + " )\n", + "\n", + " assert np.allclose(divergence_gt4py.asnumpy(), divergence_ref)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "bbcb9bf5", + "metadata": {}, + "outputs": [ + { + "ename": "AssertionError", + "evalue": "", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mAssertionError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[10], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_divergence\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[9], line 48\u001b[0m, in \u001b[0;36mtest_divergence\u001b[0;34m()\u001b[0m\n\u001b[1;32m 34\u001b[0m divergence_gt4py \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mzeros(cell_domain, allocator\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[1;32m 36\u001b[0m divergence(\n\u001b[1;32m 37\u001b[0m u,\n\u001b[1;32m 38\u001b[0m v,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 45\u001b[0m offset_provider\u001b[38;5;241m=\u001b[39m{C2E\u001b[38;5;241m.\u001b[39mvalue: c2e_connectivity},\n\u001b[1;32m 46\u001b[0m )\n\u001b[0;32m---> 48\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(divergence_gt4py\u001b[38;5;241m.\u001b[39masnumpy(), divergence_ref)\n", + "\u001b[0;31mAssertionError\u001b[0m: " + ] + } + ], + "source": [ + "test_divergence()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "markdown", + "id": "5cd78463", + "metadata": {}, + "source": [ + "## 3. Divergence in ICON" + ] + }, + { + "cell_type": "markdown", + "id": "b9a35e28", + "metadata": {}, + "source": [ + "In ICON we can find a divergence in diffusion which looks somewhat like this, but also quite a bit different:\n", + "\n", + "```fortran\n", + " DO jb = i_startblk,i_endblk\n", + "\n", + " CALL get_indices_c(p_patch, jb, i_startblk, i_endblk, &\n", + " i_startidx, i_endidx, rl_start, rl_end)\n", + " DO jk = 1, nlev\n", + " DO jc = i_startidx, i_endidx\n", + "\n", + " div(jc,jk) = p_nh_prog%vn(ieidx(jc,jb,1),jk,ieblk(jc,jb,1))*p_int%geofac_div(jc,1,jb) + &\n", + " p_nh_prog%vn(ieidx(jc,jb,2),jk,ieblk(jc,jb,2))*p_int%geofac_div(jc,2,jb) + &\n", + " p_nh_prog%vn(ieidx(jc,jb,3),jk,ieblk(jc,jb,3))*p_int%geofac_div(jc,3,jb)\n", + " ENDDO\n", + " ENDDO\n", + " ENDDO\n", + "```\n", + "\n", + "Two assumptions are necessary to derive the ICON version of the divergence starting from our version above:\n", + "* Assume that the velocity components $u$ is always orthogonal and the velocity component $v$ is always parallel to the edge, in ICON these are called $vn$ and $vt$ where the n stands for normal and the t for tangential.\n", + "* At ICON startup time merge all constants (such as cell area $A_P$ and edge length $L_e$) into one array of geometrical factors `p_int%geofac_div`, which are constant during time stepping:\n", + "\n", + "```fortran\n", + " DO jb = i_startblk, i_endblk\n", + "\n", + " CALL get_indices_c(ptr_patch, jb, i_startblk, i_endblk, &\n", + " & i_startidx, i_endidx, rl_start, rl_end)\n", + "\n", + " DO je = 1, ptr_patch%geometry_info%cell_type\n", + " DO jc = i_startidx, i_endidx\n", + "\n", + " ile = ptr_patch%cells%edge_idx(jc,jb,je)\n", + " ibe = ptr_patch%cells%edge_blk(jc,jb,je)\n", + "\n", + " ptr_int%geofac_div(jc,je,jb) = &\n", + " & ptr_patch%edges%primal_edge_length(ile,ibe) * &\n", + " & ptr_patch%cells%edge_orientation(jc,jb,je) / &\n", + " & ptr_patch%cells%area(jc,jb)\n", + "\n", + " ENDDO !cell loop\n", + " ENDDO\n", + "\n", + " END DO !block loop\n", + "\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "83dadf6c-8ade-49bb-b3ea-7dbe0d6bd5b3", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/2_divergence_exercise_solution.ipynb b/docs/user/next/workshop/exercises/2_divergence_exercise_solution.ipynb new file mode 100644 index 0000000000..eda22846eb --- /dev/null +++ b/docs/user/next/workshop/exercises/2_divergence_exercise_solution.ipynb @@ -0,0 +1,258 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c841c53b", + "metadata": {}, + "source": [ + "# 3. Divergence" + ] + }, + { + "cell_type": "markdown", + "id": "bcac0f9b", + "metadata": {}, + "source": [ + "Next we will translate a divergence stencil. We approximate the divergence of a vector field $\\mathbf{v}$ at the middle point of a cell $\\mathbf{P}$ in the following way: We take the dot product of the normal velocity $\\mathbf{n}_e$ of each direct neighbor edge of $\\mathbf{P}$ with $\\mathbf{v}_e$ which is multipled with the edge length $L_e$. The contributions from all three edges of a cell are summed up and then divided by the area of the cell $A_P$. In the next pictures we can see a graphical representation of all of the quantities involved:\n", + "\n", + "![](../images/divergence_picture.png \"Divergence\")\n", + "\n", + "And the equation:\n", + "\n", + "![](../images/divergence_formula.png \"Divergence\")\n", + "\n", + "The orientation of the edge has to factor in, since we do not know, in general, if the normal of an edge is pointed inwards or outwards of any cell we are looking at. We cannot have only outwards pointing edge normals, because if we look at two neighboring cells, the normal of their shared edge has to point outwards for one of the cells, but inwards for the other.\n", + "\n", + "![](../images/edge_orientation.png \"Edge Orientation\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "4eba62c1", + "metadata": {}, + "outputs": [], + "source": [ + "from helpers import *" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "0cb870eb", + "metadata": {}, + "outputs": [], + "source": [ + "def divergence_numpy(\n", + " c2e: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " L: np.array,\n", + " A: np.array,\n", + " edge_orientation: np.array,\n", + ") -> np.array:\n", + " uv_div = (\n", + " np.sum(\n", + " (u[c2e] * nx[c2e] + v[c2e] * ny[c2e]) * L[c2e] * edge_orientation, axis=1\n", + " )\n", + " / A\n", + " )\n", + " return uv_div" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "8fc6416f", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def divergence(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " L: gtx.Field[Dims[E], float],\n", + " A: gtx.Field[Dims[C], float],\n", + " edge_orientation: gtx.Field[Dims[C, C2EDim], float],\n", + ") -> gtx.Field[Dims[C], float]:\n", + " uv_div = (\n", + " neighbor_sum(\n", + " (u(C2E) * nx(C2E) + v(C2E) * ny(C2E)) * L(C2E) * edge_orientation,\n", + " axis=C2EDim,\n", + " )\n", + " / A\n", + " )\n", + " return uv_div" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "5dbd2f62", + "metadata": {}, + "outputs": [], + "source": [ + "def test_divergence():\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + "\n", + " cell_domain = gtx.domain({C: n_cells})\n", + " edge_domain = gtx.domain({E: n_edges})\n", + "\n", + " u = random_field(edge_domain, allocator=backend)\n", + " v = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " L = random_field(edge_domain, allocator=backend)\n", + " A = random_field(cell_domain, allocator=backend)\n", + " edge_orientation = random_sign(\n", + " gtx.domain({C: n_cells, C2EDim: 3}), allocator=backend\n", + " )\n", + "\n", + " divergence_ref = divergence_numpy(\n", + " c2e_table,\n", + " u.asnumpy(),\n", + " v.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " L.asnumpy(),\n", + " A.asnumpy(),\n", + " edge_orientation.asnumpy(),\n", + " )\n", + "\n", + " c2e_connectivity = gtx.NeighborTableOffsetProvider(\n", + " c2e_table, C, E, 3, has_skip_values=False\n", + " )\n", + "\n", + " divergence_gt4py = gtx.zeros(cell_domain, allocator=backend)\n", + "\n", + " divergence(\n", + " u,\n", + " v,\n", + " nx,\n", + " ny,\n", + " L,\n", + " A,\n", + " edge_orientation,\n", + " out=divergence_gt4py,\n", + " offset_provider={C2E.value: c2e_connectivity},\n", + " )\n", + "\n", + " assert np.allclose(divergence_gt4py.asnumpy(), divergence_ref)" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "bbcb9bf5", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Test successful\n" + ] + } + ], + "source": [ + "test_divergence()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "markdown", + "id": "5cd78463", + "metadata": {}, + "source": [ + "## 3. Divergence in ICON" + ] + }, + { + "cell_type": "markdown", + "id": "b9a35e28", + "metadata": {}, + "source": [ + "In ICON we can find a divergence in diffusion which looks somewhat like this, but also quite a bit different:\n", + "\n", + "```fortran\n", + " DO jb = i_startblk,i_endblk\n", + "\n", + " CALL get_indices_c(p_patch, jb, i_startblk, i_endblk, &\n", + " i_startidx, i_endidx, rl_start, rl_end)\n", + " DO jk = 1, nlev\n", + " DO jc = i_startidx, i_endidx\n", + "\n", + " div(jc,jk) = p_nh_prog%vn(ieidx(jc,jb,1),jk,ieblk(jc,jb,1))*p_int%geofac_div(jc,1,jb) + &\n", + " p_nh_prog%vn(ieidx(jc,jb,2),jk,ieblk(jc,jb,2))*p_int%geofac_div(jc,2,jb) + &\n", + " p_nh_prog%vn(ieidx(jc,jb,3),jk,ieblk(jc,jb,3))*p_int%geofac_div(jc,3,jb)\n", + " ENDDO\n", + " ENDDO\n", + " ENDDO\n", + "```\n", + "\n", + "Two assumptions are necessary to derive the ICON version of the divergence starting from our version above:\n", + "* Assume that the velocity components $u$ is always orthogonal and the velocity component $v$ is always parallel to the edge, in ICON these are called $vn$ and $vt$ where the n stands for normal and the t for tangential.\n", + "* At ICON startup time merge all constants (such as cell area $A_P$ and edge length $L_e$) into one array of geometrical factors `p_int%geofac_div`, which are constant during time stepping:\n", + "\n", + "```fortran\n", + " DO jb = i_startblk, i_endblk\n", + "\n", + " CALL get_indices_c(ptr_patch, jb, i_startblk, i_endblk, &\n", + " & i_startidx, i_endidx, rl_start, rl_end)\n", + "\n", + " DO je = 1, ptr_patch%geometry_info%cell_type\n", + " DO jc = i_startidx, i_endidx\n", + "\n", + " ile = ptr_patch%cells%edge_idx(jc,jb,je)\n", + " ibe = ptr_patch%cells%edge_blk(jc,jb,je)\n", + "\n", + " ptr_int%geofac_div(jc,je,jb) = &\n", + " & ptr_patch%edges%primal_edge_length(ile,ibe) * &\n", + " & ptr_patch%cells%edge_orientation(jc,jb,je) / &\n", + " & ptr_patch%cells%area(jc,jb)\n", + "\n", + " ENDDO !cell loop\n", + " ENDDO\n", + "\n", + " END DO !block loop\n", + "\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a5856c0-0b96-4fe5-87b1-a43dcd05991d", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/3_gradient_exercise.ipynb b/docs/user/next/workshop/exercises/3_gradient_exercise.ipynb new file mode 100644 index 0000000000..a0e5580eb6 --- /dev/null +++ b/docs/user/next/workshop/exercises/3_gradient_exercise.ipynb @@ -0,0 +1,192 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8cfd1e20", + "metadata": {}, + "source": [ + "# 4. Gradient" + ] + }, + { + "cell_type": "markdown", + "id": "72aa4c96", + "metadata": {}, + "source": [ + "Another example is the gradient defined at the center of a Cell $\\mathbf{P}$ of a scalar function $f$. We approximate this by taking the sum over the three edges and multiplying $f(e)$ with the edge normal $\\mathbf{n}_e$ and the edge length $L_e$ and dividing the resulting sum with the cell area $A_P$.\n", + "The result will be the two components of the gradient vector.\n", + "\n", + "![](../images/gradient_picture.png \"Divergence\")\n", + "\n", + "![](../images/gradient_formula.png \"Divergence\")" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "1bddcf59", + "metadata": {}, + "outputs": [], + "source": [ + "from helpers import *\n", + "\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "37b0bc43", + "metadata": {}, + "outputs": [], + "source": [ + "def gradient_numpy(\n", + " c2e: np.array,\n", + " f: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " L: np.array,\n", + " A: np.array,\n", + " edge_orientation: np.array,\n", + ") -> gtx.tuple[np.array, np.array]:\n", + " # edge_orientation = np.expand_dims(edge_orientation, axis=-1)\n", + " f_x = np.sum(f[c2e] * nx[c2e] * L[c2e] * edge_orientation, axis=1) / A\n", + " f_y = np.sum(f[c2e] * ny[c2e] * L[c2e] * edge_orientation, axis=1) / A\n", + " return f_x, f_y" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "afa80e49", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def gradient(\n", + " f: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " L: gtx.Field[Dims[E], float],\n", + " A: gtx.Field[Dims[C], float],\n", + " edge_orientation: gtx.Field[Dims[C, C2EDim], float],\n", + ") -> gtx.tuple[gtx.Field[Dims[C], float], gtx.Field[Dims[C], float]]:\n", + " # TODO: fix components of gradient\n", + " f_x = A\n", + " f_y = A\n", + " return f_x, f_y" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "84b02762", + "metadata": {}, + "outputs": [], + "source": [ + "def test_gradient():\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + "\n", + " cell_domain = gtx.domain({C: n_cells})\n", + " edge_domain = gtx.domain({E: n_edges})\n", + " \n", + " f = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " L = random_field(edge_domain, allocator=backend)\n", + " A = random_field(cell_domain, allocator=backend)\n", + " edge_orientation = random_sign(\n", + " gtx.domain({C: n_cells, C2EDim: 3}), allocator=backend\n", + " )\n", + "\n", + " gradient_numpy_x, gradient_numpy_y = gradient_numpy(\n", + " c2e_table,\n", + " f.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " L.asnumpy(),\n", + " A.asnumpy(),\n", + " edge_orientation.asnumpy(),\n", + " )\n", + "\n", + " c2e_connectivity = gtx.NeighborTableOffsetProvider(c2e_table, C, E, 3, has_skip_values=False)\n", + "\n", + " gradient_gt4py_x = gtx.zeros(cell_domain, allocator=backend) \n", + " gradient_gt4py_y = gtx.zeros(cell_domain, allocator=backend) \n", + "\n", + " gradient(\n", + " f,\n", + " nx,\n", + " ny,\n", + " L,\n", + " A,\n", + " edge_orientation,\n", + " out=(gradient_gt4py_x, gradient_gt4py_y),\n", + " offset_provider={C2E.value: c2e_connectivity},\n", + " )\n", + "\n", + " assert np.allclose(gradient_gt4py_x.asnumpy(), gradient_numpy_x)\n", + " assert np.allclose(gradient_gt4py_y.asnumpy(), gradient_numpy_y)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "9267fe99", + "metadata": {}, + "outputs": [ + { + "ename": "AssertionError", + "evalue": "", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mAssertionError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[12], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_gradient\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[11], line 44\u001b[0m, in \u001b[0;36mtest_gradient\u001b[0;34m()\u001b[0m\n\u001b[1;32m 31\u001b[0m gradient_gt4py_y \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mzeros(cell_domain, allocator\u001b[38;5;241m=\u001b[39mbackend) \n\u001b[1;32m 33\u001b[0m gradient(\n\u001b[1;32m 34\u001b[0m f,\n\u001b[1;32m 35\u001b[0m nx,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 41\u001b[0m offset_provider\u001b[38;5;241m=\u001b[39m{C2E\u001b[38;5;241m.\u001b[39mvalue: c2e_connectivity},\n\u001b[1;32m 42\u001b[0m )\n\u001b[0;32m---> 44\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(gradient_gt4py_x\u001b[38;5;241m.\u001b[39masnumpy(), gradient_numpy_x)\n\u001b[1;32m 45\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(gradient_gt4py_y\u001b[38;5;241m.\u001b[39masnumpy(), gradient_numpy_y)\n", + "\u001b[0;31mAssertionError\u001b[0m: " + ] + } + ], + "source": [ + "test_gradient()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "242c933f-66c7-4701-8b57-8c4fdf97adcd", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/3_gradient_exercise_solution.ipynb b/docs/user/next/workshop/exercises/3_gradient_exercise_solution.ipynb new file mode 100644 index 0000000000..64550d9b58 --- /dev/null +++ b/docs/user/next/workshop/exercises/3_gradient_exercise_solution.ipynb @@ -0,0 +1,200 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8cfd1e20", + "metadata": {}, + "source": [ + "# 4. Gradient" + ] + }, + { + "cell_type": "markdown", + "id": "72aa4c96", + "metadata": {}, + "source": [ + "Another example is the gradient defined at the center of a Cell $\\mathbf{P}$ of a scalar function $f$. We approximate this by taking the sum over the three edges and multiplying $f(e)$ with the edge normal $\\mathbf{n}_e$ and the edge length $L_e$ and dividing the resulting sum with the cell area $A_P$.\n", + "The result will be the two components of the gradient vector.\n", + "\n", + "![](../images/gradient_picture.png \"Divergence\")\n", + "\n", + "![](../images/gradient_formula.png \"Divergence\")" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "1bddcf59", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "37b0bc43", + "metadata": {}, + "outputs": [], + "source": [ + "def gradient_numpy(\n", + " c2e: np.array,\n", + " f: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " L: np.array,\n", + " A: np.array,\n", + " edge_orientation: np.array,\n", + ") -> gtx.tuple[np.array, np.array]:\n", + " # edge_orientation = np.expand_dims(edge_orientation, axis=-1)\n", + " f_x = np.sum(f[c2e] * nx[c2e] * L[c2e] * edge_orientation, axis=1) / A\n", + " f_y = np.sum(f[c2e] * ny[c2e] * L[c2e] * edge_orientation, axis=1) / A\n", + " return f_x, f_y" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "afa80e49", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def gradient(\n", + " f: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " L: gtx.Field[Dims[E], float],\n", + " A: gtx.Field[Dims[C], float],\n", + " edge_orientation: gtx.Field[Dims[C, C2EDim], float],\n", + ") -> gtx.tuple[gtx.Field[Dims[C], float], gtx.Field[Dims[C], float]]:\n", + " f_x = neighbor_sum(f(C2E) * nx(C2E) * L(C2E) * edge_orientation, axis=C2EDim) / A\n", + " f_y = neighbor_sum(f(C2E) * ny(C2E) * L(C2E) * edge_orientation, axis=C2EDim) / A\n", + " return f_x, f_y" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "84b02762", + "metadata": {}, + "outputs": [], + "source": [ + "def test_gradient():\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + "\n", + " cell_domain = gtx.domain({C: n_cells})\n", + " edge_domain = gtx.domain({E: n_edges})\n", + " \n", + " f = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " L = random_field(edge_domain, allocator=backend)\n", + " A = random_field(cell_domain, allocator=backend)\n", + " edge_orientation = random_sign(\n", + " gtx.domain({C: n_cells, C2EDim: 3}), allocator=backend\n", + " )\n", + "\n", + " gradient_numpy_x, gradient_numpy_y = gradient_numpy(\n", + " c2e_table,\n", + " f.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " L.asnumpy(),\n", + " A.asnumpy(),\n", + " edge_orientation.asnumpy(),\n", + " )\n", + "\n", + " c2e_connectivity = gtx.NeighborTableOffsetProvider(c2e_table, C, E, 3, has_skip_values=False)\n", + "\n", + " gradient_gt4py_x = gtx.zeros(cell_domain, allocator=backend) \n", + " gradient_gt4py_y = gtx.zeros(cell_domain, allocator=backend) \n", + "\n", + " gradient(\n", + " f,\n", + " nx,\n", + " ny,\n", + " L,\n", + " A,\n", + " edge_orientation,\n", + " out=(gradient_gt4py_x, gradient_gt4py_y),\n", + " offset_provider={C2E.value: c2e_connectivity},\n", + " )\n", + "\n", + " assert np.allclose(gradient_gt4py_x.asnumpy(), gradient_numpy_x)\n", + " assert np.allclose(gradient_gt4py_y.asnumpy(), gradient_numpy_y)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "9267fe99", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Test successful\n" + ] + } + ], + "source": [ + "test_gradient()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2092795a-a774-4129-b42d-0c3193961303", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/4_curl_exercise.ipynb b/docs/user/next/workshop/exercises/4_curl_exercise.ipynb new file mode 100644 index 0000000000..54806b4a3a --- /dev/null +++ b/docs/user/next/workshop/exercises/4_curl_exercise.ipynb @@ -0,0 +1,207 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9e5abeda", + "metadata": {}, + "source": [ + "# 5. Curl" + ] + }, + { + "cell_type": "markdown", + "id": "0bc751b1", + "metadata": {}, + "source": [ + "As the last example of the easier operations, we take a look at the curl of a vector field $\\mathbf{v}$ defined at a vertex $\\mathbf{N}$.\n", + "To approximate this, we once again iterate over all of the direct neighboring edges of the vertex in the center and for each edge take the dot product of the vector field $\\mathbf{v}_e$ with the edge normals $\\mathbf{n}_f$ and multiply that by the dual edge length $\\hat{L}_e$. The resulting neighbor sum is then divided by the dual area $\\hat{A}_N$, which is the area of the Voronoi cell around the Vertex $\\mathbf{N}$.\n", + "\n", + "\n", + "![](../images/curl_picture.png \"Divergence\")\n", + "\n", + "\n", + "![](../images/curl_formula.png \"Divergence\")" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "1c1af88b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "ce333ad2", + "metadata": {}, + "outputs": [], + "source": [ + "def curl_numpy(\n", + " v2e: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " dualL: np.array,\n", + " dualA: np.array,\n", + " edge_orientation: np.array,\n", + ") -> np.array:\n", + " uv_curl = (\n", + " np.sum(\n", + " (u[v2e] * nx[v2e] + v[v2e] * ny[v2e]) * dualL[v2e] * edge_orientation,\n", + " axis=1,\n", + " )\n", + " / dualA\n", + " )\n", + "\n", + " return uv_curl" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "0925bdb0", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def curl(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " dualL: gtx.Field[Dims[E], float],\n", + " dualA: gtx.Field[Dims[V], float],\n", + " edge_orientation: gtx.Field[Dims[V, V2EDim], float],\n", + ") -> gtx.Field[Dims[V], float]:\n", + " # TODO: fix curl\n", + " uv_curl = dualA\n", + "\n", + " return uv_curl" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "5b6ffc9e", + "metadata": {}, + "outputs": [], + "source": [ + "def test_curl():\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + "\n", + " edge_domain = gtx.domain({E: n_edges})\n", + " vertex_domain = gtx.domain({V: n_vertices})\n", + " \n", + " u = random_field(edge_domain, allocator=backend)\n", + " v = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " dualL = random_field(edge_domain, allocator=backend)\n", + " dualA = random_field(vertex_domain, allocator=backend)\n", + " edge_orientation = random_sign(\n", + " gtx.domain({V: n_vertices, V2EDim: 6}), allocator=backend\n", + " )\n", + "\n", + " divergence_ref = curl_numpy(\n", + " v2e_table,\n", + " u.asnumpy(),\n", + " v.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " dualL.asnumpy(),\n", + " dualA.asnumpy(),\n", + " edge_orientation.asnumpy(),\n", + " )\n", + "\n", + " v2e_connectivity = gtx.NeighborTableOffsetProvider(v2e_table, V, E, 6, has_skip_values=False)\n", + "\n", + " curl_gt4py = gtx.zeros(vertex_domain, allocator=backend) \n", + "\n", + " curl(\n", + " u, v, nx, ny, dualL, dualA, edge_orientation, out = curl_gt4py, offset_provider = {V2E.value: v2e_connectivity}\n", + " )\n", + " \n", + " assert np.allclose(curl_gt4py.asnumpy(), divergence_ref)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "ae651445", + "metadata": {}, + "outputs": [ + { + "ename": "AssertionError", + "evalue": "", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mAssertionError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_curl\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[4], line 38\u001b[0m, in \u001b[0;36mtest_curl\u001b[0;34m()\u001b[0m\n\u001b[1;32m 32\u001b[0m curl_gt4py \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mzeros(vertex_domain, allocator\u001b[38;5;241m=\u001b[39mbackend) \n\u001b[1;32m 34\u001b[0m curl(\n\u001b[1;32m 35\u001b[0m u, v, nx, ny, dualL, dualA, edge_orientation, out \u001b[38;5;241m=\u001b[39m curl_gt4py, offset_provider \u001b[38;5;241m=\u001b[39m {V2E\u001b[38;5;241m.\u001b[39mvalue: v2e_connectivity}\n\u001b[1;32m 36\u001b[0m )\n\u001b[0;32m---> 38\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(curl_gt4py\u001b[38;5;241m.\u001b[39masnumpy(), divergence_ref)\n", + "\u001b[0;31mAssertionError\u001b[0m: " + ] + } + ], + "source": [ + "test_curl()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8c57e9f9-2ee7-47c2-ae74-d863bae71aba", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/4_curl_exercise_solution.ipynb b/docs/user/next/workshop/exercises/4_curl_exercise_solution.ipynb new file mode 100644 index 0000000000..2649c5e2cd --- /dev/null +++ b/docs/user/next/workshop/exercises/4_curl_exercise_solution.ipynb @@ -0,0 +1,207 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9e5abeda", + "metadata": {}, + "source": [ + "# 5. Curl" + ] + }, + { + "cell_type": "markdown", + "id": "0bc751b1", + "metadata": {}, + "source": [ + "As the last example of the easier operations, we take a look at the curl of a vector field $\\mathbf{v}$ defined at a vertex $\\mathbf{N}$.\n", + "To approximate this, we once again iterate over all of the direct neighboring edges of the vertex in the center and for each edge take the dot product of the vector field $\\mathbf{v}_e$ with the edge normals $\\mathbf{n}_f$ and multiply that by the dual edge length $\\hat{L}_e$. The resulting neighbor sum is then divided by the dual area $\\hat{A}_N$, which is the area of the Voronoi cell around the Vertex $\\mathbf{N}$.\n", + "\n", + "\n", + "![](../images/curl_picture.png \"Divergence\")\n", + "\n", + "\n", + "![](../images/curl_formula.png \"Divergence\")" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "1c1af88b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "ce333ad2", + "metadata": {}, + "outputs": [], + "source": [ + "def curl_numpy(\n", + " v2e: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " dualL: np.array,\n", + " dualA: np.array,\n", + " edge_orientation: np.array,\n", + ") -> np.array:\n", + " uv_curl = (\n", + " np.sum(\n", + " (u[v2e] * nx[v2e] + v[v2e] * ny[v2e]) * dualL[v2e] * edge_orientation,\n", + " axis=1,\n", + " )\n", + " / dualA\n", + " )\n", + "\n", + " return uv_curl" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "0925bdb0", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def curl(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " dualL: gtx.Field[Dims[E], float],\n", + " dualA: gtx.Field[Dims[V], float],\n", + " edge_orientation: gtx.Field[Dims[V, V2EDim], float],\n", + ") -> gtx.Field[Dims[V], float]:\n", + " uv_curl = (\n", + " neighbor_sum(\n", + " (u(V2E) * nx(V2E) + v(V2E) * ny(V2E)) * dualL(V2E) * edge_orientation,\n", + " axis=V2EDim,\n", + " )\n", + " / dualA\n", + " )\n", + "\n", + " return uv_curl" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "5b6ffc9e", + "metadata": {}, + "outputs": [], + "source": [ + "def test_curl():\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + "\n", + " edge_domain = gtx.domain({E: n_edges})\n", + " vertex_domain = gtx.domain({V: n_vertices})\n", + " \n", + " u = random_field(edge_domain, allocator=backend)\n", + " v = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " dualL = random_field(edge_domain, allocator=backend)\n", + " dualA = random_field(vertex_domain, allocator=backend)\n", + " edge_orientation = random_sign(\n", + " gtx.domain({V: n_vertices, V2EDim: 6}), allocator=backend\n", + " )\n", + "\n", + " divergence_ref = curl_numpy(\n", + " v2e_table,\n", + " u.asnumpy(),\n", + " v.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " dualL.asnumpy(),\n", + " dualA.asnumpy(),\n", + " edge_orientation.asnumpy(),\n", + " )\n", + "\n", + " v2e_connectivity = gtx.NeighborTableOffsetProvider(v2e_table, V, E, 6, has_skip_values=False)\n", + "\n", + " curl_gt4py = gtx.zeros(vertex_domain, allocator=backend) \n", + "\n", + " curl(\n", + " u, v, nx, ny, dualL, dualA, edge_orientation, out = curl_gt4py, offset_provider = {V2E.value: v2e_connectivity}\n", + " )\n", + " \n", + " assert np.allclose(curl_gt4py.asnumpy(), divergence_ref)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "ae651445", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Test successful\n" + ] + } + ], + "source": [ + "test_curl()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2dc0dc6d-438b-480b-8174-f3600b600389", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/5_vector_laplace_exercise.ipynb b/docs/user/next/workshop/exercises/5_vector_laplace_exercise.ipynb new file mode 100644 index 0000000000..b976b214c3 --- /dev/null +++ b/docs/user/next/workshop/exercises/5_vector_laplace_exercise.ipynb @@ -0,0 +1,374 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "60e37e45-40d6-46eb-9a27-758bf45329a7", + "metadata": {}, + "source": [ + "# 5. Vector Laplacian" + ] + }, + { + "cell_type": "markdown", + "id": "92405e72-4aa9-4e33-8b18-46f592666f27", + "metadata": {}, + "source": [ + "Starting off from laplacian formula by taking a scalar function (scalar field) and returning a scalar **p** that measures by how much the average value of f over a surface centered at p deviates from f(p).\n", + "\n", + "![](../images/laplacian.png \"Laplacian\")\n", + "\n", + "Look at vector Laplacian formula (relevant for Navier stokes viscous stress term) by taking a vector function (vector field) and returning a vector.\n", + "\n", + "![](../images/vector_laplacian.png \"Vector Laplacian\")\n", + "\n", + "Compute normal component of vector laplacian on finite volume meshes.\n", + "\n", + "![](../images/vector_laplacian_normal_component.png \"Normal Component of Vector Laplacian\")\n", + "\n", + "Can reuse divergence and curl as defined in previous exercises, however need two more directional gradients:\n", + "\n", + "![](../images/directional_gradient_n_picture.png \"Gradient in n direction\")\n", + "\n", + "![](../images/directional_gradient_n_formula.png \"Gradient in n direction\")\n", + "\n", + "![](../images/directional_gradient_tau_picture.png \"Gradient in tau direction\")\n", + "\n", + "![](../images/directional_gradient_tau_formula.png \"Gradient in tau direction\")" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "0a8f30c8-5f66-4ca8-8509-c6ff63234703", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "66b7a45a-4ad9-48b3-9474-cf460069fce8", + "metadata": {}, + "outputs": [], + "source": [ + "def divergence_numpy(\n", + " c2e: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " L: np.array,\n", + " A: np.array,\n", + " edge_orientation: np.array,\n", + ") -> np.array:\n", + " uv_div = (\n", + " np.sum(\n", + " (u[c2e] * nx[c2e] + v[c2e] * ny[c2e]) * L[c2e] * edge_orientation, axis=1\n", + " )\n", + " / A\n", + " )\n", + " return uv_div" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "0e9b6702-d7b1-49d9-8366-29a254d79a79", + "metadata": {}, + "outputs": [], + "source": [ + "def curl_numpy(\n", + " v2e: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " dualL: np.array,\n", + " dualA: np.array,\n", + " edge_orientation: np.array,\n", + ") -> np.array:\n", + " uv_curl = (\n", + " np.sum(\n", + " (u[v2e] * nx[v2e] + v[v2e] * ny[v2e]) * dualL[v2e] * edge_orientation,\n", + " axis=1,\n", + " )\n", + " / dualA\n", + " )\n", + "\n", + " return uv_curl" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "09d7d17c", + "metadata": {}, + "outputs": [], + "source": [ + "def laplacian_numpy(\n", + " c2e: np.array,\n", + " v2e: np.array,\n", + " e2v: np.array,\n", + " e2c: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " L: np.array,\n", + " dualL: np.array,\n", + " tangent_orientation: np.array,\n", + " A: np.array,\n", + " dualA: np.array,\n", + " edge_orientation_vertex: np.array,\n", + " edge_orientation_cell: np.array,\n", + ") -> np.array:\n", + " # compute curl (on vertices)\n", + " uv_curl = curl_numpy(v2e, u, v, nx, ny, dualL, dualA, edge_orientation_vertex)\n", + "\n", + " # compute divergence (on cells)\n", + " uv_div = divergence_numpy(c2e, u, v, nx, ny, L, A, edge_orientation_cell)\n", + " \n", + " # first term of of nabla2 (gradient of curl)\n", + " grad_of_curl = (uv_curl[e2v[:, 1]] - uv_curl[e2v[:, 0]])*tangent_orientation/L\n", + "\n", + " # second term of of nabla2 (gradient of divergence)\n", + " grad_of_div = (uv_div[e2c[:, 1]] - uv_div[e2c[:, 0]])/dualL \n", + "\n", + " # finalize nabla2 (difference between the two gradients)\n", + " uv_nabla2 = grad_of_div - grad_of_curl\n", + "\n", + " return uv_nabla2" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "26df30a3-7bc5-4d60-959c-b4f40cc7fc0f", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def divergence(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " L: gtx.Field[Dims[E], float],\n", + " A: gtx.Field[Dims[C], float],\n", + " edge_orientation: gtx.Field[Dims[C, C2EDim], float],\n", + ") -> gtx.Field[Dims[C], float]:\n", + " # compute divergence\n", + " uv_div = A\n", + " \n", + " return uv_div" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "a7413223-c9f7-4665-8995-e615ab48096b", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def curl(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " dualL: gtx.Field[Dims[E], float],\n", + " dualA: gtx.Field[Dims[V], float],\n", + " edge_orientation: gtx.Field[Dims[V, V2EDim], float],\n", + ") -> gtx.Field[Dims[V], float]:\n", + " # compute curl\n", + " uv_curl = dualA\n", + "\n", + " return uv_curl" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "f36ace26", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def laplacian_fvm(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " L: gtx.Field[Dims[E], float],\n", + " dualL: gtx.Field[Dims[E], float],\n", + " tangent_orientation: gtx.Field[Dims[E], float],\n", + " A: gtx.Field[Dims[C], float],\n", + " dualA: gtx.Field[Dims[V], float],\n", + " edge_orientation_vertex: gtx.Field[Dims[V, V2EDim], float],\n", + " edge_orientation_cell: gtx.Field[Dims[C, C2EDim], float],\n", + ") -> gtx.Field[Dims[E], float]:\n", + " # compute laplacian_fvm\n", + " uv_nabla2 = L\n", + "\n", + " return uv_nabla2" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "f9cfc097", + "metadata": {}, + "outputs": [], + "source": [ + "def test_laplacian():\n", + "\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + "\n", + " edge_domain = gtx.domain({E: n_edges})\n", + " vertex_domain = gtx.domain({V: n_vertices})\n", + " cell_domain = gtx.domain({C: n_cells})\n", + "\n", + "\n", + " u = random_field(edge_domain, allocator=backend)\n", + " v = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " L = random_field(edge_domain, allocator=backend)\n", + " dualL = random_field(edge_domain, allocator=backend)\n", + " tangent_orientation = random_field(edge_domain, allocator=backend)\n", + " A = random_field(cell_domain, allocator=backend)\n", + " dualA = random_field(vertex_domain, allocator=backend)\n", + " edge_orientation_vertex = random_sign(\n", + " gtx.domain({V: n_vertices, V2EDim: 6}), allocator=backend\n", + " )\n", + " edge_orientation_cell = random_sign(\n", + " gtx.domain({C: n_cells, C2EDim: 3}), allocator=backend\n", + " )\n", + "\n", + " laplacian_ref = laplacian_numpy(\n", + " c2e_table,\n", + " v2e_table,\n", + " e2v_table,\n", + " e2c_table,\n", + " u.asnumpy(),\n", + " v.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " L.asnumpy(),\n", + " dualL.asnumpy(),\n", + " tangent_orientation.asnumpy(),\n", + " A.asnumpy(),\n", + " dualA.asnumpy(),\n", + " edge_orientation_vertex.asnumpy(),\n", + " edge_orientation_cell.asnumpy(),\n", + " )\n", + "\n", + " c2e_connectivity = gtx.NeighborTableOffsetProvider(c2e_table, C, E, 3, has_skip_values=False)\n", + " v2e_connectivity = gtx.NeighborTableOffsetProvider(v2e_table, V, E, 6, has_skip_values=False)\n", + " e2v_connectivity = gtx.NeighborTableOffsetProvider(e2v_table, E, V, 2, has_skip_values=False)\n", + " e2c_connectivity = gtx.NeighborTableOffsetProvider(e2c_table, E, C, 2, has_skip_values=False)\n", + "\n", + "\n", + " laplacian_gt4py = gtx.zeros(edge_domain, allocator=backend)\n", + "\n", + " laplacian_fvm(\n", + " u,\n", + " v,\n", + " nx,\n", + " ny,\n", + " L,\n", + " dualL,\n", + " tangent_orientation,\n", + " A,\n", + " dualA,\n", + " edge_orientation_vertex,\n", + " edge_orientation_cell,\n", + " out = laplacian_gt4py,\n", + " offset_provider = {C2E.value: c2e_connectivity,\n", + " V2E.value: v2e_connectivity,\n", + " E2V.value: e2v_connectivity,\n", + " E2C.value: e2c_connectivity,\n", + " },\n", + " )\n", + " \n", + " assert np.allclose(laplacian_gt4py.asnumpy(), laplacian_ref)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "d4079375", + "metadata": {}, + "outputs": [ + { + "ename": "AssertionError", + "evalue": "", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mAssertionError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[9], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_laplacian\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[8], line 74\u001b[0m, in \u001b[0;36mtest_laplacian\u001b[0;34m()\u001b[0m\n\u001b[1;32m 52\u001b[0m laplacian_gt4py \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mzeros(edge_domain, allocator\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[1;32m 54\u001b[0m laplacian_fvm(\n\u001b[1;32m 55\u001b[0m u,\n\u001b[1;32m 56\u001b[0m v,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 71\u001b[0m },\n\u001b[1;32m 72\u001b[0m )\n\u001b[0;32m---> 74\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(laplacian_gt4py\u001b[38;5;241m.\u001b[39masnumpy(), laplacian_ref)\n", + "\u001b[0;31mAssertionError\u001b[0m: " + ] + } + ], + "source": [ + "test_laplacian()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c2a34287-274b-4db1-8dbc-d4ef24f40eed", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/5_vector_laplace_exercise_solution.ipynb b/docs/user/next/workshop/exercises/5_vector_laplace_exercise_solution.ipynb new file mode 100644 index 0000000000..927d56d639 --- /dev/null +++ b/docs/user/next/workshop/exercises/5_vector_laplace_exercise_solution.ipynb @@ -0,0 +1,391 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "60e37e45-40d6-46eb-9a27-758bf45329a7", + "metadata": {}, + "source": [ + "# 5. Vector Laplacian" + ] + }, + { + "cell_type": "markdown", + "id": "92405e72-4aa9-4e33-8b18-46f592666f27", + "metadata": {}, + "source": [ + "Starting off from laplacian formula\n", + "\n", + "![](../images/laplacian.png \"Laplacian\")\n", + "\n", + "Look at vector Laplacian formula (relevant for Navier stokes viscous stress term)\n", + "\n", + "![](../images/vector_laplacian.png \"Vector Laplacian\")\n", + "\n", + "Compute normal component of vector laplacian on finite volume meshes\n", + "\n", + "![](../images/vector_laplacian_normal_component.png \"Normal Component of Vector Laplacian\")\n", + "\n", + "Can reuse divergence and curl as defined in previous exercises, however need two more directional gradients:\n", + "\n", + "![](../images/directional_gradient_n_picture.png \"Gradient in n direction\")\n", + "\n", + "![](../images/directional_gradient_n_formula.png \"Gradient in n direction\")\n", + "\n", + "![](../images/directional_gradient_tau_picture.png \"Gradient in tau direction\")\n", + "\n", + "![](../images/directional_gradient_tau_formula.png \"Gradient in tau direction\")" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "0a8f30c8-5f66-4ca8-8509-c6ff63234703", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "66b7a45a-4ad9-48b3-9474-cf460069fce8", + "metadata": {}, + "outputs": [], + "source": [ + "def divergence_numpy(\n", + " c2e: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " L: np.array,\n", + " A: np.array,\n", + " edge_orientation: np.array,\n", + ") -> np.array:\n", + " uv_div = (\n", + " np.sum(\n", + " (u[c2e] * nx[c2e] + v[c2e] * ny[c2e]) * L[c2e] * edge_orientation, axis=1\n", + " )\n", + " / A\n", + " )\n", + " return uv_div" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "0e9b6702-d7b1-49d9-8366-29a254d79a79", + "metadata": {}, + "outputs": [], + "source": [ + "def curl_numpy(\n", + " v2e: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " dualL: np.array,\n", + " dualA: np.array,\n", + " edge_orientation: np.array,\n", + ") -> np.array:\n", + " uv_curl = (\n", + " np.sum(\n", + " (u[v2e] * nx[v2e] + v[v2e] * ny[v2e]) * dualL[v2e] * edge_orientation,\n", + " axis=1,\n", + " )\n", + " / dualA\n", + " )\n", + "\n", + " return uv_curl" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "09d7d17c", + "metadata": {}, + "outputs": [], + "source": [ + "def laplacian_numpy(\n", + " c2e: np.array,\n", + " v2e: np.array,\n", + " e2v: np.array,\n", + " e2c: np.array,\n", + " u: np.array,\n", + " v: np.array,\n", + " nx: np.array,\n", + " ny: np.array,\n", + " L: np.array,\n", + " dualL: np.array,\n", + " tangent_orientation: np.array,\n", + " A: np.array,\n", + " dualA: np.array,\n", + " edge_orientation_vertex: np.array,\n", + " edge_orientation_cell: np.array,\n", + ") -> np.array:\n", + " # compute curl (on vertices)\n", + " uv_curl = curl_numpy(v2e, u, v, nx, ny, dualL, dualA, edge_orientation_vertex)\n", + "\n", + " # compute divergence (on cells)\n", + " uv_div = divergence_numpy(c2e, u, v, nx, ny, L, A, edge_orientation_cell)\n", + " \n", + " # first term of of nabla2 (gradient of curl)\n", + " grad_of_curl = (uv_curl[e2v[:, 1]] - uv_curl[e2v[:, 0]])*tangent_orientation/L\n", + "\n", + " # second term of of nabla2 (gradient of divergence)\n", + " grad_of_div = (uv_div[e2c[:, 1]] - uv_div[e2c[:, 0]])/dualL \n", + "\n", + " # finalize nabla2 (difference between the two gradients)\n", + " uv_nabla2 = grad_of_div - grad_of_curl\n", + "\n", + " return uv_nabla2" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "26df30a3-7bc5-4d60-959c-b4f40cc7fc0f", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def divergence(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " L: gtx.Field[Dims[E], float],\n", + " A: gtx.Field[Dims[C], float],\n", + " edge_orientation: gtx.Field[Dims[C, C2EDim], float],\n", + ") -> gtx.Field[Dims[C], float]:\n", + " uv_div = (\n", + " neighbor_sum(\n", + " (u(C2E) * nx(C2E) + v(C2E) * ny(C2E)) * L(C2E) * edge_orientation,\n", + " axis=C2EDim,\n", + " )\n", + " / A\n", + " )\n", + " return uv_div" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "a7413223-c9f7-4665-8995-e615ab48096b", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def curl(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " dualL: gtx.Field[Dims[E], float],\n", + " dualA: gtx.Field[Dims[V], float],\n", + " edge_orientation: gtx.Field[Dims[V, V2EDim], float],\n", + ") -> gtx.Field[Dims[V], float]:\n", + " uv_curl = (\n", + " neighbor_sum(\n", + " (u(V2E) * nx(V2E) + v(V2E) * ny(V2E)) * dualL(V2E) * edge_orientation,\n", + " axis=V2EDim,\n", + " )\n", + " / dualA\n", + " )\n", + "\n", + " return uv_curl" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "f36ace26", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def laplacian_fvm(\n", + " u: gtx.Field[Dims[E], float],\n", + " v: gtx.Field[Dims[E], float],\n", + " nx: gtx.Field[Dims[E], float],\n", + " ny: gtx.Field[Dims[E], float],\n", + " L: gtx.Field[Dims[E], float],\n", + " dualL: gtx.Field[Dims[E], float],\n", + " tangent_orientation: gtx.Field[Dims[E], float],\n", + " A: gtx.Field[Dims[C], float],\n", + " dualA: gtx.Field[Dims[V], float],\n", + " edge_orientation_vertex: gtx.Field[Dims[V, V2EDim], float],\n", + " edge_orientation_cell: gtx.Field[Dims[C, C2EDim], float],\n", + ") -> gtx.Field[Dims[E], float]:\n", + " \n", + " # compute curl (on vertices)\n", + " uv_curl = curl(u, v, nx, ny, dualL, dualA, edge_orientation_vertex)\n", + "\n", + " # compute divergence (on cells)\n", + " uv_div = divergence(u, v, nx, ny, L, A, edge_orientation_cell)\n", + " \n", + " # first term of of nabla2 (gradient of curl)\n", + " grad_of_curl = (uv_curl(E2V[1]) - uv_curl(E2V[0]))*tangent_orientation/L\n", + "\n", + " # second term of of nabla2 (gradient of divergence)\n", + " grad_of_div = (uv_div(E2C[1]) - uv_div(E2C[0]))/dualL \n", + "\n", + " # finalize nabla2 (difference between the two gradients)\n", + " uv_nabla2 = grad_of_div - grad_of_curl\n", + "\n", + " return uv_nabla2" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "f9cfc097", + "metadata": {}, + "outputs": [], + "source": [ + "def test_laplacian():\n", + "\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + "\n", + " edge_domain = gtx.domain({E: n_edges})\n", + " vertex_domain = gtx.domain({V: n_vertices})\n", + " cell_domain = gtx.domain({C: n_cells})\n", + "\n", + "\n", + " u = random_field(edge_domain, allocator=backend)\n", + " v = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " L = random_field(edge_domain, allocator=backend)\n", + " dualL = random_field(edge_domain, allocator=backend)\n", + " tangent_orientation = random_field(edge_domain, allocator=backend)\n", + " A = random_field(cell_domain, allocator=backend)\n", + " dualA = random_field(vertex_domain, allocator=backend)\n", + " edge_orientation_vertex = random_sign(\n", + " gtx.domain({V: n_vertices, V2EDim: 6}), allocator=backend\n", + " )\n", + " edge_orientation_cell = random_sign(\n", + " gtx.domain({C: n_cells, C2EDim: 3}), allocator=backend\n", + " )\n", + "\n", + " laplacian_ref = laplacian_numpy(\n", + " c2e_table,\n", + " v2e_table,\n", + " e2v_table,\n", + " e2c_table,\n", + " u.asnumpy(),\n", + " v.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " L.asnumpy(),\n", + " dualL.asnumpy(),\n", + " tangent_orientation.asnumpy(),\n", + " A.asnumpy(),\n", + " dualA.asnumpy(),\n", + " edge_orientation_vertex.asnumpy(),\n", + " edge_orientation_cell.asnumpy(),\n", + " )\n", + "\n", + " c2e_connectivity = gtx.NeighborTableOffsetProvider(c2e_table, C, E, 3, has_skip_values=False)\n", + " v2e_connectivity = gtx.NeighborTableOffsetProvider(v2e_table, V, E, 6, has_skip_values=False)\n", + " e2v_connectivity = gtx.NeighborTableOffsetProvider(e2v_table, E, V, 2, has_skip_values=False)\n", + " e2c_connectivity = gtx.NeighborTableOffsetProvider(e2c_table, E, C, 2, has_skip_values=False)\n", + "\n", + "\n", + " laplacian_gt4py = gtx.zeros(edge_domain, allocator=backend)\n", + "\n", + " laplacian_fvm(\n", + " u,\n", + " v,\n", + " nx,\n", + " ny,\n", + " L,\n", + " dualL,\n", + " tangent_orientation,\n", + " A,\n", + " dualA,\n", + " edge_orientation_vertex,\n", + " edge_orientation_cell,\n", + " out = laplacian_gt4py,\n", + " offset_provider = {C2E.value: c2e_connectivity,\n", + " V2E.value: v2e_connectivity,\n", + " E2V.value: e2v_connectivity,\n", + " E2C.value: e2c_connectivity,\n", + " },\n", + " )\n", + " \n", + " assert np.allclose(laplacian_gt4py.asnumpy(), laplacian_ref)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "d4079375", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Test successful\n" + ] + } + ], + "source": [ + "test_laplacian()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c2a34287-274b-4db1-8dbc-d4ef24f40eed", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/6_where_domain.ipynb b/docs/user/next/workshop/exercises/6_where_domain.ipynb new file mode 100644 index 0000000000..c23e8c121a --- /dev/null +++ b/docs/user/next/workshop/exercises/6_where_domain.ipynb @@ -0,0 +1,448 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a5ac7a99", + "metadata": {}, + "source": [ + "# Where, Offset, and domain" + ] + }, + { + "cell_type": "markdown", + "id": "3736a827", + "metadata": {}, + "source": [ + "## Conditional: where" + ] + }, + { + "cell_type": "markdown", + "id": "42bc3535", + "metadata": {}, + "source": [ + "The `where` builtin works analogously to the numpy version (https://numpy.org/doc/stable/reference/generated/numpy.where.html)\n", + "\n", + "Both require the same 3 input arguments:\n", + "- mask: a field of booleans or an expression evaluating to this type\n", + "- true branch: a tuple, a field, or a scalar\n", + "- false branch: a tuple, a field, of a scalar" + ] + }, + { + "cell_type": "markdown", + "id": "9b7552e8", + "metadata": {}, + "source": [ + "Take a simple numpy example, the `mask` here is a condition:" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "2bf82a64", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx\n", + "\n", + "backend = None\n", + "# backend = gtfn_cpu\n", + "# backend = gtfn_gpu" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "97ccaa9a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a_np array: [0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]\n", + "b_np array: [ 0. 1. 2. 3. 4. 5. 60. 70. 80. 90.]\n" + ] + } + ], + "source": [ + "a_np = np.arange(10.0)\n", + "b_np = np.where(a_np < 6.0, a_np, a_np*10.0)\n", + "print(\"a_np array: {}\".format(a_np))\n", + "print(\"b_np array: {}\".format(b_np))" + ] + }, + { + "cell_type": "markdown", + "id": "8056c3dd", + "metadata": {}, + "source": [ + "### **Task**: replicate this example in gt4py" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "80cbe9d5", + "metadata": {}, + "outputs": [ + { + "ename": "DSLError", + "evalue": "Undeclared or untyped symbol 'fieldop_where'.\n File \"/private/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56739/1995207298.py\", line 6", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mDSLError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[3], line 5\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# TODO implement the field_operator\u001b[39;00m\n\u001b[1;32m 4\u001b[0m \u001b[38;5;129;43m@gtx\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprogram\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;43;01mdef\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;21;43mprogram_where\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mb\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 6\u001b[0m \u001b[43m \u001b[49m\u001b[43mfieldop_where\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mb\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:532\u001b[0m, in \u001b[0;36mprogram..program_inner\u001b[0;34m(definition)\u001b[0m\n\u001b[1;32m 531\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprogram_inner\u001b[39m(definition: types\u001b[38;5;241m.\u001b[39mFunctionType) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Program:\n\u001b[0;32m--> 532\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mProgram\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_function\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 533\u001b[0m \u001b[43m \u001b[49m\u001b[43mdefinition\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mDEFAULT_BACKEND\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mis\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43meve\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mNOTHING\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrid_type\u001b[49m\n\u001b[1;32m 534\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:193\u001b[0m, in \u001b[0;36mProgram.from_function\u001b[0;34m(cls, definition, backend, grid_type)\u001b[0m\n\u001b[1;32m 191\u001b[0m closure_vars \u001b[38;5;241m=\u001b[39m get_closure_vars_from_function(definition)\n\u001b[1;32m 192\u001b[0m annotations \u001b[38;5;241m=\u001b[39m typing\u001b[38;5;241m.\u001b[39mget_type_hints(definition)\n\u001b[0;32m--> 193\u001b[0m past_node \u001b[38;5;241m=\u001b[39m \u001b[43mProgramParser\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43msource_def\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 194\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mcls\u001b[39m(\n\u001b[1;32m 195\u001b[0m past_node\u001b[38;5;241m=\u001b[39mpast_node,\n\u001b[1;32m 196\u001b[0m closure_vars\u001b[38;5;241m=\u001b[39mclosure_vars,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 199\u001b[0m grid_type\u001b[38;5;241m=\u001b[39mgrid_type,\n\u001b[1;32m 200\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/dialect_parser.py:74\u001b[0m, in \u001b[0;36mDialectParser.apply\u001b[0;34m(cls, source_definition, closure_vars, annotations)\u001b[0m\n\u001b[1;32m 72\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m RemoveDocstrings\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[1;32m 73\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m FixMissingLocations\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[0;32m---> 74\u001b[0m output_ast \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_postprocess_dialect_ast\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 75\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 76\u001b[0m \u001b[43m \u001b[49m\u001b[43msource_definition\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msource_definition\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 77\u001b[0m \u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 78\u001b[0m \u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mannotations\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 79\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_preprocess_definition_ast\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdefinition_ast\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 80\u001b[0m \u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 81\u001b[0m \u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 82\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 84\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m output_ast\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_past.py:42\u001b[0m, in \u001b[0;36mProgramParser._postprocess_dialect_ast\u001b[0;34m(cls, output_node, closure_vars, annotations)\u001b[0m\n\u001b[1;32m 37\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m 38\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_postprocess_dialect_ast\u001b[39m(\n\u001b[1;32m 39\u001b[0m \u001b[38;5;28mcls\u001b[39m, output_node: past\u001b[38;5;241m.\u001b[39mProgram, closure_vars: \u001b[38;5;28mdict\u001b[39m[\u001b[38;5;28mstr\u001b[39m, Any], annotations: \u001b[38;5;28mdict\u001b[39m[\u001b[38;5;28mstr\u001b[39m, Any]\n\u001b[1;32m 40\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m past\u001b[38;5;241m.\u001b[39mProgram:\n\u001b[1;32m 41\u001b[0m output_node \u001b[38;5;241m=\u001b[39m ClosureVarTypeDeduction\u001b[38;5;241m.\u001b[39mapply(output_node, closure_vars)\n\u001b[0;32m---> 42\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mProgramTypeDeduction\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutput_node\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/past_passes/type_deduction.py:97\u001b[0m, in \u001b[0;36mProgramTypeDeduction.apply\u001b[0;34m(cls, node)\u001b[0m\n\u001b[1;32m 95\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m 96\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mapply\u001b[39m(\u001b[38;5;28mcls\u001b[39m, node: past\u001b[38;5;241m.\u001b[39mProgram) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m past\u001b[38;5;241m.\u001b[39mProgram:\n\u001b[0;32m---> 97\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/past_passes/type_deduction.py:112\u001b[0m, in \u001b[0;36mProgramTypeDeduction.visit_Program\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 100\u001b[0m params \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mparams, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m 102\u001b[0m definition_type \u001b[38;5;241m=\u001b[39m ts\u001b[38;5;241m.\u001b[39mFunctionType(\n\u001b[1;32m 103\u001b[0m pos_only_args\u001b[38;5;241m=\u001b[39m[],\n\u001b[1;32m 104\u001b[0m pos_or_kw_args\u001b[38;5;241m=\u001b[39m{\u001b[38;5;28mstr\u001b[39m(param\u001b[38;5;241m.\u001b[39mid): param\u001b[38;5;241m.\u001b[39mtype \u001b[38;5;28;01mfor\u001b[39;00m param \u001b[38;5;129;01min\u001b[39;00m params},\n\u001b[1;32m 105\u001b[0m kw_only_args\u001b[38;5;241m=\u001b[39m{},\n\u001b[1;32m 106\u001b[0m returns\u001b[38;5;241m=\u001b[39mts\u001b[38;5;241m.\u001b[39mVoidType(),\n\u001b[1;32m 107\u001b[0m )\n\u001b[1;32m 108\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m past\u001b[38;5;241m.\u001b[39mProgram(\n\u001b[1;32m 109\u001b[0m \u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mid, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs),\n\u001b[1;32m 110\u001b[0m \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mts_ffront\u001b[38;5;241m.\u001b[39mProgramType(definition\u001b[38;5;241m=\u001b[39mdefinition_type),\n\u001b[1;32m 111\u001b[0m params\u001b[38;5;241m=\u001b[39mparams,\n\u001b[0;32m--> 112\u001b[0m body\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 113\u001b[0m closure_vars\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mclosure_vars, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs),\n\u001b[1;32m 114\u001b[0m location\u001b[38;5;241m=\u001b[39mnode\u001b[38;5;241m.\u001b[39mlocation,\n\u001b[1;32m 115\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:181\u001b[0m, in \u001b[0;36mNodeTranslator.generic_visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m new_node\n\u001b[1;32m 177\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mlist\u001b[39m, \u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mset\u001b[39m, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mSet)) \u001b[38;5;129;01mor\u001b[39;00m (\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28misinstance\u001b[39m(node, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mSequence) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mbytes\u001b[39m))\n\u001b[1;32m 179\u001b[0m ):\n\u001b[1;32m 180\u001b[0m \u001b[38;5;66;03m# Sequence or set: create a new container instance with the new values\u001b[39;00m\n\u001b[0;32m--> 181\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__class__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# type: ignore\u001b[39;49;00m\n\u001b[1;32m 182\u001b[0m \u001b[43m \u001b[49m\u001b[43mnew_child\u001b[49m\n\u001b[1;32m 183\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mchild\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtrees\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43miter_children_values\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 184\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[43mnew_child\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m:=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchild\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mis\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mNOTHING\u001b[49m\n\u001b[1;32m 185\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 187\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mdict\u001b[39m, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mMapping)):\n\u001b[1;32m 188\u001b[0m \u001b[38;5;66;03m# Mapping: create a new mapping instance with the new values\u001b[39;00m\n\u001b[1;32m 189\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m( \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[1;32m 190\u001b[0m {\n\u001b[1;32m 191\u001b[0m name: new_child\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 194\u001b[0m }\n\u001b[1;32m 195\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:184\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m new_node\n\u001b[1;32m 177\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mlist\u001b[39m, \u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mset\u001b[39m, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mSet)) \u001b[38;5;129;01mor\u001b[39;00m (\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28misinstance\u001b[39m(node, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mSequence) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mbytes\u001b[39m))\n\u001b[1;32m 179\u001b[0m ):\n\u001b[1;32m 180\u001b[0m \u001b[38;5;66;03m# Sequence or set: create a new container instance with the new values\u001b[39;00m\n\u001b[1;32m 181\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m( \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[1;32m 182\u001b[0m new_child\n\u001b[1;32m 183\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m child \u001b[38;5;129;01min\u001b[39;00m trees\u001b[38;5;241m.\u001b[39miter_children_values(node)\n\u001b[0;32m--> 184\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (new_child \u001b[38;5;241m:=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchild\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m) \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m NOTHING\n\u001b[1;32m 185\u001b[0m )\n\u001b[1;32m 187\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mdict\u001b[39m, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mMapping)):\n\u001b[1;32m 188\u001b[0m \u001b[38;5;66;03m# Mapping: create a new mapping instance with the new values\u001b[39;00m\n\u001b[1;32m 189\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m( \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[1;32m 190\u001b[0m {\n\u001b[1;32m 191\u001b[0m name: new_child\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 194\u001b[0m }\n\u001b[1;32m 195\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/past_passes/type_deduction.py:185\u001b[0m, in \u001b[0;36mProgramTypeDeduction.visit_Call\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 184\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvisit_Call\u001b[39m(\u001b[38;5;28mself\u001b[39m, node: past\u001b[38;5;241m.\u001b[39mCall, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m--> 185\u001b[0m new_func \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 186\u001b[0m new_args \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m 187\u001b[0m new_kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mkwargs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/past_passes/type_deduction.py:245\u001b[0m, in \u001b[0;36mProgramTypeDeduction.visit_Name\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 243\u001b[0m symtable \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m 244\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node\u001b[38;5;241m.\u001b[39mid \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m symtable \u001b[38;5;129;01mor\u001b[39;00m symtable[node\u001b[38;5;241m.\u001b[39mid]\u001b[38;5;241m.\u001b[39mtype \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 245\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m errors\u001b[38;5;241m.\u001b[39mDSLError(node\u001b[38;5;241m.\u001b[39mlocation, \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUndeclared or untyped symbol \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnode\u001b[38;5;241m.\u001b[39mid\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 247\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m past\u001b[38;5;241m.\u001b[39mName(\u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39mnode\u001b[38;5;241m.\u001b[39mid, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39msymtable[node\u001b[38;5;241m.\u001b[39mid]\u001b[38;5;241m.\u001b[39mtype, location\u001b[38;5;241m=\u001b[39mnode\u001b[38;5;241m.\u001b[39mlocation)\n", + "\u001b[0;31mDSLError\u001b[0m: Undeclared or untyped symbol 'fieldop_where'.\n File \"/private/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56739/1995207298.py\", line 6" + ] + } + ], + "source": [ + "# TODO implement the field_operator\n", + "\n", + "\n", + "@gtx.program(backend=backend)\n", + "def program_where(a: gtx.Field[Dims[K], float], b: gtx.Field[Dims[K], float]):\n", + " fieldop_where(a, out=b)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "c45e711e", + "metadata": {}, + "outputs": [], + "source": [ + "def test_where():\n", + " a = gtx.as_field([K], np.arange(10.0), allocator=backend)\n", + " b = gtx.as_field([K], np.zeros(shape=10), allocator=backend)\n", + " program_where(a, b, offset_provider={})\n", + " \n", + " assert np.allclose(b_np, b.asnumpy())" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "c196687b", + "metadata": {}, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'program_where' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_where\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[4], line 4\u001b[0m, in \u001b[0;36mtest_where\u001b[0;34m()\u001b[0m\n\u001b[1;32m 2\u001b[0m a \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mas_field([K], np\u001b[38;5;241m.\u001b[39marange(\u001b[38;5;241m10.0\u001b[39m), allocator\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[1;32m 3\u001b[0m b \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mas_field([K], np\u001b[38;5;241m.\u001b[39mzeros(shape\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m), allocator\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[0;32m----> 4\u001b[0m \u001b[43mprogram_where\u001b[49m(a, b, offset_provider\u001b[38;5;241m=\u001b[39m{})\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(b_np, b\u001b[38;5;241m.\u001b[39masnumpy())\n", + "\u001b[0;31mNameError\u001b[0m: name 'program_where' is not defined" + ] + } + ], + "source": [ + "test_where()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "markdown", + "id": "bcd232ab", + "metadata": {}, + "source": [ + "## Domain" + ] + }, + { + "cell_type": "markdown", + "id": "086273df", + "metadata": {}, + "source": [ + "The same operation can be performed in gt4py by including the `domain` keyowrd argument on `field_operator` call\n", + "\n", + "### **Task**: implement the same operation as above using `domain` instead of `where`" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "ebe2cd9d", + "metadata": {}, + "outputs": [ + { + "ename": "ValueError", + "evalue": "'' type is not supported.", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mValueError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[6], line 7\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;129m@gtx\u001b[39m\u001b[38;5;241m.\u001b[39mfield_operator\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfieldop_domain\u001b[39m(a: gtx\u001b[38;5;241m.\u001b[39mField[Dims[K], \u001b[38;5;28mfloat\u001b[39m]) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m gtx\u001b[38;5;241m.\u001b[39mField[Dims[K], \u001b[38;5;28mfloat\u001b[39m]:\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m a \u001b[38;5;241m*\u001b[39m \u001b[38;5;241m10.0\u001b[39m\n\u001b[1;32m 6\u001b[0m \u001b[38;5;129;43m@gtx\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprogram\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m----> 7\u001b[0m \u001b[38;5;28;43;01mdef\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;21;43mprogram_domain\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mb\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 8\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m \u001b[38;5;66;03m# TODO write the call to fieldop_domain\u001b[39;00m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:532\u001b[0m, in \u001b[0;36mprogram..program_inner\u001b[0;34m(definition)\u001b[0m\n\u001b[1;32m 531\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprogram_inner\u001b[39m(definition: types\u001b[38;5;241m.\u001b[39mFunctionType) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Program:\n\u001b[0;32m--> 532\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mProgram\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_function\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 533\u001b[0m \u001b[43m \u001b[49m\u001b[43mdefinition\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mDEFAULT_BACKEND\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mis\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43meve\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mNOTHING\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrid_type\u001b[49m\n\u001b[1;32m 534\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:193\u001b[0m, in \u001b[0;36mProgram.from_function\u001b[0;34m(cls, definition, backend, grid_type)\u001b[0m\n\u001b[1;32m 191\u001b[0m closure_vars \u001b[38;5;241m=\u001b[39m get_closure_vars_from_function(definition)\n\u001b[1;32m 192\u001b[0m annotations \u001b[38;5;241m=\u001b[39m typing\u001b[38;5;241m.\u001b[39mget_type_hints(definition)\n\u001b[0;32m--> 193\u001b[0m past_node \u001b[38;5;241m=\u001b[39m \u001b[43mProgramParser\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43msource_def\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 194\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mcls\u001b[39m(\n\u001b[1;32m 195\u001b[0m past_node\u001b[38;5;241m=\u001b[39mpast_node,\n\u001b[1;32m 196\u001b[0m closure_vars\u001b[38;5;241m=\u001b[39mclosure_vars,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 199\u001b[0m grid_type\u001b[38;5;241m=\u001b[39mgrid_type,\n\u001b[1;32m 200\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/dialect_parser.py:79\u001b[0m, in \u001b[0;36mDialectParser.apply\u001b[0;34m(cls, source_definition, closure_vars, annotations)\u001b[0m\n\u001b[1;32m 72\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m RemoveDocstrings\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[1;32m 73\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m FixMissingLocations\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[1;32m 74\u001b[0m output_ast \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m_postprocess_dialect_ast(\n\u001b[1;32m 75\u001b[0m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 76\u001b[0m \u001b[43m \u001b[49m\u001b[43msource_definition\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msource_definition\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 77\u001b[0m \u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 78\u001b[0m \u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mannotations\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m---> 79\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_preprocess_definition_ast\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdefinition_ast\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 80\u001b[0m closure_vars,\n\u001b[1;32m 81\u001b[0m annotations,\n\u001b[1;32m 82\u001b[0m )\n\u001b[1;32m 84\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m output_ast\n", + "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/ast.py:410\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 408\u001b[0m method \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvisit_\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m 409\u001b[0m visitor \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, method, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneric_visit)\n\u001b[0;32m--> 410\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_past.py:59\u001b[0m, in \u001b[0;36mProgramParser.visit_FunctionDef\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 44\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvisit_FunctionDef\u001b[39m(\u001b[38;5;28mself\u001b[39m, node: ast\u001b[38;5;241m.\u001b[39mFunctionDef) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m past\u001b[38;5;241m.\u001b[39mProgram:\n\u001b[1;32m 45\u001b[0m closure_symbols: \u001b[38;5;28mlist\u001b[39m[past\u001b[38;5;241m.\u001b[39mSymbol] \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m 46\u001b[0m past\u001b[38;5;241m.\u001b[39mSymbol(\n\u001b[1;32m 47\u001b[0m \u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39mname,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 52\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m name, val \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mclosure_vars\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m 53\u001b[0m ]\n\u001b[1;32m 55\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m past\u001b[38;5;241m.\u001b[39mProgram(\n\u001b[1;32m 56\u001b[0m \u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39mnode\u001b[38;5;241m.\u001b[39mname,\n\u001b[1;32m 57\u001b[0m \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mts\u001b[38;5;241m.\u001b[39mDeferredType(constraint\u001b[38;5;241m=\u001b[39mts_ffront\u001b[38;5;241m.\u001b[39mProgramType),\n\u001b[1;32m 58\u001b[0m params\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39margs),\n\u001b[0;32m---> 59\u001b[0m body\u001b[38;5;241m=\u001b[39m[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node) \u001b[38;5;28;01mfor\u001b[39;00m node \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mbody],\n\u001b[1;32m 60\u001b[0m closure_vars\u001b[38;5;241m=\u001b[39mclosure_symbols,\n\u001b[1;32m 61\u001b[0m location\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_location(node),\n\u001b[1;32m 62\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_past.py:59\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 44\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvisit_FunctionDef\u001b[39m(\u001b[38;5;28mself\u001b[39m, node: ast\u001b[38;5;241m.\u001b[39mFunctionDef) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m past\u001b[38;5;241m.\u001b[39mProgram:\n\u001b[1;32m 45\u001b[0m closure_symbols: \u001b[38;5;28mlist\u001b[39m[past\u001b[38;5;241m.\u001b[39mSymbol] \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m 46\u001b[0m past\u001b[38;5;241m.\u001b[39mSymbol(\n\u001b[1;32m 47\u001b[0m \u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39mname,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 52\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m name, val \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mclosure_vars\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m 53\u001b[0m ]\n\u001b[1;32m 55\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m past\u001b[38;5;241m.\u001b[39mProgram(\n\u001b[1;32m 56\u001b[0m \u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39mnode\u001b[38;5;241m.\u001b[39mname,\n\u001b[1;32m 57\u001b[0m \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mts\u001b[38;5;241m.\u001b[39mDeferredType(constraint\u001b[38;5;241m=\u001b[39mts_ffront\u001b[38;5;241m.\u001b[39mProgramType),\n\u001b[1;32m 58\u001b[0m params\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39margs),\n\u001b[0;32m---> 59\u001b[0m body\u001b[38;5;241m=\u001b[39m[\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m node \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mbody],\n\u001b[1;32m 60\u001b[0m closure_vars\u001b[38;5;241m=\u001b[39mclosure_symbols,\n\u001b[1;32m 61\u001b[0m location\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_location(node),\n\u001b[1;32m 62\u001b[0m )\n", + "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/ast.py:410\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 408\u001b[0m method \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvisit_\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m 409\u001b[0m visitor \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, method, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneric_visit)\n\u001b[0;32m--> 410\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_past.py:77\u001b[0m, in \u001b[0;36mProgramParser.visit_Expr\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 76\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvisit_Expr\u001b[39m(\u001b[38;5;28mself\u001b[39m, node: ast\u001b[38;5;241m.\u001b[39mExpr) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m past\u001b[38;5;241m.\u001b[39mLocatedNode:\n\u001b[0;32m---> 77\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/ast.py:410\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 408\u001b[0m method \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvisit_\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m 409\u001b[0m visitor \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, method, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneric_visit)\n\u001b[0;32m--> 410\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_past.py:172\u001b[0m, in \u001b[0;36mProgramParser.visit_Constant\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 171\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvisit_Constant\u001b[39m(\u001b[38;5;28mself\u001b[39m, node: ast\u001b[38;5;241m.\u001b[39mConstant) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m past\u001b[38;5;241m.\u001b[39mConstant:\n\u001b[0;32m--> 172\u001b[0m symbol_type \u001b[38;5;241m=\u001b[39m \u001b[43mtype_translation\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_value\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 173\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m past\u001b[38;5;241m.\u001b[39mConstant(value\u001b[38;5;241m=\u001b[39mnode\u001b[38;5;241m.\u001b[39mvalue, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39msymbol_type, location\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_location(node))\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/type_system/type_translation.py:200\u001b[0m, in \u001b[0;36mfrom_value\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m 198\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 199\u001b[0m type_ \u001b[38;5;241m=\u001b[39m xtyping\u001b[38;5;241m.\u001b[39minfer_type(value, annotate_callable_kwargs\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m--> 200\u001b[0m symbol_type \u001b[38;5;241m=\u001b[39m \u001b[43mfrom_type_hint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtype_\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 202\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(symbol_type, (ts\u001b[38;5;241m.\u001b[39mDataType, ts\u001b[38;5;241m.\u001b[39mCallableType, ts\u001b[38;5;241m.\u001b[39mOffsetType, ts\u001b[38;5;241m.\u001b[39mDimensionType)):\n\u001b[1;32m 203\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m symbol_type\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/type_system/type_translation.py:160\u001b[0m, in \u001b[0;36mfrom_type_hint\u001b[0;34m(type_hint, globalns, localns)\u001b[0m\n\u001b[1;32m 152\u001b[0m \u001b[38;5;66;03m# TODO(tehrengruber): print better error when no return type annotation is given\u001b[39;00m\n\u001b[1;32m 153\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ts\u001b[38;5;241m.\u001b[39mFunctionType(\n\u001b[1;32m 154\u001b[0m pos_only_args\u001b[38;5;241m=\u001b[39margs,\n\u001b[1;32m 155\u001b[0m pos_or_kw_args\u001b[38;5;241m=\u001b[39mkwargs,\n\u001b[1;32m 156\u001b[0m kw_only_args\u001b[38;5;241m=\u001b[39m{}, \u001b[38;5;66;03m# TODO\u001b[39;00m\n\u001b[1;32m 157\u001b[0m returns\u001b[38;5;241m=\u001b[39mrecursive_make_symbol(return_type),\n\u001b[1;32m 158\u001b[0m )\n\u001b[0;32m--> 160\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtype_hint\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m type is not supported.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "\u001b[0;31mValueError\u001b[0m: '' type is not supported." + ] + } + ], + "source": [ + "@gtx.field_operator\n", + "def fieldop_domain(a: gtx.Field[Dims[K], float]) -> gtx.Field[Dims[K], float]:\n", + " return a * 10.0\n", + "\n", + "\n", + "@gtx.program(backend=backend)\n", + "def program_domain(a: gtx.Field[Dims[K], float], b: gtx.Field[Dims[K], float]):\n", + " ... # TODO write the call to fieldop_domain" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "30fd5b6c", + "metadata": {}, + "outputs": [], + "source": [ + "def test_domain():\n", + " a = gtx.as_field([K], np.arange(10.0), allocator=backend)\n", + " b = gtx.as_field([K], np.arange(10.0), allocator=backend)\n", + " program_domain(a, b, offset_provider={})\n", + "\n", + " assert np.allclose(b_np, b.asnumpy())" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "f1827310", + "metadata": {}, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'program_domain' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_domain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[7], line 4\u001b[0m, in \u001b[0;36mtest_domain\u001b[0;34m()\u001b[0m\n\u001b[1;32m 2\u001b[0m a \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mas_field([K], np\u001b[38;5;241m.\u001b[39marange(\u001b[38;5;241m10.0\u001b[39m), allocator\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[1;32m 3\u001b[0m b \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mas_field([K], np\u001b[38;5;241m.\u001b[39marange(\u001b[38;5;241m10.0\u001b[39m), allocator\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[0;32m----> 4\u001b[0m \u001b[43mprogram_domain\u001b[49m(a, b, offset_provider\u001b[38;5;241m=\u001b[39m{})\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(b_np, b\u001b[38;5;241m.\u001b[39masnumpy())\n", + "\u001b[0;31mNameError\u001b[0m: name 'program_domain' is not defined" + ] + } + ], + "source": [ + "test_domain()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "markdown", + "id": "dc05317c", + "metadata": {}, + "source": [ + "## where and domain\n", + "\n", + "A combination of `where` and `domain` is useful in cases when an offset is used which exceeds the field size.\n", + "\n", + "e.g. a field `a: gtx.Field[Dims[K], float]` with shape (10,) is applied an offset (`Koff`)." + ] + }, + { + "cell_type": "markdown", + "id": "255fc343", + "metadata": {}, + "source": [ + "### **Task**: combine `domain` and `where` to account for extra indices\n", + "\n", + "Edit the code below such that:\n", + " 1. operations on field `a` are performed only up until the 8th index\n", + " 2. the domain is properly set accound for the offset" + ] + }, + { + "cell_type": "markdown", + "id": "caaa8a7d", + "metadata": {}, + "source": [ + "#### Python reference" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "47fe9be1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a_np_result array: [ 1. 3. 5. 7. 9. 11. 13. 15. 8. 0.]\n" + ] + } + ], + "source": [ + "a_np_result = np.zeros(shape=10)\n", + "for i in range(len(a_np)):\n", + " if a_np[i] < 8.0:\n", + " a_np_result[i] = a_np[i + 1] + a_np[i]\n", + " elif i < 9:\n", + " a_np_result[i] = a_np[i]\n", + "print(\"a_np_result array: {}\".format(a_np_result))" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "cac80fbb", + "metadata": {}, + "outputs": [ + { + "ename": "DSLError", + "evalue": "Must return a value, not None\n File \"/private/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56739/3624757481.py\", line 3", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mDSLError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[10], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;129;43m@gtx\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfield_operator\u001b[49m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;43;01mdef\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;21;43mfieldop_domain_where\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m>\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mreturn\u001b[39;49;00m \u001b[38;5;66;03m# TODO\u001b[39;00m\n\u001b[1;32m 5\u001b[0m \u001b[38;5;129m@gtx\u001b[39m\u001b[38;5;241m.\u001b[39mprogram(backend\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprogram_domain_where\u001b[39m(a: gtx\u001b[38;5;241m.\u001b[39mField[Dims[K], \u001b[38;5;28mfloat\u001b[39m], b: gtx\u001b[38;5;241m.\u001b[39mField[Dims[K], \u001b[38;5;28mfloat\u001b[39m]):\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:794\u001b[0m, in \u001b[0;36mfield_operator\u001b[0;34m(definition, backend, grid_type)\u001b[0m\n\u001b[1;32m 789\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfield_operator_inner\u001b[39m(definition: types\u001b[38;5;241m.\u001b[39mFunctionType) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m FieldOperator[foast\u001b[38;5;241m.\u001b[39mFieldOperator]:\n\u001b[1;32m 790\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m FieldOperator\u001b[38;5;241m.\u001b[39mfrom_function(\n\u001b[1;32m 791\u001b[0m definition, DEFAULT_BACKEND \u001b[38;5;28;01mif\u001b[39;00m backend \u001b[38;5;129;01mis\u001b[39;00m eve\u001b[38;5;241m.\u001b[39mNOTHING \u001b[38;5;28;01melse\u001b[39;00m backend, grid_type\n\u001b[1;32m 792\u001b[0m )\n\u001b[0;32m--> 794\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m field_operator_inner \u001b[38;5;28;01mif\u001b[39;00m definition \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[43mfield_operator_inner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdefinition\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:790\u001b[0m, in \u001b[0;36mfield_operator..field_operator_inner\u001b[0;34m(definition)\u001b[0m\n\u001b[1;32m 789\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfield_operator_inner\u001b[39m(definition: types\u001b[38;5;241m.\u001b[39mFunctionType) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m FieldOperator[foast\u001b[38;5;241m.\u001b[39mFieldOperator]:\n\u001b[0;32m--> 790\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mFieldOperator\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_function\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 791\u001b[0m \u001b[43m \u001b[49m\u001b[43mdefinition\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mDEFAULT_BACKEND\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mis\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43meve\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mNOTHING\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrid_type\u001b[49m\n\u001b[1;32m 792\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:589\u001b[0m, in \u001b[0;36mFieldOperator.from_function\u001b[0;34m(cls, definition, backend, grid_type, operator_node_cls, operator_attributes)\u001b[0m\n\u001b[1;32m 587\u001b[0m closure_vars \u001b[38;5;241m=\u001b[39m get_closure_vars_from_function(definition)\n\u001b[1;32m 588\u001b[0m annotations \u001b[38;5;241m=\u001b[39m typing\u001b[38;5;241m.\u001b[39mget_type_hints(definition)\n\u001b[0;32m--> 589\u001b[0m foast_definition_node \u001b[38;5;241m=\u001b[39m \u001b[43mFieldOperatorParser\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43msource_def\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 590\u001b[0m loc \u001b[38;5;241m=\u001b[39m foast_definition_node\u001b[38;5;241m.\u001b[39mlocation\n\u001b[1;32m 591\u001b[0m operator_attribute_nodes \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m 592\u001b[0m key: foast\u001b[38;5;241m.\u001b[39mConstant(value\u001b[38;5;241m=\u001b[39mvalue, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mtype_translation\u001b[38;5;241m.\u001b[39mfrom_value(value), location\u001b[38;5;241m=\u001b[39mloc)\n\u001b[1;32m 593\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m operator_attributes\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m 594\u001b[0m }\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/dialect_parser.py:79\u001b[0m, in \u001b[0;36mDialectParser.apply\u001b[0;34m(cls, source_definition, closure_vars, annotations)\u001b[0m\n\u001b[1;32m 72\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m RemoveDocstrings\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[1;32m 73\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m FixMissingLocations\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[1;32m 74\u001b[0m output_ast \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m_postprocess_dialect_ast(\n\u001b[1;32m 75\u001b[0m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 76\u001b[0m \u001b[43m \u001b[49m\u001b[43msource_definition\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msource_definition\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 77\u001b[0m \u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 78\u001b[0m \u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mannotations\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m---> 79\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_preprocess_definition_ast\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdefinition_ast\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 80\u001b[0m closure_vars,\n\u001b[1;32m 81\u001b[0m annotations,\n\u001b[1;32m 82\u001b[0m )\n\u001b[1;32m 84\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m output_ast\n", + "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/ast.py:410\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 408\u001b[0m method \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvisit_\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m 409\u001b[0m visitor \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, method, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneric_visit)\n\u001b[0;32m--> 410\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:170\u001b[0m, in \u001b[0;36mFieldOperatorParser.visit_FunctionDef\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 160\u001b[0m \u001b[38;5;28;01mcontinue\u001b[39;00m\n\u001b[1;32m 161\u001b[0m closure_var_symbols\u001b[38;5;241m.\u001b[39mappend(\n\u001b[1;32m 162\u001b[0m foast\u001b[38;5;241m.\u001b[39mSymbol(\n\u001b[1;32m 163\u001b[0m \u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39mname,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 167\u001b[0m )\n\u001b[1;32m 168\u001b[0m )\n\u001b[0;32m--> 170\u001b[0m new_body \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_visit_stmts\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_location\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 172\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m deduce_stmt_return_kind(new_body) \u001b[38;5;241m==\u001b[39m StmtReturnKind\u001b[38;5;241m.\u001b[39mNO_RETURN:\n\u001b[1;32m 173\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m errors\u001b[38;5;241m.\u001b[39mDSLError(loc, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mFunction\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m is expected to return a value.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:409\u001b[0m, in \u001b[0;36mFieldOperatorParser._visit_stmts\u001b[0;34m(self, stmts, location, **kwargs)\u001b[0m\n\u001b[1;32m 405\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_visit_stmts\u001b[39m(\n\u001b[1;32m 406\u001b[0m \u001b[38;5;28mself\u001b[39m, stmts: \u001b[38;5;28mlist\u001b[39m[ast\u001b[38;5;241m.\u001b[39mstmt], location: eve\u001b[38;5;241m.\u001b[39mSourceLocation, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m 407\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m foast\u001b[38;5;241m.\u001b[39mBlockStmt:\n\u001b[1;32m 408\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m foast\u001b[38;5;241m.\u001b[39mBlockStmt(\n\u001b[0;32m--> 409\u001b[0m stmts\u001b[38;5;241m=\u001b[39m[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(el, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;28;01mfor\u001b[39;00m el \u001b[38;5;129;01min\u001b[39;00m stmts \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(el, ast\u001b[38;5;241m.\u001b[39mPass)],\n\u001b[1;32m 410\u001b[0m location\u001b[38;5;241m=\u001b[39mlocation,\n\u001b[1;32m 411\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:409\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 405\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_visit_stmts\u001b[39m(\n\u001b[1;32m 406\u001b[0m \u001b[38;5;28mself\u001b[39m, stmts: \u001b[38;5;28mlist\u001b[39m[ast\u001b[38;5;241m.\u001b[39mstmt], location: eve\u001b[38;5;241m.\u001b[39mSourceLocation, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m 407\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m foast\u001b[38;5;241m.\u001b[39mBlockStmt:\n\u001b[1;32m 408\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m foast\u001b[38;5;241m.\u001b[39mBlockStmt(\n\u001b[0;32m--> 409\u001b[0m stmts\u001b[38;5;241m=\u001b[39m[\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m el \u001b[38;5;129;01min\u001b[39;00m stmts \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(el, ast\u001b[38;5;241m.\u001b[39mPass)],\n\u001b[1;32m 410\u001b[0m location\u001b[38;5;241m=\u001b[39mlocation,\n\u001b[1;32m 411\u001b[0m )\n", + "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/ast.py:410\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 408\u001b[0m method \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvisit_\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m 409\u001b[0m visitor \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, method, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneric_visit)\n\u001b[0;32m--> 410\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:316\u001b[0m, in \u001b[0;36mFieldOperatorParser.visit_Return\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 314\u001b[0m loc \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_location(node)\n\u001b[1;32m 315\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m node\u001b[38;5;241m.\u001b[39mvalue:\n\u001b[0;32m--> 316\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m errors\u001b[38;5;241m.\u001b[39mDSLError(loc, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMust return a value, not None\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 317\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m foast\u001b[38;5;241m.\u001b[39mReturn(value\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mvalue), location\u001b[38;5;241m=\u001b[39mloc)\n", + "\u001b[0;31mDSLError\u001b[0m: Must return a value, not None\n File \"/private/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56739/3624757481.py\", line 3" + ] + } + ], + "source": [ + "@gtx.field_operator\n", + "def fieldop_domain_where(a: gtx.Field[Dims[K], float]) -> gtx.Field[Dims[K], float]:\n", + " return # TODO\n", + "\n", + "@gtx.program(backend=backend)\n", + "def program_domain_where(a: gtx.Field[Dims[K], float], b: gtx.Field[Dims[K], float]):\n", + " ... # TODO " + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "ed8254a8", + "metadata": {}, + "outputs": [], + "source": [ + "def test_domain_where(): \n", + " a = gtx.as_field([K], np.arange(10.0), allocator=backend)\n", + " b = gtx.as_field([K], np.zeros(shape=10), allocator=backend)\n", + " program_domain_where(a, b, offset_provider={\"Koff\": K})\n", + " \n", + " assert np.allclose(a_np_result, b.asnumpy())" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "ed393959", + "metadata": {}, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'program_domain_where' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[12], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_domain_where\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[11], line 4\u001b[0m, in \u001b[0;36mtest_domain_where\u001b[0;34m()\u001b[0m\n\u001b[1;32m 2\u001b[0m a \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mas_field([K], np\u001b[38;5;241m.\u001b[39marange(\u001b[38;5;241m10.0\u001b[39m), allocator\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[1;32m 3\u001b[0m b \u001b[38;5;241m=\u001b[39m gtx\u001b[38;5;241m.\u001b[39mas_field([K], np\u001b[38;5;241m.\u001b[39mzeros(shape\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m), allocator\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[0;32m----> 4\u001b[0m \u001b[43mprogram_domain_where\u001b[49m(a, b, offset_provider\u001b[38;5;241m=\u001b[39m{\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mKoff\u001b[39m\u001b[38;5;124m\"\u001b[39m: K})\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(a_np_result, b\u001b[38;5;241m.\u001b[39masnumpy())\n", + "\u001b[0;31mNameError\u001b[0m: name 'program_domain_where' is not defined" + ] + } + ], + "source": [ + "test_domain_where()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b37a1eec", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/6_where_domain_solutions.ipynb b/docs/user/next/workshop/exercises/6_where_domain_solutions.ipynb new file mode 100644 index 0000000000..f42701473c --- /dev/null +++ b/docs/user/next/workshop/exercises/6_where_domain_solutions.ipynb @@ -0,0 +1,382 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a5ac7a99", + "metadata": {}, + "source": [ + "# Where, Offset, and domain" + ] + }, + { + "cell_type": "markdown", + "id": "3736a827", + "metadata": {}, + "source": [ + "## Conditional: where" + ] + }, + { + "cell_type": "markdown", + "id": "42bc3535", + "metadata": {}, + "source": [ + "The `where` builtin works analogously to the numpy version (https://numpy.org/doc/stable/reference/generated/numpy.where.html)\n", + "\n", + "Both require the same 3 input arguments:\n", + "- mask: a field of booleans or an expression evaluating to this type\n", + "- true branch: a tuple, a field, or a scalar\n", + "- false branch: a tuple, a field, of a scalar" + ] + }, + { + "cell_type": "markdown", + "id": "9b7552e8", + "metadata": {}, + "source": [ + "Take a simple numpy example, the `mask` here is a condition:" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "2bf82a64", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx\n", + "\n", + "backend = None\n", + "# backend = gtfn_cpu\n", + "# backend = gtfn_gpu" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "97ccaa9a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a_np array: [0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]\n", + "b_np array: [ 0. 1. 2. 3. 4. 5. 60. 70. 80. 90.]\n" + ] + } + ], + "source": [ + "a_np = np.arange(10.0)\n", + "b_np = np.where(a_np < 6.0, a_np, a_np*10.0)\n", + "print(\"a_np array: {}\".format(a_np))\n", + "print(\"b_np array: {}\".format(b_np))" + ] + }, + { + "cell_type": "markdown", + "id": "8056c3dd", + "metadata": {}, + "source": [ + "### **Task**: replicate this example in gt4py" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "80cbe9d5", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def fieldop_where(a: gtx.Field[Dims[K], float]) -> gtx.Field[Dims[K], float]:\n", + " return where(a < 6.0, a, a*10.0)\n", + "\n", + "@gtx.program(backend=backend)\n", + "def program_where(a: gtx.Field[Dims[K], float], b: gtx.Field[Dims[K], float]):\n", + " fieldop_where(a, out=b) " + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "c45e711e", + "metadata": {}, + "outputs": [], + "source": [ + "def test_where():\n", + " a = gtx.as_field([K], np.arange(10.0), allocator=backend)\n", + " b = gtx.as_field([K], np.zeros(shape=10), allocator=backend)\n", + " program_where(a, b, offset_provider={})\n", + " \n", + " assert np.allclose(b_np, b.asnumpy())" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "c196687b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Test successful\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56692/2672671145.py:4: UserWarning: Field View Program 'program_where': Using Python execution, consider selecting a perfomance backend.\n", + " program_where(a, b, offset_provider={})\n" + ] + } + ], + "source": [ + "test_where()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "markdown", + "id": "bcd232ab", + "metadata": {}, + "source": [ + "## Domain" + ] + }, + { + "cell_type": "markdown", + "id": "086273df", + "metadata": {}, + "source": [ + "The same operation can be performed in gt4py by including the `domain` keyowrd argument on `field_operator` call\n", + "\n", + "### **Task**: implement the same operation as above using `domain` instead of `where`" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "ebe2cd9d", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def fieldop_domain(a: gtx.Field[Dims[K], float]) -> gtx.Field[Dims[K], float]:\n", + " return a*10.0\n", + "\n", + "@gtx.program(backend=backend)\n", + "def program_domain(a: gtx.Field[Dims[K], float],\n", + " b: gtx.Field[Dims[K], float]):\n", + " fieldop_domain(a, out=b, domain={K: (6, 10)}) " + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "30fd5b6c", + "metadata": {}, + "outputs": [], + "source": [ + "def test_domain():\n", + " a = gtx.as_field([K], np.arange(10.0), allocator=backend)\n", + " b = gtx.as_field([K], np.arange(10.0), allocator=backend)\n", + " program_domain(a, b, offset_provider={})\n", + "\n", + " assert np.allclose(b_np, b.asnumpy())" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "f1827310", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Test successful\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56692/2292331765.py:4: UserWarning: Field View Program 'program_domain': Using Python execution, consider selecting a perfomance backend.\n", + " program_domain(a, b, offset_provider={})\n" + ] + } + ], + "source": [ + "test_domain()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "markdown", + "id": "dc05317c", + "metadata": {}, + "source": [ + "## where and domain\n", + "\n", + "A combination of `where` and `domain` is useful in cases when an offset is used which exceeds the field size.\n", + "\n", + "e.g. a field `a: gtx.Field[Dims[K], float]` with shape (10,) is applied an offset (`Koff`)." + ] + }, + { + "cell_type": "markdown", + "id": "255fc343", + "metadata": {}, + "source": [ + "### **Task**: combine `domain` and `where` to account for extra indices\n", + "\n", + "Edit the code below such that:\n", + " 1. operations on field `a` are performed only up until the 8th index\n", + " 2. the domain is properly set accound for the offset" + ] + }, + { + "cell_type": "markdown", + "id": "caaa8a7d", + "metadata": {}, + "source": [ + "#### Python reference" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "47fe9be1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a_np_result array: [ 1. 3. 5. 7. 9. 11. 13. 15. 8. 0.]\n" + ] + } + ], + "source": [ + "a_np_result = np.zeros(shape=10)\n", + "for i in range(len(a_np)):\n", + " if a_np[i] < 8.0:\n", + " a_np_result[i] = a_np[i + 1] + a_np[i]\n", + " elif i < 9:\n", + " a_np_result[i] = a_np[i]\n", + "print(\"a_np_result array: {}\".format(a_np_result))" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "cac80fbb", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def fieldop_domain_where(a: gtx.Field[Dims[K], float]) -> gtx.Field[Dims[K], float]:\n", + " return where(a<8.0, a(Koff[1])+a, a)\n", + "\n", + "@gtx.program(backend=backend)\n", + "def program_domain_where(a: gtx.Field[Dims[K], float], b: gtx.Field[Dims[K], float]):\n", + " fieldop_domain_where(a, out=b, domain={K: (0, 9)}) " + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "ed8254a8", + "metadata": {}, + "outputs": [], + "source": [ + "def test_domain_where(): \n", + " a = gtx.as_field([K], np.arange(10.0), allocator=backend)\n", + " b = gtx.as_field([K], np.zeros(shape=10), allocator=backend)\n", + " program_domain_where(a, b, offset_provider={\"Koff\": K})\n", + " \n", + " assert np.allclose(a_np_result, b.asnumpy())" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "ed393959", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Test successful\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56692/57164705.py:4: UserWarning: Field View Program 'program_domain_where': Using Python execution, consider selecting a perfomance backend.\n", + " program_domain_where(a, b, offset_provider={\"Koff\": K})\n" + ] + } + ], + "source": [ + "test_domain_where()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b37a1eec", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/7_scan_operator.ipynb b/docs/user/next/workshop/exercises/7_scan_operator.ipynb new file mode 100644 index 0000000000..90982c352d --- /dev/null +++ b/docs/user/next/workshop/exercises/7_scan_operator.ipynb @@ -0,0 +1,343 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9ba87bfb", + "metadata": {}, + "source": [ + "## Scan operator" + ] + }, + { + "cell_type": "markdown", + "id": "2ee9d989", + "metadata": {}, + "source": [ + "The unique feature of this operator is that it provides the return state of the previous iteration as its first argument (i.e., the result from the previous grid point). In other words, all the arguments of the current `return` will be available (as a tuple) in the next iteration from the first argument of the defined function. \n", + "\n", + "Example: A FORTRAN pseudocode for integrating a moisture variable (e.g., cloud water or water vapour) over a column could look as follows:\n", + "\n", + "\n", + "```FORTRAN\n", + "SUBROUTINE column_integral( var_in, rho, dz, var_out, ie, je, ke )\n", + " ! Return the column integral of a moist species.\n", + " INTEGER, INTENT (IN) :: &\n", + " ie, je, ke ! array dimensions of the I/O-fields (horizontal, horizontal, vertical)\n", + "\n", + " REAL (KIND=wp), INTENT (OUT) :: &\n", + " q_colsum (ie,je) ! Vertically-integrated mass of water species\n", + "\n", + " REAL (KIND=wp), INTENT (IN) :: &\n", + " rho (ie,je,ke), & \n", + " dz (ie,je,ke), & ! height of model half levels\n", + " var_in (ie,je,ke) ! humidity mass concentration at time-level nnow\n", + " \n", + " !$acc parallel present( iq ) if (lzacc)\n", + " !$acc loop gang\n", + " DO j=1,je\n", + " !$acc loop vector\n", + " DO i=1,ie\n", + " q_sum(i,j) = 0.0\n", + " END DO\n", + " END DO\n", + " !$acc end parallel\n", + " \n", + " \n", + " !$acc parallel present( iq, rho, hhl, q ) if (lzacc)\n", + " DO k = 1, ke ! Vertical loop\n", + " !$acc loop gang\n", + " DO j=1,je\n", + " !$acc loop vector\n", + " DO i=1,ie\n", + " q_colsum(i,j) = q_colsum(i,j) + var_in(i,j,k) * rho(i,j,k)* dz(i,j,k)\n", + " END DO\n", + " END DO\n", + " END DO\n", + " !$acc end parallel\n", + "END SUBROUTINE column_integral\n", + "```\n", + "\n", + "Where:\n", + "- `var_in` is the 3D variable that will be summed up\n", + "- `q_colsum` is the resulting 2D variable\n", + "- `rho` the air density\n", + "- `dz`the thickness of the vertical layers\n", + "\n", + "In the first loop nest, `column_sum` is set to zero for all grid columns. The vertical dependency enters on the RHS of the second loop nest `q_colsum(i,j) = q_colsum(i,j) + ...`\n", + "\n", + "Using the `scan_operator` this operation would be written like this:\n", + "\n", + "```python\n", + "@scan_operator(axis=KDim, forward=True, init=0.0)\n", + "def column_integral(float: state, float: var, float: rho, float: dz)\n", + " \"\"\"Return the column integral of a moist species.\"\"\"\n", + " return var * rho * dz + state\n", + "```\n", + "\n", + "Here the vertical dependency is expressed by the first function argument (`state`). This argument carries the return from the previous k-level and does not need to be specified when the function is called (similar to the `self` argument of Python classes). The argument is intialized to `init=0.0` in the function decorator (first loop nest above) and the dimension of the integral is specified with `axis=KDim`.\n", + "\n", + "\n", + "```python\n", + "q_colsum = column_integral(qv, rho, dz)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "c9e31bff", + "metadata": {}, + "source": [ + "#### Exercise: port a toy cloud microphysics scheme from python/numpy using the template of a `scan_operator` below" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "bd2fd309", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx\n", + "\n", + "backend = None\n", + "# backend = gtfn_cpu\n", + "# backend = gtfn_gpu" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "74338168", + "metadata": {}, + "outputs": [], + "source": [ + "def toy_microphysics_numpy(qc, qr, autoconversion_rate=0.1, sedimentaion_constant=0.05):\n", + " \"\"\"A toy model of a microphysics scheme contaning autoconversion and scavenging\"\"\"\n", + "\n", + " sedimentation_flux = 0.0\n", + "\n", + " for cell, k in np.ndindex(qc.shape):\n", + " # Autoconversion: Cloud Drops -> Rain Drops\n", + " autoconversion_tendency = qc[cell, k] * autoconversion_rate\n", + "\n", + " qc[cell, k] -= autoconversion_tendency\n", + " qr[cell, k] += autoconversion_tendency\n", + "\n", + " ## Apply sedimentation flux from level above\n", + " qr[cell, k] += sedimentation_flux\n", + "\n", + " ## Remove mass due to sedimentation flux from the current cell\n", + " qr[cell, k] -= sedimentation_flux" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "69bf6022", + "metadata": {}, + "outputs": [ + { + "ename": "DSLError", + "evalue": "Must return a value, not None\n File \"/private/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56865/3074024806.py\", line 20", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mDSLError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[3], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;129;43m@gtx\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscan_operator\u001b[49m\u001b[43m(\u001b[49m\u001b[43maxis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mK\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mforward\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minit\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m0.0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0.0\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0.0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;43;01mdef\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;21;43m_graupel_toy_scan\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[43m \u001b[49m\u001b[43mstate\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mtuple\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqc_in\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqr_in\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\n\u001b[1;32m 4\u001b[0m \u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m>\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mtuple\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 5\u001b[0m \u001b[43m \u001b[49m\u001b[43mautoconversion_rate\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0.1\u001b[39;49m\n\u001b[1;32m 6\u001b[0m \u001b[43m \u001b[49m\u001b[43msedimentaion_constant\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0.05\u001b[39;49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:862\u001b[0m, in \u001b[0;36mscan_operator..scan_operator_inner\u001b[0;34m(definition)\u001b[0m\n\u001b[1;32m 861\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mscan_operator_inner\u001b[39m(definition: types\u001b[38;5;241m.\u001b[39mFunctionType) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m FieldOperator:\n\u001b[0;32m--> 862\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mFieldOperator\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_function\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 863\u001b[0m \u001b[43m \u001b[49m\u001b[43mdefinition\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 864\u001b[0m \u001b[43m \u001b[49m\u001b[43mDEFAULT_BACKEND\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mis\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43meve\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mNOTHING\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 865\u001b[0m \u001b[43m \u001b[49m\u001b[43mgrid_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 866\u001b[0m \u001b[43m \u001b[49m\u001b[43moperator_node_cls\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfoast\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mScanOperator\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 867\u001b[0m \u001b[43m \u001b[49m\u001b[43moperator_attributes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m{\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43maxis\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mforward\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mforward\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43minit\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43minit\u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 868\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:589\u001b[0m, in \u001b[0;36mFieldOperator.from_function\u001b[0;34m(cls, definition, backend, grid_type, operator_node_cls, operator_attributes)\u001b[0m\n\u001b[1;32m 587\u001b[0m closure_vars \u001b[38;5;241m=\u001b[39m get_closure_vars_from_function(definition)\n\u001b[1;32m 588\u001b[0m annotations \u001b[38;5;241m=\u001b[39m typing\u001b[38;5;241m.\u001b[39mget_type_hints(definition)\n\u001b[0;32m--> 589\u001b[0m foast_definition_node \u001b[38;5;241m=\u001b[39m \u001b[43mFieldOperatorParser\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43msource_def\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 590\u001b[0m loc \u001b[38;5;241m=\u001b[39m foast_definition_node\u001b[38;5;241m.\u001b[39mlocation\n\u001b[1;32m 591\u001b[0m operator_attribute_nodes \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m 592\u001b[0m key: foast\u001b[38;5;241m.\u001b[39mConstant(value\u001b[38;5;241m=\u001b[39mvalue, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mtype_translation\u001b[38;5;241m.\u001b[39mfrom_value(value), location\u001b[38;5;241m=\u001b[39mloc)\n\u001b[1;32m 593\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m operator_attributes\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m 594\u001b[0m }\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/dialect_parser.py:79\u001b[0m, in \u001b[0;36mDialectParser.apply\u001b[0;34m(cls, source_definition, closure_vars, annotations)\u001b[0m\n\u001b[1;32m 72\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m RemoveDocstrings\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[1;32m 73\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m FixMissingLocations\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[1;32m 74\u001b[0m output_ast \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m_postprocess_dialect_ast(\n\u001b[1;32m 75\u001b[0m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 76\u001b[0m \u001b[43m \u001b[49m\u001b[43msource_definition\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msource_definition\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 77\u001b[0m \u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 78\u001b[0m \u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mannotations\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m---> 79\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_preprocess_definition_ast\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdefinition_ast\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m 80\u001b[0m closure_vars,\n\u001b[1;32m 81\u001b[0m annotations,\n\u001b[1;32m 82\u001b[0m )\n\u001b[1;32m 84\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m output_ast\n", + "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/ast.py:410\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 408\u001b[0m method \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvisit_\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m 409\u001b[0m visitor \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, method, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneric_visit)\n\u001b[0;32m--> 410\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:170\u001b[0m, in \u001b[0;36mFieldOperatorParser.visit_FunctionDef\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 160\u001b[0m \u001b[38;5;28;01mcontinue\u001b[39;00m\n\u001b[1;32m 161\u001b[0m closure_var_symbols\u001b[38;5;241m.\u001b[39mappend(\n\u001b[1;32m 162\u001b[0m foast\u001b[38;5;241m.\u001b[39mSymbol(\n\u001b[1;32m 163\u001b[0m \u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39mname,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 167\u001b[0m )\n\u001b[1;32m 168\u001b[0m )\n\u001b[0;32m--> 170\u001b[0m new_body \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_visit_stmts\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_location\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 172\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m deduce_stmt_return_kind(new_body) \u001b[38;5;241m==\u001b[39m StmtReturnKind\u001b[38;5;241m.\u001b[39mNO_RETURN:\n\u001b[1;32m 173\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m errors\u001b[38;5;241m.\u001b[39mDSLError(loc, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mFunction\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m is expected to return a value.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:409\u001b[0m, in \u001b[0;36mFieldOperatorParser._visit_stmts\u001b[0;34m(self, stmts, location, **kwargs)\u001b[0m\n\u001b[1;32m 405\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_visit_stmts\u001b[39m(\n\u001b[1;32m 406\u001b[0m \u001b[38;5;28mself\u001b[39m, stmts: \u001b[38;5;28mlist\u001b[39m[ast\u001b[38;5;241m.\u001b[39mstmt], location: eve\u001b[38;5;241m.\u001b[39mSourceLocation, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m 407\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m foast\u001b[38;5;241m.\u001b[39mBlockStmt:\n\u001b[1;32m 408\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m foast\u001b[38;5;241m.\u001b[39mBlockStmt(\n\u001b[0;32m--> 409\u001b[0m stmts\u001b[38;5;241m=\u001b[39m[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(el, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;28;01mfor\u001b[39;00m el \u001b[38;5;129;01min\u001b[39;00m stmts \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(el, ast\u001b[38;5;241m.\u001b[39mPass)],\n\u001b[1;32m 410\u001b[0m location\u001b[38;5;241m=\u001b[39mlocation,\n\u001b[1;32m 411\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:409\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 405\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_visit_stmts\u001b[39m(\n\u001b[1;32m 406\u001b[0m \u001b[38;5;28mself\u001b[39m, stmts: \u001b[38;5;28mlist\u001b[39m[ast\u001b[38;5;241m.\u001b[39mstmt], location: eve\u001b[38;5;241m.\u001b[39mSourceLocation, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m 407\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m foast\u001b[38;5;241m.\u001b[39mBlockStmt:\n\u001b[1;32m 408\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m foast\u001b[38;5;241m.\u001b[39mBlockStmt(\n\u001b[0;32m--> 409\u001b[0m stmts\u001b[38;5;241m=\u001b[39m[\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m el \u001b[38;5;129;01min\u001b[39;00m stmts \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(el, ast\u001b[38;5;241m.\u001b[39mPass)],\n\u001b[1;32m 410\u001b[0m location\u001b[38;5;241m=\u001b[39mlocation,\n\u001b[1;32m 411\u001b[0m )\n", + "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/ast.py:410\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m 408\u001b[0m method \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mvisit_\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m 409\u001b[0m visitor \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, method, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgeneric_visit)\n\u001b[0;32m--> 410\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:316\u001b[0m, in \u001b[0;36mFieldOperatorParser.visit_Return\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 314\u001b[0m loc \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_location(node)\n\u001b[1;32m 315\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m node\u001b[38;5;241m.\u001b[39mvalue:\n\u001b[0;32m--> 316\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m errors\u001b[38;5;241m.\u001b[39mDSLError(loc, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMust return a value, not None\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 317\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m foast\u001b[38;5;241m.\u001b[39mReturn(value\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mvalue), location\u001b[38;5;241m=\u001b[39mloc)\n", + "\u001b[0;31mDSLError\u001b[0m: Must return a value, not None\n File \"/private/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56865/3074024806.py\", line 20" + ] + } + ], + "source": [ + "@gtx.scan_operator(axis=K, forward=True, init=(0.0, 0.0, 0.0))\n", + "def _graupel_toy_scan(\n", + " state: tuple[float, float, float], qc_in: float, qr_in: float\n", + ") -> tuple[float, float, float]:\n", + " autoconversion_rate = 0.1\n", + " sedimentaion_constant = 0.05\n", + "\n", + " # unpack state of previous iteration\n", + " # TODO\n", + " \n", + " # Autoconversion: Cloud Drops -> Rain Drops\n", + " # TODO\n", + " \n", + " ## Add sedimentation flux from level above\n", + " # TODO\n", + "\n", + " # Remove mass due to sedimentation flux\n", + " # TODO\n", + "\n", + " return # TODO" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "0de41cb8", + "metadata": {}, + "outputs": [ + { + "ename": "DSLError", + "evalue": "Undeclared symbol '_graupel_toy_scan'.\n File \"/private/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56865/2647114660.py\", line 8", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mDSLError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[4], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;129;43m@gtx\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfield_operator\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;43;01mdef\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;21;43mgraupel_toy_scan\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[43m \u001b[49m\u001b[43mqc\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mC\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqr\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mC\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\n\u001b[1;32m 4\u001b[0m \u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m>\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mtuple\u001b[39;49m\u001b[43m[\u001b[49m\n\u001b[1;32m 5\u001b[0m \u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mC\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 6\u001b[0m \u001b[43m \u001b[49m\u001b[43mgtx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mField\u001b[49m\u001b[43m[\u001b[49m\u001b[43mDims\u001b[49m\u001b[43m[\u001b[49m\u001b[43mC\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mK\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mfloat\u001b[39;49m\u001b[43m]\u001b[49m\n\u001b[1;32m 7\u001b[0m \u001b[43m]\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 8\u001b[0m \u001b[43m \u001b[49m\u001b[43mqc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43m_graupel_toy_scan\u001b[49m\u001b[43m(\u001b[49m\u001b[43mqc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqr\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 10\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mreturn\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mqc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqr\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:790\u001b[0m, in \u001b[0;36mfield_operator..field_operator_inner\u001b[0;34m(definition)\u001b[0m\n\u001b[1;32m 789\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfield_operator_inner\u001b[39m(definition: types\u001b[38;5;241m.\u001b[39mFunctionType) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m FieldOperator[foast\u001b[38;5;241m.\u001b[39mFieldOperator]:\n\u001b[0;32m--> 790\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mFieldOperator\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_function\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 791\u001b[0m \u001b[43m \u001b[49m\u001b[43mdefinition\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mDEFAULT_BACKEND\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mis\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43meve\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mNOTHING\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrid_type\u001b[49m\n\u001b[1;32m 792\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:589\u001b[0m, in \u001b[0;36mFieldOperator.from_function\u001b[0;34m(cls, definition, backend, grid_type, operator_node_cls, operator_attributes)\u001b[0m\n\u001b[1;32m 587\u001b[0m closure_vars \u001b[38;5;241m=\u001b[39m get_closure_vars_from_function(definition)\n\u001b[1;32m 588\u001b[0m annotations \u001b[38;5;241m=\u001b[39m typing\u001b[38;5;241m.\u001b[39mget_type_hints(definition)\n\u001b[0;32m--> 589\u001b[0m foast_definition_node \u001b[38;5;241m=\u001b[39m \u001b[43mFieldOperatorParser\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43msource_def\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 590\u001b[0m loc \u001b[38;5;241m=\u001b[39m foast_definition_node\u001b[38;5;241m.\u001b[39mlocation\n\u001b[1;32m 591\u001b[0m operator_attribute_nodes \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m 592\u001b[0m key: foast\u001b[38;5;241m.\u001b[39mConstant(value\u001b[38;5;241m=\u001b[39mvalue, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39mtype_translation\u001b[38;5;241m.\u001b[39mfrom_value(value), location\u001b[38;5;241m=\u001b[39mloc)\n\u001b[1;32m 593\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m operator_attributes\u001b[38;5;241m.\u001b[39mitems()\n\u001b[1;32m 594\u001b[0m }\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/dialect_parser.py:74\u001b[0m, in \u001b[0;36mDialectParser.apply\u001b[0;34m(cls, source_definition, closure_vars, annotations)\u001b[0m\n\u001b[1;32m 72\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m RemoveDocstrings\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[1;32m 73\u001b[0m definition_ast \u001b[38;5;241m=\u001b[39m FixMissingLocations\u001b[38;5;241m.\u001b[39mapply(definition_ast)\n\u001b[0;32m---> 74\u001b[0m output_ast \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_postprocess_dialect_ast\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 75\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m 76\u001b[0m \u001b[43m \u001b[49m\u001b[43msource_definition\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msource_definition\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 77\u001b[0m \u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 78\u001b[0m \u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mannotations\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 79\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_preprocess_definition_ast\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdefinition_ast\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 80\u001b[0m \u001b[43m \u001b[49m\u001b[43mclosure_vars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 81\u001b[0m \u001b[43m \u001b[49m\u001b[43mannotations\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 82\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 84\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m output_ast\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/func_to_foast.py:101\u001b[0m, in \u001b[0;36mFieldOperatorParser._postprocess_dialect_ast\u001b[0;34m(cls, foast_node, closure_vars, annotations)\u001b[0m\n\u001b[1;32m 99\u001b[0m foast_node \u001b[38;5;241m=\u001b[39m DeadClosureVarElimination\u001b[38;5;241m.\u001b[39mapply(foast_node)\n\u001b[1;32m 100\u001b[0m foast_node \u001b[38;5;241m=\u001b[39m ClosureVarTypeDeduction\u001b[38;5;241m.\u001b[39mapply(foast_node, closure_vars)\n\u001b[0;32m--> 101\u001b[0m foast_node \u001b[38;5;241m=\u001b[39m \u001b[43mFieldOperatorTypeDeduction\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mapply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfoast_node\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 102\u001b[0m foast_node \u001b[38;5;241m=\u001b[39m UnpackedAssignPass\u001b[38;5;241m.\u001b[39mapply(foast_node)\n\u001b[1;32m 104\u001b[0m \u001b[38;5;66;03m# check deduced matches annotated return type\u001b[39;00m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/foast_passes/type_deduction.py:252\u001b[0m, in \u001b[0;36mFieldOperatorTypeDeduction.apply\u001b[0;34m(cls, node)\u001b[0m\n\u001b[1;32m 250\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m 251\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mapply\u001b[39m(\u001b[38;5;28mcls\u001b[39m, node: foast\u001b[38;5;241m.\u001b[39mFunctionDefinition) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m foast\u001b[38;5;241m.\u001b[39mFunctionDefinition:\n\u001b[0;32m--> 252\u001b[0m typed_foast_node \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 254\u001b[0m FieldOperatorTypeDeductionCompletnessValidator\u001b[38;5;241m.\u001b[39mapply(typed_foast_node)\n\u001b[1;32m 256\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m typed_foast_node\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/foast_passes/type_deduction.py:260\u001b[0m, in \u001b[0;36mFieldOperatorTypeDeduction.visit_FunctionDefinition\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 258\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvisit_FunctionDefinition\u001b[39m(\u001b[38;5;28mself\u001b[39m, node: foast\u001b[38;5;241m.\u001b[39mFunctionDefinition, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m 259\u001b[0m new_params \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mparams, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 260\u001b[0m new_body \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbody\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 261\u001b[0m new_closure_vars \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mclosure_vars, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m 262\u001b[0m return_type \u001b[38;5;241m=\u001b[39m deduce_stmt_return_type(new_body)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:161\u001b[0m, in \u001b[0;36mNodeTranslator.generic_visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 157\u001b[0m memo \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__memo__\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m 159\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, concepts\u001b[38;5;241m.\u001b[39mNode):\n\u001b[1;32m 160\u001b[0m new_node \u001b[38;5;241m=\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m( \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[0;32m--> 161\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m{\n\u001b[1;32m 162\u001b[0m name: new_child\n\u001b[1;32m 163\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m name, child \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39miter_children_items()\n\u001b[1;32m 164\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (new_child \u001b[38;5;241m:=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(child, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)) \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m NOTHING\n\u001b[1;32m 165\u001b[0m },\n\u001b[1;32m 166\u001b[0m )\n\u001b[1;32m 167\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mPRESERVED_ANNEX_ATTRS \u001b[38;5;129;01mand\u001b[39;00m (old_annex \u001b[38;5;241m:=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(node, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__node_annex__\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)):\n\u001b[1;32m 168\u001b[0m \u001b[38;5;66;03m# access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter\u001b[39;00m\n\u001b[1;32m 169\u001b[0m new_annex_dict \u001b[38;5;241m=\u001b[39m new_node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__dict__\u001b[39m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:164\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 157\u001b[0m memo \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__memo__\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m 159\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, concepts\u001b[38;5;241m.\u001b[39mNode):\n\u001b[1;32m 160\u001b[0m new_node \u001b[38;5;241m=\u001b[39m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m( \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[1;32m 161\u001b[0m \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m{\n\u001b[1;32m 162\u001b[0m name: new_child\n\u001b[1;32m 163\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m name, child \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39miter_children_items()\n\u001b[0;32m--> 164\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (new_child \u001b[38;5;241m:=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchild\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m) \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m NOTHING\n\u001b[1;32m 165\u001b[0m },\n\u001b[1;32m 166\u001b[0m )\n\u001b[1;32m 167\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mPRESERVED_ANNEX_ATTRS \u001b[38;5;129;01mand\u001b[39;00m (old_annex \u001b[38;5;241m:=\u001b[39m \u001b[38;5;28mgetattr\u001b[39m(node, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__node_annex__\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m)):\n\u001b[1;32m 168\u001b[0m \u001b[38;5;66;03m# access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter\u001b[39;00m\n\u001b[1;32m 169\u001b[0m new_annex_dict \u001b[38;5;241m=\u001b[39m new_node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__dict__\u001b[39m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:181\u001b[0m, in \u001b[0;36mNodeTranslator.generic_visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m new_node\n\u001b[1;32m 177\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mlist\u001b[39m, \u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mset\u001b[39m, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mSet)) \u001b[38;5;129;01mor\u001b[39;00m (\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28misinstance\u001b[39m(node, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mSequence) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mbytes\u001b[39m))\n\u001b[1;32m 179\u001b[0m ):\n\u001b[1;32m 180\u001b[0m \u001b[38;5;66;03m# Sequence or set: create a new container instance with the new values\u001b[39;00m\n\u001b[0;32m--> 181\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__class__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# type: ignore\u001b[39;49;00m\n\u001b[1;32m 182\u001b[0m \u001b[43m \u001b[49m\u001b[43mnew_child\u001b[49m\n\u001b[1;32m 183\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mchild\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtrees\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43miter_children_values\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 184\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[43mnew_child\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m:=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchild\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mis\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mNOTHING\u001b[49m\n\u001b[1;32m 185\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 187\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mdict\u001b[39m, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mMapping)):\n\u001b[1;32m 188\u001b[0m \u001b[38;5;66;03m# Mapping: create a new mapping instance with the new values\u001b[39;00m\n\u001b[1;32m 189\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m( \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[1;32m 190\u001b[0m {\n\u001b[1;32m 191\u001b[0m name: new_child\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 194\u001b[0m }\n\u001b[1;32m 195\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:184\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m new_node\n\u001b[1;32m 177\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mlist\u001b[39m, \u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mset\u001b[39m, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mSet)) \u001b[38;5;129;01mor\u001b[39;00m (\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28misinstance\u001b[39m(node, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mSequence) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mbytes\u001b[39m))\n\u001b[1;32m 179\u001b[0m ):\n\u001b[1;32m 180\u001b[0m \u001b[38;5;66;03m# Sequence or set: create a new container instance with the new values\u001b[39;00m\n\u001b[1;32m 181\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m( \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[1;32m 182\u001b[0m new_child\n\u001b[1;32m 183\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m child \u001b[38;5;129;01min\u001b[39;00m trees\u001b[38;5;241m.\u001b[39miter_children_values(node)\n\u001b[0;32m--> 184\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (new_child \u001b[38;5;241m:=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchild\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m) \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m NOTHING\n\u001b[1;32m 185\u001b[0m )\n\u001b[1;32m 187\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node, (\u001b[38;5;28mdict\u001b[39m, collections\u001b[38;5;241m.\u001b[39mabc\u001b[38;5;241m.\u001b[39mMapping)):\n\u001b[1;32m 188\u001b[0m \u001b[38;5;66;03m# Mapping: create a new mapping instance with the new values\u001b[39;00m\n\u001b[1;32m 189\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m node\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m( \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[1;32m 190\u001b[0m {\n\u001b[1;32m 191\u001b[0m name: new_child\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 194\u001b[0m }\n\u001b[1;32m 195\u001b[0m )\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/foast_passes/type_deduction.py:370\u001b[0m, in \u001b[0;36mFieldOperatorTypeDeduction.visit_TupleTargetAssign\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 366\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvisit_TupleTargetAssign\u001b[39m(\n\u001b[1;32m 367\u001b[0m \u001b[38;5;28mself\u001b[39m, node: foast\u001b[38;5;241m.\u001b[39mTupleTargetAssign, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m 368\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m foast\u001b[38;5;241m.\u001b[39mTupleTargetAssign:\n\u001b[1;32m 369\u001b[0m TargetType \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlist\u001b[39m[foast\u001b[38;5;241m.\u001b[39mStarred \u001b[38;5;241m|\u001b[39m foast\u001b[38;5;241m.\u001b[39mSymbol]\n\u001b[0;32m--> 370\u001b[0m values \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 372\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(values\u001b[38;5;241m.\u001b[39mtype, ts\u001b[38;5;241m.\u001b[39mTupleType):\n\u001b[1;32m 373\u001b[0m num_elts: \u001b[38;5;28mint\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(values\u001b[38;5;241m.\u001b[39mtype\u001b[38;5;241m.\u001b[39mtypes)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/foast_passes/type_deduction.py:670\u001b[0m, in \u001b[0;36mFieldOperatorTypeDeduction.visit_Call\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 669\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvisit_Call\u001b[39m(\u001b[38;5;28mself\u001b[39m, node: foast\u001b[38;5;241m.\u001b[39mCall, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m foast\u001b[38;5;241m.\u001b[39mCall:\n\u001b[0;32m--> 670\u001b[0m new_func \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 671\u001b[0m new_args \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m 672\u001b[0m new_kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvisit(node\u001b[38;5;241m.\u001b[39mkwargs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/traits.py:168\u001b[0m, in \u001b[0;36mVisitorWithSymbolTableTrait.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope \u001b[38;5;241m:=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m node\u001b[38;5;241m.\u001b[39mannex):\n\u001b[1;32m 166\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mnew_child(node\u001b[38;5;241m.\u001b[39mannex\u001b[38;5;241m.\u001b[39msymtable)\n\u001b[0;32m--> 168\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvisit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m new_scope:\n\u001b[1;32m 171\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mparents\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/eve/visitors.py:121\u001b[0m, in \u001b[0;36mNodeVisitor.visit\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 118\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node_class \u001b[38;5;129;01mis\u001b[39;00m concepts\u001b[38;5;241m.\u001b[39mNode:\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m--> 121\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvisitor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/foast_passes/type_deduction.py:354\u001b[0m, in \u001b[0;36mFieldOperatorTypeDeduction.visit_Name\u001b[0;34m(self, node, **kwargs)\u001b[0m\n\u001b[1;32m 352\u001b[0m symtable \u001b[38;5;241m=\u001b[39m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msymtable\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m 353\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m node\u001b[38;5;241m.\u001b[39mid \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m symtable \u001b[38;5;129;01mor\u001b[39;00m symtable[node\u001b[38;5;241m.\u001b[39mid]\u001b[38;5;241m.\u001b[39mtype \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 354\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m errors\u001b[38;5;241m.\u001b[39mDSLError(node\u001b[38;5;241m.\u001b[39mlocation, \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUndeclared symbol \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnode\u001b[38;5;241m.\u001b[39mid\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 356\u001b[0m symbol \u001b[38;5;241m=\u001b[39m symtable[node\u001b[38;5;241m.\u001b[39mid]\n\u001b[1;32m 357\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m foast\u001b[38;5;241m.\u001b[39mName(\u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39mnode\u001b[38;5;241m.\u001b[39mid, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39msymbol\u001b[38;5;241m.\u001b[39mtype, location\u001b[38;5;241m=\u001b[39mnode\u001b[38;5;241m.\u001b[39mlocation)\n", + "\u001b[0;31mDSLError\u001b[0m: Undeclared symbol '_graupel_toy_scan'.\n File \"/private/var/folders/2b/_2y31vzs4sl_7rngh2yghbpw0000gn/T/ipykernel_56865/2647114660.py\", line 8" + ] + } + ], + "source": [ + "@gtx.field_operator(backend=backend)\n", + "def graupel_toy_scan(\n", + " qc: gtx.Field[Dims[C, K], float], qr: gtx.Field[Dims[C, K], float]\n", + ") -> tuple[\n", + " gtx.Field[Dims[C, K], float],\n", + " gtx.Field[Dims[C, K], float]\n", + "]:\n", + " qc, qr, _ = _graupel_toy_scan(qc, qr)\n", + "\n", + " return qc, qr" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "4e0dc8d2", + "metadata": {}, + "outputs": [], + "source": [ + "def test_scan_operator():\n", + " cell_k_domain = gtx.domain({C: n_cells, K: n_levels})\n", + " \n", + " qc = random_field(cell_k_domain, allocator=backend)\n", + " qr = random_field(cell_k_domain, allocator=backend)\n", + "\n", + " qc_new = gtx.zeros(cell_k_domain, allocator=backend)\n", + " qr_new = gtx.zeros(cell_k_domain, allocator=backend)\n", + "\n", + " # Initialize Numpy fields from GT4Py fields\n", + " qc_numpy = qc.asnumpy().copy()\n", + " qr_numpy = qr.asnumpy().copy()\n", + "\n", + " # Execute the Numpy version of scheme\n", + " toy_microphysics_numpy(qc_numpy, qr_numpy)\n", + "\n", + " # Execute the GT4Py version of scheme\n", + " graupel_toy_scan(qc, qr, out=(qc_new, qr_new), offset_provider={})\n", + "\n", + " # Compare results\n", + " assert np.allclose(qc_new.asnumpy(), qc_numpy)\n", + " assert np.allclose(qr_new.asnumpy(), qr_numpy)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "a76a6be7", + "metadata": {}, + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'graupel_toy_scan' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[6], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_scan_operator\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[5], line 18\u001b[0m, in \u001b[0;36mtest_scan_operator\u001b[0;34m()\u001b[0m\n\u001b[1;32m 15\u001b[0m toy_microphysics_numpy(qc_numpy, qr_numpy)\n\u001b[1;32m 17\u001b[0m \u001b[38;5;66;03m# Execute the GT4Py version of scheme\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[43mgraupel_toy_scan\u001b[49m(qc, qr, out\u001b[38;5;241m=\u001b[39m(qc_new, qr_new), offset_provider\u001b[38;5;241m=\u001b[39m{})\n\u001b[1;32m 20\u001b[0m \u001b[38;5;66;03m# Compare results\u001b[39;00m\n\u001b[1;32m 21\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(qc_new\u001b[38;5;241m.\u001b[39masnumpy(), qc_numpy)\n", + "\u001b[0;31mNameError\u001b[0m: name 'graupel_toy_scan' is not defined" + ] + } + ], + "source": [ + "test_scan_operator()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db590d8a", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/7_scan_operator_solutions.ipynb b/docs/user/next/workshop/exercises/7_scan_operator_solutions.ipynb new file mode 100644 index 0000000000..e7831f3687 --- /dev/null +++ b/docs/user/next/workshop/exercises/7_scan_operator_solutions.ipynb @@ -0,0 +1,299 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9ba87bfb", + "metadata": {}, + "source": [ + "## Scan operator" + ] + }, + { + "cell_type": "markdown", + "id": "2ee9d989", + "metadata": {}, + "source": [ + "The unique feature of this operator is that it provides the return state of the previous iteration as its first argument (i.e., the result from the previous grid point). In other words, all the arguments of the current `return` will be available (as a tuple) in the next iteration from the first argument of the defined function. \n", + "\n", + "Example: A FORTRAN pseudocode for integrating a moisture variable (e.g., cloud water or water vapour) over a column could look as follows:\n", + "\n", + "\n", + "```FORTRAN\n", + "SUBROUTINE column_integral( var_in, rho, dz, var_out, ie, je, ke )\n", + " ! Return the column integral of a moist species.\n", + " INTEGER, INTENT (IN) :: &\n", + " ie, je, ke ! array dimensions of the I/O-fields (horizontal, horizontal, vertical)\n", + "\n", + " REAL (KIND=wp), INTENT (OUT) :: &\n", + " q_colsum (ie,je) ! Vertically-integrated mass of water species\n", + "\n", + " REAL (KIND=wp), INTENT (IN) :: &\n", + " rho (ie,je,ke), & \n", + " dz (ie,je,ke), & ! height of model half levels\n", + " var_in (ie,je,ke) ! humidity mass concentration at time-level nnow\n", + " \n", + " !$acc parallel present( iq ) if (lzacc)\n", + " !$acc loop gang\n", + " DO j=1,je\n", + " !$acc loop vector\n", + " DO i=1,ie\n", + " q_sum(i,j) = 0.0\n", + " END DO\n", + " END DO\n", + " !$acc end parallel\n", + " \n", + " \n", + " !$acc parallel present( iq, rho, hhl, q ) if (lzacc)\n", + " DO k = 1, ke ! Vertical loop\n", + " !$acc loop gang\n", + " DO j=1,je\n", + " !$acc loop vector\n", + " DO i=1,ie\n", + " q_colsum(i,j) = q_colsum(i,j) + var_in(i,j,k) * rho(i,j,k)* dz(i,j,k)\n", + " END DO\n", + " END DO\n", + " END DO\n", + " !$acc end parallel\n", + "END SUBROUTINE column_integral\n", + "```\n", + "\n", + "Where:\n", + "- `var_in` is the 3D variable that will be summed up\n", + "- `q_colsum` is the resulting 2D variable\n", + "- `rho` the air density\n", + "- `dz`the thickness of the vertical layers\n", + "\n", + "In the first loop nest, `column_sum` is set to zero for all grid columns. The vertical dependency enters on the RHS of the second loop nest `q_colsum(i,j) = q_colsum(i,j) + ...`\n", + "\n", + "Using the `scan_operator` this operation would be written like this:\n", + "\n", + "```python\n", + "@scan_operator(axis=KDim, forward=True, init=0.0)\n", + "def column_integral(float: state, float: var, float: rho, float: dz)\n", + " \"\"\"Return the column integral of a moist species.\"\"\"\n", + " return var * rho * dz + state\n", + "```\n", + "\n", + "Here the vertical dependency is expressed by the first function argument (`state`). This argument carries the return from the previous k-level and does not need to be specified when the function is called (similar to the `self` argument of Python classes). The argument is intialized to `init=0.0` in the function decorator (first loop nest above) and the dimension of the integral is specified with `axis=KDim`.\n", + "\n", + "\n", + "```python\n", + "q_colsum = column_integral(qv, rho, dz)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "c9e31bff", + "metadata": {}, + "source": [ + "#### Exercise: port a toy cloud microphysics scheme from python/numpy using the template of a `scan_operator` below" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "bd2fd309", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx\n", + "\n", + "backend = None\n", + "# backend = gtfn_cpu\n", + "# backend = gtfn_gpu" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "74338168", + "metadata": {}, + "outputs": [], + "source": [ + "def toy_microphysics_numpy(qc, qr, autoconversion_rate=0.1, sedimentaion_constant=0.05):\n", + " \"\"\"A toy model of a microphysics scheme contaning autoconversion and scavenging\"\"\"\n", + "\n", + " sedimentation_flux = 0.0\n", + "\n", + " for cell, k in np.ndindex(qc.shape):\n", + " # Autoconversion: Cloud Drops -> Rain Drops\n", + " autoconversion_tendency = qc[cell, k] * autoconversion_rate\n", + "\n", + " qc[cell, k] -= autoconversion_tendency\n", + " qr[cell, k] += autoconversion_tendency\n", + "\n", + " ## Apply sedimentation flux from level above\n", + " qr[cell, k] += sedimentation_flux\n", + "\n", + " ## Remove mass due to sedimentation flux from the current cell\n", + " qr[cell, k] -= sedimentation_flux" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "69bf6022", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.scan_operator(axis=K, forward=True, init=(0.0, 0.0, 0.0))\n", + "def _graupel_toy_scan(\n", + " state: tuple[float, float, float], qc_in: float, qr_in: float\n", + ") -> tuple[float, float, float]:\n", + " autoconversion_rate = 0.1\n", + " sedimentaion_constant = 0.05\n", + "\n", + " # unpack state of previous iteration\n", + " _, _, sedimentation_flux = state\n", + "\n", + " # Autoconversion: Cloud Drops -> Rain Drops\n", + " autoconv_t = qc_in * autoconversion_rate\n", + " qc = qc_in - autoconv_t\n", + " qr = qr_in + autoconv_t\n", + "\n", + " ## Add sedimentation flux from level above\n", + " qr = qr + sedimentation_flux\n", + "\n", + " # Remove mass due to sedimentation flux\n", + " qr = qr - sedimentation_flux\n", + "\n", + " return qc, qr, sedimentation_flux" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "0de41cb8", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator(backend=backend)\n", + "def graupel_toy_scan(\n", + " qc: gtx.Field[Dims[C, K], float], qr: gtx.Field[Dims[C, K], float]\n", + ") -> tuple[\n", + " gtx.Field[Dims[C, K], float],\n", + " gtx.Field[Dims[C, K], float]\n", + "]:\n", + " qc, qr, _ = _graupel_toy_scan(qc, qr)\n", + "\n", + " return qc, qr" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "4e0dc8d2", + "metadata": {}, + "outputs": [], + "source": [ + "def test_scan_operator():\n", + " cell_k_domain = gtx.domain({C: n_cells, K: n_levels})\n", + " \n", + " qc = random_field(cell_k_domain, allocator=backend)\n", + " qr = random_field(cell_k_domain, allocator=backend)\n", + "\n", + " qc_new = gtx.zeros(cell_k_domain, allocator=backend)\n", + " qr_new = gtx.zeros(cell_k_domain, allocator=backend)\n", + "\n", + " # Initialize Numpy fields from GT4Py fields\n", + " qc_numpy = qc.asnumpy().copy()\n", + " qr_numpy = qr.asnumpy().copy()\n", + "\n", + " # Execute the Numpy version of scheme\n", + " toy_microphysics_numpy(qc_numpy, qr_numpy)\n", + "\n", + " # Execute the GT4Py version of scheme\n", + " graupel_toy_scan(qc, qr, out=(qc_new, qr_new), offset_provider={})\n", + "\n", + " # Compare results\n", + " assert np.allclose(qc_new.asnumpy(), qc_numpy)\n", + " assert np.allclose(qr_new.asnumpy(), qr_numpy)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "a76a6be7", + "metadata": {}, + "outputs": [ + { + "ename": "AssertionError", + "evalue": "", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mAssertionError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[7], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtest_scan_operator\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest successful\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n", + "Cell \u001b[0;32mIn[6], line 18\u001b[0m, in \u001b[0;36mtest_scan_operator\u001b[0;34m()\u001b[0m\n\u001b[1;32m 15\u001b[0m toy_microphysics_numpy(qc_numpy, qr_numpy)\n\u001b[1;32m 17\u001b[0m \u001b[38;5;66;03m# Execute the GT4Py version of scheme\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[43mgraupel_toy_scan\u001b[49m\u001b[43m(\u001b[49m\u001b[43mqc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mqc_new\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqr_new\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moffset_provider\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m{\u001b[49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 20\u001b[0m \u001b[38;5;66;03m# Compare results\u001b[39;00m\n\u001b[1;32m 21\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m np\u001b[38;5;241m.\u001b[39mallclose(qc_new\u001b[38;5;241m.\u001b[39masnumpy(), qc_numpy)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:757\u001b[0m, in \u001b[0;36mFieldOperator.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 755\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 756\u001b[0m op \u001b[38;5;241m=\u001b[39m embedded_operators\u001b[38;5;241m.\u001b[39mEmbeddedOperator(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdefinition)\n\u001b[0;32m--> 757\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43membedded_operators\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfield_operator_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43mop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/embedded/operators.py:113\u001b[0m, in \u001b[0;36mfield_operator_call\u001b[0;34m(op, args, kwargs)\u001b[0m\n\u001b[1;32m 110\u001b[0m new_context_kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mclosure_column_range\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m _get_vertical_range(out_domain)\n\u001b[1;32m 112\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m embedded_context\u001b[38;5;241m.\u001b[39mnew_context(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mnew_context_kwargs) \u001b[38;5;28;01mas\u001b[39;00m ctx:\n\u001b[0;32m--> 113\u001b[0m res \u001b[38;5;241m=\u001b[39m \u001b[43mctx\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 114\u001b[0m _tuple_assign_field(\n\u001b[1;32m 115\u001b[0m out,\n\u001b[1;32m 116\u001b[0m res,\n\u001b[1;32m 117\u001b[0m domain\u001b[38;5;241m=\u001b[39mout_domain,\n\u001b[1;32m 118\u001b[0m )\n\u001b[1;32m 119\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 120\u001b[0m \u001b[38;5;66;03m# called from other field_operator or missing `out` argument\u001b[39;00m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/embedded/operators.py:36\u001b[0m, in \u001b[0;36mEmbeddedOperator.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs: _P\u001b[38;5;241m.\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: _P\u001b[38;5;241m.\u001b[39mkwargs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m _R:\n\u001b[0;32m---> 36\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "Cell \u001b[0;32mIn[5], line 8\u001b[0m, in \u001b[0;36mgraupel_toy_scan\u001b[0;34m(qc, qr)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;129m@gtx\u001b[39m\u001b[38;5;241m.\u001b[39mfield_operator(backend\u001b[38;5;241m=\u001b[39mbackend)\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mgraupel_toy_scan\u001b[39m(\n\u001b[1;32m 3\u001b[0m qc: gtx\u001b[38;5;241m.\u001b[39mField[Dims[C, K], \u001b[38;5;28mfloat\u001b[39m], qr: gtx\u001b[38;5;241m.\u001b[39mField[Dims[C, K], \u001b[38;5;28mfloat\u001b[39m]\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 6\u001b[0m gtx\u001b[38;5;241m.\u001b[39mField[Dims[C, K], \u001b[38;5;28mfloat\u001b[39m]\n\u001b[1;32m 7\u001b[0m ]:\n\u001b[0;32m----> 8\u001b[0m qc, qr, _ \u001b[38;5;241m=\u001b[39m \u001b[43m_graupel_toy_scan\u001b[49m\u001b[43m(\u001b[49m\u001b[43mqc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mqr\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 10\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m qc, qr\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/ffront/decorator.py:757\u001b[0m, in \u001b[0;36mFieldOperator.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 755\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 756\u001b[0m op \u001b[38;5;241m=\u001b[39m embedded_operators\u001b[38;5;241m.\u001b[39mEmbeddedOperator(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdefinition)\n\u001b[0;32m--> 757\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43membedded_operators\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfield_operator_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43mop\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/embedded/operators.py:124\u001b[0m, in \u001b[0;36mfield_operator_call\u001b[0;34m(op, args, kwargs)\u001b[0m\n\u001b[1;32m 121\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moffset_provider\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m kwargs:\n\u001b[1;32m 122\u001b[0m \u001b[38;5;66;03m# assuming we wanted to call the field_operator as program, otherwise `offset_provider` would not be there\u001b[39;00m\n\u001b[1;32m 123\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m errors\u001b[38;5;241m.\u001b[39mMissingArgumentError(\u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mout\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m--> 124\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mop\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/embedded/operators.py:81\u001b[0m, in \u001b[0;36mScanOperator.__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 79\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 80\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m hpos \u001b[38;5;129;01min\u001b[39;00m embedded_common\u001b[38;5;241m.\u001b[39miterate_domain(non_scan_domain):\n\u001b[0;32m---> 81\u001b[0m \u001b[43mscan_loop\u001b[49m\u001b[43m(\u001b[49m\u001b[43mhpos\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 83\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m res\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/embedded/operators.py:71\u001b[0m, in \u001b[0;36mScanOperator.__call__..scan_loop\u001b[0;34m(hpos)\u001b[0m\n\u001b[1;32m 69\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m scan_range[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mforward \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mreversed\u001b[39m(scan_range[\u001b[38;5;241m1\u001b[39m]):\n\u001b[1;32m 70\u001b[0m pos \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m*\u001b[39mhpos, (scan_axis, k))\n\u001b[0;32m---> 71\u001b[0m new_args \u001b[38;5;241m=\u001b[39m [_tuple_at(pos, arg) \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args]\n\u001b[1;32m 72\u001b[0m new_kwargs \u001b[38;5;241m=\u001b[39m {k: _tuple_at(pos, v) \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m kwargs\u001b[38;5;241m.\u001b[39mitems()}\n\u001b[1;32m 73\u001b[0m acc \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfun(acc, \u001b[38;5;241m*\u001b[39mnew_args, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mnew_kwargs)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/embedded/operators.py:71\u001b[0m, in \u001b[0;36m\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m 69\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m scan_range[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mforward \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mreversed\u001b[39m(scan_range[\u001b[38;5;241m1\u001b[39m]):\n\u001b[1;32m 70\u001b[0m pos \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m*\u001b[39mhpos, (scan_axis, k))\n\u001b[0;32m---> 71\u001b[0m new_args \u001b[38;5;241m=\u001b[39m [\u001b[43m_tuple_at\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpos\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43marg\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args]\n\u001b[1;32m 72\u001b[0m new_kwargs \u001b[38;5;241m=\u001b[39m {k: _tuple_at(pos, v) \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m kwargs\u001b[38;5;241m.\u001b[39mitems()}\n\u001b[1;32m 73\u001b[0m acc \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfun(acc, \u001b[38;5;241m*\u001b[39mnew_args, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mnew_kwargs)\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/embedded/operators.py:198\u001b[0m, in \u001b[0;36m_tuple_at\u001b[0;34m(pos, field)\u001b[0m\n\u001b[1;32m 195\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m core_defs\u001b[38;5;241m.\u001b[39mis_scalar_type(res)\n\u001b[1;32m 196\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m res\n\u001b[0;32m--> 198\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfield\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/utils.py:90\u001b[0m, in \u001b[0;36mtree_map..impl\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m 87\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mall\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(arg, \u001b[38;5;28mtuple\u001b[39m) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(args[\u001b[38;5;241m0\u001b[39m]) \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mlen\u001b[39m(arg) \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m args)\n\u001b[1;32m 88\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtuple\u001b[39m(impl(\u001b[38;5;241m*\u001b[39marg) \u001b[38;5;28;01mfor\u001b[39;00m arg \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39margs))\n\u001b[0;32m---> 90\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 91\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcast\u001b[49m\u001b[43m(\u001b[49m\u001b[43m_P\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 92\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/gt4py-update_types/gt4py/src/gt4py/next/embedded/operators.py:195\u001b[0m, in \u001b[0;36m_tuple_at..impl\u001b[0;34m(field)\u001b[0m\n\u001b[1;32m 192\u001b[0m \u001b[38;5;129m@utils\u001b[39m\u001b[38;5;241m.\u001b[39mtree_map\n\u001b[1;32m 193\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mimpl\u001b[39m(field: common\u001b[38;5;241m.\u001b[39mField \u001b[38;5;241m|\u001b[39m core_defs\u001b[38;5;241m.\u001b[39mScalar) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m core_defs\u001b[38;5;241m.\u001b[39mScalar:\n\u001b[1;32m 194\u001b[0m res \u001b[38;5;241m=\u001b[39m field[pos] \u001b[38;5;28;01mif\u001b[39;00m common\u001b[38;5;241m.\u001b[39mis_field(field) \u001b[38;5;28;01melse\u001b[39;00m field\n\u001b[0;32m--> 195\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m core_defs\u001b[38;5;241m.\u001b[39mis_scalar_type(res)\n\u001b[1;32m 196\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m res\n", + "\u001b[0;31mAssertionError\u001b[0m: " + ] + } + ], + "source": [ + "test_scan_operator()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db590d8a", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/8_diffusion_exercise_solution.ipynb b/docs/user/next/workshop/exercises/8_diffusion_exercise_solution.ipynb new file mode 100644 index 0000000000..e12c51973a --- /dev/null +++ b/docs/user/next/workshop/exercises/8_diffusion_exercise_solution.ipynb @@ -0,0 +1,229 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 17, + "id": "037b7de4-e92c-439b-81fb-9be718c62f17", + "metadata": {}, + "outputs": [], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": 119, + "id": "f36ace26", + "metadata": {}, + "outputs": [], + "source": [ + "def diffusion_step_numpy(\n", + " e2c2v: np.array,\n", + " v2e: np.array,\n", + " TE: np.array,\n", + " TE_t: np.array,\n", + " inv_primal_edge_length: np.array,\n", + " inv_vert_vert_length: np.array,\n", + " nnbhV: np.array,\n", + " boundary_edge: np.array,\n", + " kappa: float,\n", + " dt: float,\n", + ") -> tuple[np.array, np.array]:\n", + "\n", + " # initialize\n", + " TEinit = TE\n", + " inv_primal_edge_length = inv_primal_edge_length[:, np.newaxis]\n", + " inv_vert_vert_length = inv_vert_vert_length[:, np.newaxis]\n", + "\n", + " # predict\n", + " TE = TEinit + 0.5*dt*TE_t\n", + "\n", + " # interpolate temperature from edges to vertices\n", + " TV = np.sum(TE[v2e], axis=1) / nnbhV\n", + "\n", + " # compute nabla2 using the finite differences\n", + " TEnabla2 = np.sum(\n", + " TV[e2c2v] * inv_primal_edge_length ** 2\n", + " + TV[e2c2v] * inv_vert_vert_length ** 2,\n", + " axis=1,\n", + " )\n", + "\n", + " TEnabla2 = TEnabla2 - (\n", + " (2.0 * TE * inv_primal_edge_length ** 2)\n", + " + (2.0 * TE * inv_vert_vert_length ** 2)\n", + " )\n", + "\n", + " # build ODEs\n", + " TE_t = np.where(\n", + " boundary_edge,\n", + " 0.,\n", + " kappa*TEnabla2,\n", + " )\n", + "\n", + " # correct\n", + " TE = TEinit + dt*TE_t\n", + " return TE_t, TE" + ] + }, + { + "cell_type": "code", + "execution_count": 120, + "id": "c59c4ffb-b31c-4193-bead-e213c23ba858", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def diffusion_step(\n", + " TE: gtx.Field[Dims[E], float],\n", + " TE_t: gtx.Field[Dims[E], float],\n", + " inv_primal_edge_length: gtx.Field[Dims[E], float],\n", + " inv_vert_vert_length: gtx.Field[Dims[E], float],\n", + " nnbhV: gtx.Field[Dims[V], float],\n", + " boundary_edge: gtx.Field[Dims[E], bool],\n", + " kappa: float,\n", + " dt: float,\n", + ") -> gtx.tuple[\n", + " gtx.Field[Dims[E], float],\n", + " gtx.Field[Dims[E], float],\n", + "]:\n", + "\n", + " # initialize\n", + " TEinit = TE\n", + "\n", + " # predict\n", + " TE = TEinit + 0.5*dt*TE_t\n", + "\n", + " # interpolate temperature from edges to vertices\n", + " TV = neighbor_sum(TE(V2E), axis=V2EDim) / nnbhV\n", + "\n", + " # compute nabla2 using the finite differences\n", + " TEnabla2 = neighbor_sum(\n", + " (TV(E2C2V) * inv_primal_edge_length ** 2\n", + " + TV(E2C2V) * inv_vert_vert_length ** 2),\n", + " axis=E2C2VDim\n", + " )\n", + "\n", + " TEnabla2 = TEnabla2 - (\n", + " (2.0 * TE * inv_primal_edge_length ** 2)\n", + " + (2.0 * TE * inv_vert_vert_length ** 2)\n", + " )\n", + "\n", + " # build ODEs\n", + " TE_t = where(\n", + " boundary_edge,\n", + " 0.,\n", + " kappa*TEnabla2,\n", + " )\n", + "\n", + " # correct\n", + " TE = TEinit + dt*TE_t\n", + " \n", + " return TE_t, TE" + ] + }, + { + "cell_type": "code", + "execution_count": 127, + "id": "f9cfc097", + "metadata": {}, + "outputs": [], + "source": [ + "def test_diffusion_step():\n", + " backend = None\n", + " # backend = gtfn_cpu\n", + " # backend = gtfn_gpu\n", + " \n", + " cell_domain = gtx.domain({C: n_cells})\n", + " edge_domain = gtx.domain({E: n_edges})\n", + " vertex_domain = gtx.domain({V: n_vertices})\n", + " \n", + " u = random_field(edge_domain, allocator=backend)\n", + " v = random_field(edge_domain, allocator=backend)\n", + " nx = random_field(edge_domain, allocator=backend)\n", + " ny = random_field(edge_domain, allocator=backend)\n", + " L = random_field(edge_domain, allocator=backend)\n", + " dualL = random_field(vertex_domain, allocator=backend)\n", + " divergence_gt4py_1 = gtx.zeros(edge_domain, allocator=backend)\n", + " divergence_gt4py_2 = gtx.zeros(edge_domain, allocator=backend)\n", + " kappa = 1.0\n", + " dt = 1.0\n", + "\n", + " divergence_ref_1, divergence_ref_2 = diffusion_step_numpy( \n", + " e2c2v_table,\n", + " v2e_table,\n", + " u.asnumpy(),\n", + " v.asnumpy(),\n", + " nx.asnumpy(),\n", + " ny.asnumpy(),\n", + " dualL.asnumpy(),\n", + " L.asnumpy(),\n", + " kappa,\n", + " dt\n", + " )\n", + "\n", + " e2c2v_connectivity = gtx.NeighborTableOffsetProvider(e2c2v_table, E, V, 4, has_skip_values=False)\n", + " v2e_connectivity = gtx.NeighborTableOffsetProvider(v2e_table, V, E, 6, has_skip_values=False)\n", + "\n", + " diffusion_step(\n", + " u, v, nx, ny, dualL, L, kappa, dt, out=(divergence_gt4py_1, divergence_gt4py_2), offset_provider = {E2C2V.value: e2c2v_connectivity, V2E.value: v2e_connectivity}\n", + " )\n", + " \n", + " assert np.allclose(divergence_gt4py_1.asnumpy(), divergence_ref_1)\n", + " assert np.allclose(divergence_gt4py_2.asnumpy(), divergence_ref_2)" + ] + }, + { + "cell_type": "code", + "execution_count": 128, + "id": "d4079375", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Test successful\n" + ] + } + ], + "source": [ + "test_diffusion_step()\n", + "print(\"Test successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3456b12c", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/exercises/helpers.py b/docs/user/next/workshop/exercises/helpers.py new file mode 100644 index 0000000000..7b74711977 --- /dev/null +++ b/docs/user/next/workshop/exercises/helpers.py @@ -0,0 +1,384 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2023, ETH Zurich +# All rights reserved. +# +# This file is part of the GT4Py project and the GridTools framework. +# GT4Py is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +import gt4py.next as gtx +from gt4py.next.iterator.embedded import MutableLocatedField +from gt4py.next import neighbor_sum, where, Dims +from gt4py.next import Dimension, DimensionKind, FieldOffset +from gt4py.next.program_processors.runners import roundtrip +from gt4py.next.program_processors.runners.gtfn import ( + run_gtfn_cached as gtfn_cpu, + run_gtfn_gpu as gtfn_gpu, +) + + +def random_mask( + sizes, + *dims, + dtype=None, +) -> MutableLocatedField: + arr = np.full(shape=sizes, fill_value=False).flatten() + arr[: int(arr.size * 0.5)] = True + np.random.shuffle(arr) + arr = np.reshape(arr, newshape=sizes) + if dtype: + arr = arr.astype(dtype) + return gtx.as_field([*dims], (arr)) + + +def random_field( + domain: gtx.Domain, low: float = -1.0, high: float = 1.0, *, allocator=None +) -> MutableLocatedField: + return gtx.as_field( + domain, + np.random.default_rng().uniform(low=low, high=high, size=domain.shape), + allocator=allocator, + ) + + +def random_sign(domain: gtx.Domain, allocator=None, dtype=float) -> MutableLocatedField: + return gtx.as_field( + domain, + np.asarray(np.random.randint(0, high=2, size=domain.shape) * 2 - 1, dtype=dtype), + allocator=allocator, + ) + + +def ripple_field(domain: gtx.Domain, *, allocator=None) -> MutableLocatedField: + assert domain.ndim == 2 + nx, ny = domain.shape + x = np.linspace(0, 1, nx) + y = np.linspace(0, 1, ny) + xx, yy = np.meshgrid(x, y) + data = ( + 5.0 + + 8.0 * (2.0 + np.cos(np.pi * (xx + 1.5 * yy)) + np.sin(2 * np.pi * (xx + 1.5 * yy))) / 4.0 + ) + + return gtx.as_field(domain, data, allocator=allocator) + + +# For simplicity we use a triangulated donut in the horizontal. + +# 0v---0e-- 1v---3e-- 2v---6e-- 0v +# | \ 0c | \ 1c | \2c +# | \1e | \4e | \7e +# |2e \ |5e \ |8e \ +# | 3c \ | 4c \ | 5c\ +# 3v---9e-- 4v--12e-- 5v--15e-- 3v +# | \ 6c | \ 7c | \ 8c +# | \10e | \13e | \16e +# |11e \ |14e \ |17e \ +# | 9c \ | 10c \ | 11c \ +# 6v--18e-- 7v--21e-- 8v--24e-- 6v +# | \12c | \ 13c | \ 14c +# | \19e | \22e | \25e +# |20e \ |23e \ |26e \ +# | 15c \ | 16c \ | 17c \ +# 0v 1v 2v 0v + + +n_edges = 27 +n_vertices = 9 +n_cells = 18 +n_levels = 10 + + +e2c2v_table = np.asarray([ + [0, 1, 4, 6], # 0 + [0, 4, 1, 3], # 1 + [0, 3, 4, 2], # 2 + [1, 2, 5, 7], # 3 + [1, 5, 2, 4], # 4 + [1, 4, 5, 0], # 5 + [2, 0, 3, 8], # 6 + [2, 3, 5, 0], # 7 + [2, 5, 1, 3], # 8 + [3, 4, 0, 7], # 9 + [3, 7, 4, 6], # 10 + [3, 6, 7, 5], # 11 + [4, 5, 8, 1], # 12 + [4, 8, 7, 5], # 13 + [4, 7, 3, 8], # 14 + [5, 3, 6, 2], # 15 + [6, 5, 3, 8], # 16 + [8, 5, 6, 4], # 17 + [6, 7, 3, 1], # 18 + [6, 1, 7, 0], # 19 + [6, 0, 1, 8], # 20 + [7, 8, 2, 4], # 21 + [7, 2, 8, 1], # 22 + [7, 1, 2, 6], # 23 + [8, 6, 0, 5], # 24 + [8, 0, 6, 2], # 25 + [8, 2, 0, 6], # 26 +]) + +e2c_table = np.asarray([ + [0, 15], + [0, 3], + [3, 2], + [1, 16], + [1, 4], + [0, 4], + [2, 17], + [2, 5], + [1, 5], + [3, 6], + [6, 9], + [9, 8], + [4, 7], + [7, 10], + [6, 10], + [5, 8], + [8, 11], + [7, 11], + [9, 12], + [12, 15], + [15, 14], + [10, 13], + [13, 16], + [12, 16], + [11, 14], + [14, 17], + [13, 17], +]) + +e2v_table = np.asarray([ + [0, 1], + [0, 4], + [0, 3], + [1, 2], + [1, 5], + [1, 4], + [2, 0], + [2, 3], + [2, 5], + [3, 4], + [3, 7], + [3, 6], + [4, 5], + [4, 8], + [4, 7], + [5, 3], + [5, 6], + [5, 8], + [6, 7], + [6, 1], + [6, 0], + [7, 8], + [7, 2], + [7, 1], + [8, 6], + [8, 0], + [8, 2], +]) + +e2c2e_table = np.asarray([ + [1, 5, 19, 20], + [0, 5, 2, 9], + [1, 9, 6, 7], + [4, 8, 22, 23], + [3, 8, 5, 12], + [0, 1, 4, 12], + [7, 2, 25, 26], + [6, 2, 8, 15], + [3, 4, 7, 15], + [1, 2, 10, 14], + [9, 14, 11, 18], + [10, 18, 15, 16], + [4, 5, 13, 17], + [12, 17, 14, 21], + [9, 10, 13, 21], + [7, 8, 16, 11], + [15, 11, 17, 24], + [12, 13, 16, 24], + [10, 11, 19, 23], + [18, 23, 20, 0], + [19, 0, 24, 25], + [13, 14, 22, 26], + [21, 26, 23, 3], + [18, 19, 22, 3], + [16, 17, 25, 20], + [24, 20, 26, 6], + [25, 6, 21, 22], +]) + +e2c2eO_table = np.asarray([ + [0, 1, 5, 19, 20], + [0, 1, 5, 2, 9], + [1, 2, 9, 6, 7], + [3, 4, 8, 22, 23], + [3, 4, 8, 5, 12], + [0, 1, 5, 4, 12], + [6, 7, 2, 25, 26], + [6, 7, 2, 8, 15], + [3, 4, 8, 7, 15], + [1, 2, 9, 10, 14], + [9, 10, 14, 11, 18], + [10, 11, 18, 15, 16], + [4, 5, 12, 13, 17], + [12, 13, 17, 14, 21], + [9, 10, 14, 13, 21], + [7, 8, 15, 16, 11], + [15, 16, 11, 17, 24], + [12, 13, 17, 16, 24], + [10, 11, 18, 19, 23], + [18, 19, 23, 20, 0], + [19, 20, 0, 24, 25], + [13, 14, 21, 22, 26], + [21, 22, 26, 23, 3], + [18, 19, 23, 22, 3], + [16, 17, 24, 25, 20], + [24, 25, 20, 26, 6], + [25, 26, 6, 21, 22], +]) + +c2e_table = np.asarray([ + [0, 1, 5], # cell 0 + [3, 4, 8], # cell 1 + [6, 7, 2], # cell 2 + [1, 2, 9], # cell 3 + [4, 5, 12], # cell 4 + [7, 8, 15], # cell 5 + [9, 10, 14], # cell 6 + [12, 13, 17], # cell 7 + [15, 16, 11], # cell 8 + [10, 11, 18], # cell 9 + [13, 14, 21], # cell 10 + [16, 17, 24], # cell 11 + [18, 19, 23], # cell 12 + [21, 22, 26], # cell 13 + [24, 25, 20], # cell 14 + [19, 20, 0], # cell 15 + [22, 23, 3], # cell 16 + [25, 26, 6], # cell 17 +]) + +v2c_table = np.asarray([ + [17, 14, 3, 0, 2, 15], + [0, 4, 1, 12, 16, 15], + [1, 5, 2, 16, 13, 17], + [3, 6, 9, 5, 8, 2], + [6, 10, 7, 4, 0, 3], + [7, 11, 8, 5, 1, 4], + [9, 12, 15, 8, 11, 14], + [12, 16, 13, 10, 6, 9], + [13, 17, 14, 11, 7, 10], +]) + +v2e_table = np.asarray([ + [0, 1, 2, 6, 25, 20], + [3, 4, 5, 0, 23, 19], + [6, 7, 8, 3, 22, 26], + [9, 10, 11, 15, 7, 2], + [12, 13, 14, 9, 1, 5], + [15, 16, 17, 12, 4, 8], + [18, 19, 20, 24, 16, 11], + [21, 22, 23, 18, 10, 14], + [24, 25, 26, 21, 13, 17], +]) + +diamond_table = np.asarray([ + [0, 1, 4, 6], # 0 + [0, 4, 1, 3], + [0, 3, 4, 2], + [1, 2, 5, 7], # 3 + [1, 5, 2, 4], + [1, 4, 5, 0], + [2, 0, 3, 8], # 6 + [2, 3, 0, 5], + [2, 5, 1, 3], + [3, 4, 0, 7], # 9 + [3, 7, 4, 6], + [3, 6, 5, 7], + [4, 5, 1, 8], # 12 + [4, 8, 5, 7], + [4, 7, 3, 8], + [5, 3, 2, 6], # 15 + [5, 6, 3, 8], + [5, 8, 4, 6], + [6, 7, 3, 1], # 18 + [6, 1, 7, 0], + [6, 0, 1, 8], + [7, 8, 4, 2], # 21 + [7, 2, 8, 1], + [7, 1, 6, 2], + [8, 6, 5, 0], # 24 + [8, 0, 6, 2], + [8, 2, 7, 0], +]) + +c2e2cO_table = np.asarray([ + [15, 4, 3, 0], + [16, 5, 4, 1], + [17, 3, 5, 2], + [0, 6, 2, 3], + [1, 7, 0, 4], + [2, 8, 1, 5], + [3, 10, 9, 6], + [4, 11, 10, 7], + [5, 9, 11, 8], + [6, 12, 8, 9], + [7, 13, 6, 10], + [8, 14, 7, 11], + [9, 16, 15, 12], + [10, 17, 16, 13], + [11, 15, 17, 14], + [12, 0, 14, 15], + [13, 1, 12, 16], + [14, 2, 13, 17], +]) + +c2e2c_table = np.asarray([ + [15, 4, 3], + [16, 5, 4], + [17, 3, 5], + [0, 6, 2], + [1, 7, 0], + [2, 8, 1], + [3, 10, 9], + [4, 11, 10], + [5, 9, 11], + [6, 12, 8], + [7, 13, 6], + [8, 14, 7], + [9, 16, 15], + [10, 17, 16], + [11, 15, 17], + [12, 0, 14], + [13, 1, 12], + [14, 2, 13], +]) + + +C = Dimension("C") +V = Dimension("V") +E = Dimension("E") +K = Dimension("K", kind=gtx.DimensionKind.VERTICAL) + +C2EDim = Dimension("C2E", kind=DimensionKind.LOCAL) +C2E = FieldOffset("C2E", source=E, target=(C, C2EDim)) +V2EDim = Dimension("V2E", kind=DimensionKind.LOCAL) +V2E = FieldOffset("V2E", source=E, target=(V, V2EDim)) +E2VDim = Dimension("E2V", kind=DimensionKind.LOCAL) +E2V = FieldOffset("E2V", source=V, target=(E, E2VDim)) +E2CDim = Dimension("E2C", kind=DimensionKind.LOCAL) +E2C = FieldOffset("E2C", source=C, target=(E, E2CDim)) +E2C2VDim = Dimension("E2C2V", kind=DimensionKind.LOCAL) +E2C2V = FieldOffset("E2C2V", source=V, target=(E, E2C2VDim)) +Coff = FieldOffset("Coff", source=C, target=(C,)) # delete this? +Koff = FieldOffset("Koff", source=K, target=(K,)) diff --git a/docs/user/next/connectivity_cell_field.svg b/docs/user/next/workshop/images/connectivity_cell_field.svg similarity index 100% rename from docs/user/next/connectivity_cell_field.svg rename to docs/user/next/workshop/images/connectivity_cell_field.svg diff --git a/docs/user/next/connectivity_edge_0th_cell.svg b/docs/user/next/workshop/images/connectivity_edge_0th_cell.svg similarity index 100% rename from docs/user/next/connectivity_edge_0th_cell.svg rename to docs/user/next/workshop/images/connectivity_edge_0th_cell.svg diff --git a/docs/user/next/connectivity_edge_cell_sum.svg b/docs/user/next/workshop/images/connectivity_edge_cell_sum.svg similarity index 100% rename from docs/user/next/connectivity_edge_cell_sum.svg rename to docs/user/next/workshop/images/connectivity_edge_cell_sum.svg diff --git a/docs/user/next/connectivity_numbered_grid.svg b/docs/user/next/workshop/images/connectivity_numbered_grid.svg similarity index 100% rename from docs/user/next/connectivity_numbered_grid.svg rename to docs/user/next/workshop/images/connectivity_numbered_grid.svg diff --git a/docs/user/next/workshop/images/curl_formula.png b/docs/user/next/workshop/images/curl_formula.png new file mode 100644 index 0000000000..ff5fbc4e86 Binary files /dev/null and b/docs/user/next/workshop/images/curl_formula.png differ diff --git a/docs/user/next/workshop/images/curl_picture.png b/docs/user/next/workshop/images/curl_picture.png new file mode 100644 index 0000000000..061f565fd5 Binary files /dev/null and b/docs/user/next/workshop/images/curl_picture.png differ diff --git a/docs/user/next/workshop/images/diamond.png b/docs/user/next/workshop/images/diamond.png new file mode 100644 index 0000000000..d6cb425242 Binary files /dev/null and b/docs/user/next/workshop/images/diamond.png differ diff --git a/docs/user/next/workshop/images/directional_gradient_n_formula.png b/docs/user/next/workshop/images/directional_gradient_n_formula.png new file mode 100644 index 0000000000..e73e971645 Binary files /dev/null and b/docs/user/next/workshop/images/directional_gradient_n_formula.png differ diff --git a/docs/user/next/workshop/images/directional_gradient_n_picture.png b/docs/user/next/workshop/images/directional_gradient_n_picture.png new file mode 100644 index 0000000000..b1d39a9436 Binary files /dev/null and b/docs/user/next/workshop/images/directional_gradient_n_picture.png differ diff --git a/docs/user/next/workshop/images/directional_gradient_tau_formula.png b/docs/user/next/workshop/images/directional_gradient_tau_formula.png new file mode 100644 index 0000000000..3873aa10b9 Binary files /dev/null and b/docs/user/next/workshop/images/directional_gradient_tau_formula.png differ diff --git a/docs/user/next/workshop/images/directional_gradient_tau_picture.png b/docs/user/next/workshop/images/directional_gradient_tau_picture.png new file mode 100644 index 0000000000..7451919531 Binary files /dev/null and b/docs/user/next/workshop/images/directional_gradient_tau_picture.png differ diff --git a/docs/user/next/workshop/images/divergence_formula.png b/docs/user/next/workshop/images/divergence_formula.png new file mode 100644 index 0000000000..e7cc7d0f5a Binary files /dev/null and b/docs/user/next/workshop/images/divergence_formula.png differ diff --git a/docs/user/next/workshop/images/divergence_picture.png b/docs/user/next/workshop/images/divergence_picture.png new file mode 100644 index 0000000000..b64c293cf2 Binary files /dev/null and b/docs/user/next/workshop/images/divergence_picture.png differ diff --git a/docs/user/next/workshop/images/e2v-add-code.png b/docs/user/next/workshop/images/e2v-add-code.png new file mode 100644 index 0000000000..a8c1a5a6fb Binary files /dev/null and b/docs/user/next/workshop/images/e2v-add-code.png differ diff --git a/docs/user/next/workshop/images/e2v-add.png b/docs/user/next/workshop/images/e2v-add.png new file mode 100644 index 0000000000..22fc54f1b3 Binary files /dev/null and b/docs/user/next/workshop/images/e2v-add.png differ diff --git a/docs/user/next/workshop/images/edge_orientation.png b/docs/user/next/workshop/images/edge_orientation.png new file mode 100644 index 0000000000..572471b929 Binary files /dev/null and b/docs/user/next/workshop/images/edge_orientation.png differ diff --git a/docs/user/next/workshop/images/field-code.png b/docs/user/next/workshop/images/field-code.png new file mode 100644 index 0000000000..2fd85b1662 Binary files /dev/null and b/docs/user/next/workshop/images/field-code.png differ diff --git a/docs/user/next/workshop/images/field.png b/docs/user/next/workshop/images/field.png new file mode 100644 index 0000000000..0b0b10ded4 Binary files /dev/null and b/docs/user/next/workshop/images/field.png differ diff --git a/docs/user/next/workshop/images/gradient_formula.png b/docs/user/next/workshop/images/gradient_formula.png new file mode 100644 index 0000000000..3f9219a037 Binary files /dev/null and b/docs/user/next/workshop/images/gradient_formula.png differ diff --git a/docs/user/next/workshop/images/gradient_picture.png b/docs/user/next/workshop/images/gradient_picture.png new file mode 100644 index 0000000000..3550b24c2f Binary files /dev/null and b/docs/user/next/workshop/images/gradient_picture.png differ diff --git a/docs/user/next/workshop/images/icon_structure_1.png b/docs/user/next/workshop/images/icon_structure_1.png new file mode 100644 index 0000000000..23d7ead6a3 Binary files /dev/null and b/docs/user/next/workshop/images/icon_structure_1.png differ diff --git a/docs/user/next/workshop/images/icon_structure_2.png b/docs/user/next/workshop/images/icon_structure_2.png new file mode 100644 index 0000000000..f3fc6449c3 Binary files /dev/null and b/docs/user/next/workshop/images/icon_structure_2.png differ diff --git a/docs/user/next/workshop/images/laplacian.png b/docs/user/next/workshop/images/laplacian.png new file mode 100644 index 0000000000..2e754df78f Binary files /dev/null and b/docs/user/next/workshop/images/laplacian.png differ diff --git a/docs/user/next/workshop/images/local-field.png b/docs/user/next/workshop/images/local-field.png new file mode 100644 index 0000000000..72a1d10cee Binary files /dev/null and b/docs/user/next/workshop/images/local-field.png differ diff --git a/docs/user/next/workshop/images/logos/c2sm_logo.gif b/docs/user/next/workshop/images/logos/c2sm_logo.gif new file mode 100644 index 0000000000..4c31074aee Binary files /dev/null and b/docs/user/next/workshop/images/logos/c2sm_logo.gif differ diff --git a/docs/user/next/workshop/images/logos/cscs_logo.jpeg b/docs/user/next/workshop/images/logos/cscs_logo.jpeg new file mode 100644 index 0000000000..9f400289e3 Binary files /dev/null and b/docs/user/next/workshop/images/logos/cscs_logo.jpeg differ diff --git a/docs/user/next/workshop/images/logos/exclaim_logo.png b/docs/user/next/workshop/images/logos/exclaim_logo.png new file mode 100644 index 0000000000..7056608eb3 Binary files /dev/null and b/docs/user/next/workshop/images/logos/exclaim_logo.png differ diff --git a/docs/user/next/workshop/images/logos/mch_logo.svg b/docs/user/next/workshop/images/logos/mch_logo.svg new file mode 100644 index 0000000000..183b89c2ab --- /dev/null +++ b/docs/user/next/workshop/images/logos/mch_logo.svg @@ -0,0 +1,830 @@ + + + + + + + + image/svg+xml + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/user/next/workshop/images/remap-field-code.png b/docs/user/next/workshop/images/remap-field-code.png new file mode 100644 index 0000000000..f55cdd2d91 Binary files /dev/null and b/docs/user/next/workshop/images/remap-field-code.png differ diff --git a/docs/user/next/workshop/images/remap-field.png b/docs/user/next/workshop/images/remap-field.png new file mode 100644 index 0000000000..ebd4f6e2b9 Binary files /dev/null and b/docs/user/next/workshop/images/remap-field.png differ diff --git a/docs/user/next/workshop/images/scan_operator.png b/docs/user/next/workshop/images/scan_operator.png new file mode 100644 index 0000000000..f0c1d03636 Binary files /dev/null and b/docs/user/next/workshop/images/scan_operator.png differ diff --git a/docs/user/next/workshop/images/simple_offset.png b/docs/user/next/workshop/images/simple_offset.png new file mode 100644 index 0000000000..660abe8764 Binary files /dev/null and b/docs/user/next/workshop/images/simple_offset.png differ diff --git a/docs/user/next/workshop/images/toolchain_details_existing_path.png b/docs/user/next/workshop/images/toolchain_details_existing_path.png new file mode 100644 index 0000000000..70eeaca68e Binary files /dev/null and b/docs/user/next/workshop/images/toolchain_details_existing_path.png differ diff --git a/docs/user/next/workshop/images/vector_laplacian.png b/docs/user/next/workshop/images/vector_laplacian.png new file mode 100644 index 0000000000..f27354915c Binary files /dev/null and b/docs/user/next/workshop/images/vector_laplacian.png differ diff --git a/docs/user/next/workshop/images/vector_laplacian_formula.png b/docs/user/next/workshop/images/vector_laplacian_formula.png new file mode 100644 index 0000000000..f27354915c Binary files /dev/null and b/docs/user/next/workshop/images/vector_laplacian_formula.png differ diff --git a/docs/user/next/workshop/images/vector_laplacian_normal_component.png b/docs/user/next/workshop/images/vector_laplacian_normal_component.png new file mode 100644 index 0000000000..2def4060e9 Binary files /dev/null and b/docs/user/next/workshop/images/vector_laplacian_normal_component.png differ diff --git a/docs/user/next/workshop/jupyter_intro.ipynb b/docs/user/next/workshop/jupyter_intro.ipynb new file mode 100644 index 0000000000..c1100223a1 --- /dev/null +++ b/docs/user/next/workshop/jupyter_intro.ipynb @@ -0,0 +1,136 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Introduction to jupyter lab\n", + "\n", + "In this notebook, we want to give you an introduction to jupyter lab. The goal is to learn all prerequisites so that you can solve the exercises." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Jupyter Lab gives us a convenient environment to work with notebooks. The current web interface is Jupyter Lab. It has some very nice features that we want to go through:\n", + "\n", + "#### File Explorer\n", + "\n", + "On the left hand side you will find a simple file explorer with the following useful functionality.\n", + "\n", + "You can navigate between folders by double clicking folders. At the top, you will see the current path and you can navigate back to the starting location.\n", + "\n", + "Double clicking a file will open it. Depending on the file type it will be opened differently. Jupyter Lab also has a basic text editor, so you can open textual files.\n", + "\n", + "It is likely that you opened jupyter lab through mybinder or by running a docker image. In that case, it's useful to download files onto your local machine. You can do that by right clicking a file and then selecting download. Vice versa, you can also upload files by clicking the upload file button (up arrow icon) at the top button toolbar of the file explorer.\n", + "\n", + "The file explorer is quite useful and you're generally encouraged to look for features yourselves.\n", + "\n", + "#### Launcher\n", + "\n", + "The next feature we want to look at is the launcher. It is accessible through the file explorere. The plus icon at the top toolbar of the file explorer will open the launcher. With the launcher you can start new python notebooks, python consoles and even a basic terminal. This way you can, e.g., experiment with dusk & dawn more interactively by using a python console or a terminal. Additionally, you can create your own python notebooks for experimentation and if you want, you can also download it.\n", + "\n", + "#### Windowing\n", + "\n", + "Jupyter lab has pretty good support for having different _windows_ open. If you open multiple notebooks, terminals or consoles, they will all be accessible from the top bar. Almost at the top of this web interface. By dragging the window title, you can arrange the windows however you want. E.g., side by side or top/down arrangements are possible as well. You are also encouraged to arrange them to your liking. Some useful layouts could be to open an exercise notebook side by side with the corresponding exercise slides or to have a python console below your exercise notebook, to experiment with dusk/dawn's functionality while solving the exercise." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Jupyter Lab's notebooks\n", + "\n", + "What you are currently reading is called a python notebook. Here we want to give you a quick introduction to that.\n", + "\n", + "A notebook consists of cells. These cells can contain different types of content. E.g., markdown text (as is the case for this cell) or code.\n", + "Below is an example of a cell with code:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# This is a cell with python code.\n", + "print(\"Hello, World!\")\n", + "\n", + "# We can write python code which will be executed and we'll see the output of that code below:" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The output of a cell always corresponds to an execution of that cell. If the code of the cell gets modified, we'll only see the modifications in the output if we execute the cell again. At the top of this notebook, you find a toolbar with various actions we can carry out on cells. This toolbar includes the possibility to execute a cell (the _play_ button). Try it for the cell below:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(\"\".join(map(lambda c: chr(ord(c) - 2), \"Kv\\\"yqtmgf#\")))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You should see as the output \"It worked!\".\n", + "\n", + "You can add new cells by clicking the plus button in the _cell's toolbar_. The type of the cell can be changed with the drop down menu in the same toolbar. Code cells are probably the most useful for you.\n", + "\n", + "Sometimes it can take some time for a cell to execute. While the cell executes, you will see a star ('[\\*]') on its left hand side. Try it with the cell below:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from time import sleep\n", + "\n", + "for word in \"While the cell executes, you will see a star ('[*]') on its left hand side.\".split():\n", + " print(word + \" \", end=\"\")\n", + " sleep(0.8)\n", + "print()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Feel free to experiment with the top toolbar a bit.\n", + "\n", + "The order of executing the cell matters, since each cell can introduce and use variables for other cells. Sometimes, your code doesn't work anymore, because the cells were executed in a bad order. You can reset this easily, by clicking the button to restart and execute the whole notebook again (the two left arrows)." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + }, + "output_auto_scroll": false + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/user/next/workshop/slides/slides_1.ipynb b/docs/user/next/workshop/slides/slides_1.ipynb new file mode 100644 index 0000000000..ff3ca764e4 --- /dev/null +++ b/docs/user/next/workshop/slides/slides_1.ipynb @@ -0,0 +1,380 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9726610d", + "metadata": {}, + "source": [ + "\"cscs\" \"c2sm\"\n", + "\"exclaim\" \"mch\"" + ] + }, + { + "cell_type": "markdown", + "id": "ecab97a5", + "metadata": {}, + "source": [ + "# GT4Py workshop" + ] + }, + { + "cell_type": "markdown", + "id": "9046c2b6", + "metadata": {}, + "source": [ + "## GT4Py: GridTools for Python\n", + "\n", + "GT4Py is a Python library for generating high performance implementations of stencil kernels from a high-level definition using regular Python functions.\n", + "\n", + "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.\n", + "\n", + "**NOTE:** The `gt4py.next` subpackage contains a new and currently experimental version of GT4Py.\n", + "\n", + "## Description\n", + "\n", + "GT4Py is a Python library for expressing computational motifs as found in weather and climate applications.\n", + "\n", + "These computations are expressed in a domain specific language (DSL) which is translated to high-performance implementations for CPUs and GPUs.\n", + "\n", + "In addition, GT4Py provides functions to allocate arrays with memory layout suited for a particular backend.\n", + "\n", + "The following backends are supported:\n", + "\n", + "- `None` aka _embedded_: runs the DSL code directly via the Python interpreter (experimental)\n", + "- `gtfn_cpu` and `gtfn_gpu`: transpiles the DSL to C++ code using the GridTools library\n", + "- `dace`: uses the DaCe library to generate optimized code (experimental)\n", + "\n", + "In this workshop we will mainly use the _embedded_ backend.\n", + "\n", + "## Current efforts\n", + "\n", + "GT4Py is being used to port the ICON model from FORTRAN. Currently the **dycore**, **diffusion**, and **microphysics** are complete.\n", + "\n", + "The ultimate goal is to have a more flexible and modularized model that can be run on CSCS Alps infrastructure as well as other hardware.\n", + "\n", + "Other models ported using GT4Py are ECMWF's FVM, in global (with `gt4py.next` and local area configuration (with `gt4py.cartesian`) and GFDL's FV3 (with `gt4py.cartesian`; original port by AI2)." + ] + }, + { + "cell_type": "markdown", + "id": "79cec88a", + "metadata": {}, + "source": [ + "## Installation and setup\n", + "\n", + "Get an account from https://docs.google.com/document/d/1SuMr2sEdsGHGcnSFczNLGdTVYvNuuXBpCqB-3zL1E9c/edit?usp=sharing and mark with your name.\n", + "\n", + "After cloning the repository to $SCRATCH and setting a symlink to your home-directory\n", + "\n", + "```\n", + "cd $SCRATCH\n", + "git clone --branch gt4py-workshop https://github.com/GridTools/gt4py.git\n", + "cd $HOME\n", + "ln -s $SCRATCH/gt4py\n", + "```\n", + "\n", + "you can install the library with pip.\n", + "\n", + "Make sure that GT4Py is in the expected location, remove `#` and run the cell)" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "320699c6", + "metadata": {}, + "outputs": [], + "source": [ + "#! pip install $HOME/gt4py" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "c2741342", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "f0aa1e50", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "import numpy as np\n", + "import gt4py.next as gtx\n", + "from gt4py.next import float64, Dims" + ] + }, + { + "cell_type": "markdown", + "id": "374952f2", + "metadata": {}, + "source": [ + "## Key concepts and application structure\n", + "\n", + "- [Fields](#Fields),\n", + "- [Field operators](#Field-operators), and\n", + "- [Programs](#Programs)." + ] + }, + { + "cell_type": "markdown", + "id": "831781be", + "metadata": {}, + "source": [ + "### Fields\n", + "\n", + "Fields are **multi-dimensional array** defined over a set of dimensions and a dtype: `gtx.Field[Dims[dimensions], dtype]`.\n", + "\n", + "|
|\n", + "| :----------------------------------------------------------: |\n", + "|
|" + ] + }, + { + "cell_type": "markdown", + "id": "55803c7d", + "metadata": {}, + "source": [ + "Fields can be constructed with the following functions, inspired by numpy:\n", + "\n", + "- `zeros`\n", + "- `full` to fill with a given value\n", + "- `as_field` to convert from numpy or cupy arrays\n", + "\n", + "The first argument is the domain of the field, which can be constructed from a mapping from `Dimension` to range.\n", + "\n", + "Optional we can pass\n", + "\n", + "- `dtype` the description of type of the field\n", + "- `allocator` which describes how and where (e.g. GPU) the buffer is allocated.\n", + "\n", + "Note: `as_field` can also take a sequence of Dimensions and infer the shape" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "69fc3531", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a definition: \n", + " ⟨Domain(Cell[horizontal]=(0:5), K[vertical]=(0:6)) → DType(scalar_type=, tensor_shape=())⟩\n", + "a array: \n", + " [[0. 0. 0. 0. 0. 0.]\n", + " [0. 0. 0. 0. 0. 0.]\n", + " [0. 0. 0. 0. 0. 0.]\n", + " [0. 0. 0. 0. 0. 0.]\n", + " [0. 0. 0. 0. 0. 0.]]\n", + "b array: \n", + " [[3. 3. 3. 3. 3. 3.]\n", + " [3. 3. 3. 3. 3. 3.]\n", + " [3. 3. 3. 3. 3. 3.]\n", + " [3. 3. 3. 3. 3. 3.]\n", + " [3. 3. 3. 3. 3. 3.]]\n", + "c array: \n", + " [[ 0. 1. 2. 3. 4. 5.]\n", + " [10. 11. 12. 13. 14. 15.]\n", + " [20. 21. 22. 23. 24. 25.]\n", + " [30. 31. 32. 33. 34. 35.]\n", + " [40. 41. 42. 43. 44. 45.]]\n" + ] + } + ], + "source": [ + "Cell = gtx.Dimension(\"Cell\")\n", + "K = gtx.Dimension(\"K\", kind=gtx.DimensionKind.VERTICAL)\n", + "\n", + "domain = gtx.domain({Cell: 5, K: 6})\n", + "\n", + "a = gtx.zeros(domain, dtype=float64)\n", + "b = gtx.full(domain, fill_value=3.0, dtype=float64)\n", + "c = gtx.as_field([Cell, K], np.fromfunction(lambda c, k: c*10+k, shape=domain.shape))\n", + "\n", + "print(\"a definition: \\n {}\".format(a))\n", + "print(\"a array: \\n {}\".format(a.asnumpy()))\n", + "print(\"b array: \\n {}\".format(b.asnumpy()))\n", + "print(\"c array: \\n {}\".format(c.asnumpy()))" + ] + }, + { + "cell_type": "markdown", + "id": "59284eb5", + "metadata": {}, + "source": [ + "### Field operators\n", + "\n", + "Field operators perform operations on a set of fields, i.e. elementwise addition or reduction along a dimension.\n", + "\n", + "They are written as Python functions by using the `@field_operator` decorator." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "4199bc10", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def add(a: gtx.Field[Dims[Cell, K], float64],\n", + " b: gtx.Field[Dims[Cell, K], float64]) -> gtx.Field[Dims[Cell, K], float64]:\n", + " return a + b" + ] + }, + { + "cell_type": "markdown", + "id": "fa3e7c47", + "metadata": {}, + "source": [ + "Direct calls to field operators require two additional arguments:\n", + "\n", + "- `out`: a field to write the return value to\n", + "- `offset_provider`: empty dict for now, explanation will follow" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "c2d61c1a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "result array \n", + " [[3. 3. 3. 3. 3. 3.]\n", + " [3. 3. 3. 3. 3. 3.]\n", + " [3. 3. 3. 3. 3. 3.]\n", + " [3. 3. 3. 3. 3. 3.]\n", + " [3. 3. 3. 3. 3. 3.]]\n" + ] + } + ], + "source": [ + "result = gtx.zeros(domain)\n", + "add(a, b, out=result, offset_provider={})\n", + "\n", + "print(\"result array \\n {}\".format(result.asnumpy()))" + ] + }, + { + "cell_type": "markdown", + "id": "85aeccf6", + "metadata": {}, + "source": [ + "### Programs" + ] + }, + { + "cell_type": "markdown", + "id": "6b6e8526", + "metadata": {}, + "source": [ + "Programs are used to call field operators to mutate the latter's output arguments.\n", + "\n", + "They are written as Python functions by using the `@program` decorator.\n", + "\n", + "This example below calls the `add` field operator twice:" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "1287916a", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.program\n", + "def run_add(a : gtx.Field[Dims[Cell, K], float64],\n", + " b : gtx.Field[Dims[Cell, K], float64],\n", + " result : gtx.Field[Dims[Cell, K], float64]):\n", + " add(a, b, out=result)\n", + " add(b, result, out=result)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "d27f8366", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "result array: \n", + " [[6. 6. 6. 6. 6. 6.]\n", + " [6. 6. 6. 6. 6. 6.]\n", + " [6. 6. 6. 6. 6. 6.]\n", + " [6. 6. 6. 6. 6. 6.]\n", + " [6. 6. 6. 6. 6. 6.]]\n" + ] + } + ], + "source": [ + "result = gtx.zeros(domain, dtype=float64)\n", + "run_add(a, b, result, offset_provider={})\n", + "\n", + "print(\"result array: \\n {}\".format(result.asnumpy()))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "131f1247-314e-4ecb-97a6-2204a4418c17", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/slides/slides_2.ipynb b/docs/user/next/workshop/slides/slides_2.ipynb new file mode 100644 index 0000000000..903182b7bc --- /dev/null +++ b/docs/user/next/workshop/slides/slides_2.ipynb @@ -0,0 +1,435 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2d564f4a", + "metadata": {}, + "source": [ + "\"cscs\" \"c2sm\"\n", + "\"exclaim\" \"mch\"" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "be426d3c", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "f46d4bcc", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "import numpy as np\n", + "import gt4py.next as gtx\n", + "from gt4py.next import float64, neighbor_sum, Dims" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "8724b3c1", + "metadata": {}, + "outputs": [], + "source": [ + "Cell = gtx.Dimension(\"Cell\")\n", + "K = gtx.Dimension(\"K\", kind=gtx.DimensionKind.VERTICAL)" + ] + }, + { + "cell_type": "markdown", + "id": "fa5157be", + "metadata": {}, + "source": [ + "## Offsets\n", + "Fields can be shifted with a (Cartesian) offset.\n", + "\n", + "Take the following array:" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "0f816fbc", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a_off array: \n", + " [1. 1. 2. 3. 5. 8.]\n" + ] + } + ], + "source": [ + "a_off = gtx.as_field([K], np.array([1.0, 1.0, 2.0, 3.0, 5.0, 8.0]))\n", + "\n", + "print(\"a_off array: \\n {}\".format(a_off.asnumpy()))" + ] + }, + { + "cell_type": "markdown", + "id": "fd78cff7", + "metadata": {}, + "source": [ + "Visually, offsetting this field by 1 would result in the following:\n", + "\n", + "| ![Coff](../images/simple_offset.png) |\n", + "| :------------------------: |\n", + "| _CellDim Offset (Coff)_ |\n", + "\n", + "In GT4Py we express this by" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "1e8670ca", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "result field: \n", + " ⟨Domain(K[vertical]=(0:6)) → DType(scalar_type=, tensor_shape=())⟩ \n", + " [1. 2. 3. 5. 8. 0.]\n" + ] + } + ], + "source": [ + "Koff = gtx.FieldOffset(\"Koff\", source=K, target=(K,))\n", + "\n", + "@gtx.field_operator\n", + "def a_offset(a_off: gtx.Field[Dims[K], float64]) -> gtx.Field[Dims[K], float64]:\n", + " return a_off(Koff[1])\n", + "\n", + "result = gtx.zeros(gtx.domain({K: 6}))\n", + "\n", + "a_offset(a_off, out=result[:-1], offset_provider={\"Koff\": K})\n", + "print(f\"result field: \\n {result} \\n {result.asnumpy()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "aa841d31", + "metadata": {}, + "source": [ + "## Defining the mesh and its connectivities\n", + "Take an unstructured mesh with numbered cells (in red) and edges (in blue).\n", + "\n", + "| ![grid_topo](../images/connectivity_numbered_grid.svg) |\n", + "| :------------------------------------------: |\n", + "| _The mesh with the indices_ |" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "64621d61", + "metadata": {}, + "outputs": [], + "source": [ + "Cell = gtx.Dimension(\"Cell\")\n", + "Edge = gtx.Dimension(\"Edge\")" + ] + }, + { + "cell_type": "markdown", + "id": "8da1ce25", + "metadata": {}, + "source": [ + "Connectivity among mesh elements is expressed through connectivity tables.\n", + "\n", + "For example, `e2c_table` lists for each edge its adjacent cells. \n", + "\n", + "Similarly, `c2e_table` lists the edges that are neighbors to a particular cell.\n", + "\n", + "Note that if an edge is lying at the border, one entry will be filled with -1." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "9e15475d", + "metadata": {}, + "outputs": [], + "source": [ + "e2c_table = np.array([\n", + " [0, -1], # edge 0 (neighbours: cell 0)\n", + " [2, -1], # edge 1\n", + " [2, -1], # edge 2\n", + " [3, -1], # edge 3\n", + " [4, -1], # edge 4\n", + " [5, -1], # edge 5\n", + " [0, 5], # edge 6 (neighbours: cell 0, cell 5)\n", + " [0, 1], # edge 7\n", + " [1, 2], # edge 8\n", + " [1, 3], # edge 9\n", + " [3, 4], # edge 10\n", + " [4, 5] # edge 11\n", + "])\n", + "\n", + "c2e_table = np.array([\n", + " [0, 6, 7], # cell 0 (neighbors: edge 0, edge 6, edge 7)\n", + " [7, 8, 9], # cell 1\n", + " [1, 2, 8], # cell 2\n", + " [3, 9, 10], # cell 3\n", + " [4, 10, 11], # cell 4\n", + " [5, 6, 11], # cell 5\n", + "])" + ] + }, + { + "cell_type": "markdown", + "id": "50ee54b2", + "metadata": {}, + "source": [ + "#### Using connectivities in field operators\n", + "\n", + "Let's start by defining two fields: one over the cells and another one over the edges. The field over cells serves input as for subsequent calculations and is therefore filled up with values, whereas the field over the edges stores the output of the calculations and is therefore left blank." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "19a2c01a", + "metadata": {}, + "outputs": [], + "source": [ + "cell_field = gtx.as_field([Cell], np.array([1.0, 1.0, 2.0, 3.0, 5.0, 8.0]))\n", + "edge_field = gtx.as_field([Edge], np.zeros((12,)))" + ] + }, + { + "cell_type": "markdown", + "id": "d4184fbc", + "metadata": {}, + "source": [ + "| ![cell_values](../images/connectivity_cell_field.svg) |\n", + "| :-----------------------------------------: |\n", + "| _Cell values_ |" + ] + }, + { + "cell_type": "markdown", + "id": "fbeb2868", + "metadata": {}, + "source": [ + "`field_offset` is used to remap fields over one domain to another domain, e.g. cells -> edges." + ] + }, + { + "cell_type": "markdown", + "id": "1c02e0a8", + "metadata": {}, + "source": [ + "Field remappings are just composition of mappings\n", + "- Field defined on cells: $f_C: C \\to \\mathbb{R}$\n", + "- Connectivity from _edges to cells_: $c_{E \\to C_0}$\n", + "- We define a new field on edges composing both mappings\n", + "$$ f_E: E \\to \\mathbb{R}, e \\mapsto (f_C \\circ c_{E \\to C_0})(e) := f_c(c_{E \\to C_0}(e)) $$\n", + "- In point-free notation: $f_E = f_C(c_{E \\to C_0}) \\Rightarrow$ `f_c(E2C[0])`\n", + "\n", + "\n", + "We extend the connectivities to refer to more than just one neighbor\n", + "- `E2CDim` is the local dimension of all cell neighbors of an edge\n", + "\n", + "$$ c_{E \\to C}: E \\times E2CDim \\to C $$\n", + "$$ f_E(e, l) := f_C(c_{E \\to C}(e, l)), e \\in E, l \\in \\{0,1\\} $$" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "2f87a176", + "metadata": {}, + "outputs": [], + "source": [ + "E2CDim = gtx.Dimension(\"E2C\", kind=gtx.DimensionKind.LOCAL)\n", + "E2C = gtx.FieldOffset(\"E2C\", source=Cell, target=(Edge, E2CDim))" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "6d30a5e1", + "metadata": {}, + "outputs": [], + "source": [ + "E2C_offset_provider = gtx.NeighborTableOffsetProvider(e2c_table, Edge, Cell, 2)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "d62f6c98", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0th adjacent cell's value: [1. 2. 2. 3. 5. 8. 1. 1. 1. 1. 3. 5.]\n" + ] + } + ], + "source": [ + "@gtx.field_operator\n", + "def nearest_cell_to_edge(cell_field: gtx.Field[Dims[Cell], float64]) -> gtx.Field[Dims[Edge], float64]:\n", + " return cell_field(E2C[0]) # 0th index to isolate edge dimension\n", + "\n", + "@gtx.program # uses skip_values, therefore we cannot use embedded\n", + "def run_nearest_cell_to_edge(cell_field: gtx.Field[Dims[Cell], float64], edge_field: gtx.Field[Dims[Edge], float64]):\n", + " nearest_cell_to_edge(cell_field, out=edge_field)\n", + "\n", + "run_nearest_cell_to_edge(cell_field, edge_field, offset_provider={\"E2C\": E2C_offset_provider})\n", + "\n", + "print(\"0th adjacent cell's value: {}\".format(edge_field.asnumpy()))" + ] + }, + { + "cell_type": "markdown", + "id": "57e47a67", + "metadata": {}, + "source": [ + "Running the above snippet results in the following edge field:\n", + "\n", + "| ![nearest_cell_values](../images/connectivity_numbered_grid.svg) | $\\mapsto$ | ![grid_topo](../images/connectivity_edge_0th_cell.svg) |\n", + "| :----------------------------------------------------: | :-------: | :------------------------------------------: |\n", + "| _Domain (edges)_ | | _Edge values_ |" + ] + }, + { + "cell_type": "markdown", + "id": "c94ef7e7", + "metadata": {}, + "source": [ + "### Another example: E2V\n", + "\n", + "Creating fields on edges from fields on vertices using an **E2V** connectivity:\n", + "\n", + "|
|\n", + "| :-----------------------------------------: |\n", + "|
|\n", + "\n", + "We can create two edge fields from the same vertex field, by taking the values from the start or from the end vertex, and then you can operate wi to the \n", + "|
|\n", + "| :-----------------------------------------: |\n", + "|
|" + ] + }, + { + "cell_type": "markdown", + "id": "aee8bf01", + "metadata": {}, + "source": [ + "### Using reductions on connected mesh elements\n", + "\n", + "To sum up all the cells adjacent to an edge the `neighbor_sum` builtin function can be called to operate along the `E2CDim` dimension." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "d51eae8b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "sum of adjacent cells: [ 1. 2. 2. 3. 5. 8. 9. 2. 3. 4. 8. 13.]\n" + ] + } + ], + "source": [ + "@gtx.field_operator\n", + "def sum_adjacent_cells(cell_field : gtx.Field[Dims[Cell], float64]) -> gtx.Field[Dims[Edge], float64]:\n", + " return neighbor_sum(cell_field(E2C), axis=E2CDim)\n", + "\n", + "@gtx.program # uses skip_values, therefore we cannot use embedded\n", + "def run_sum_adjacent_cells(cell_field : gtx.Field[Dims[Cell], float64], edge_field: gtx.Field[Dims[Edge], float64]):\n", + " sum_adjacent_cells(cell_field, out=edge_field)\n", + "\n", + "run_sum_adjacent_cells(cell_field, edge_field, offset_provider={\"E2C\": E2C_offset_provider})\n", + "\n", + "print(\"sum of adjacent cells: {}\".format(edge_field.asnumpy()))" + ] + }, + { + "cell_type": "markdown", + "id": "aadeae6d", + "metadata": {}, + "source": [ + "For the border edges, the results are unchanged compared to the previous example, but the inner edges now contain the sum of the two adjacent cells:\n", + "\n", + "| ![nearest_cell_values](../images/connectivity_numbered_grid.svg) | $\\mapsto$ | ![cell_values](../images/connectivity_edge_cell_sum.svg) |\n", + "| :----------------------------------------------------: | :-------: | :--------------------------------------------: |\n", + "| _Domain (edges)_ | | _Edge values_ |" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dcc661aa-8300-4c4c-9437-a85e6d319fae", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3906e41a-964c-4157-a3e8-29f8969162f4", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/slides/slides_3.ipynb b/docs/user/next/workshop/slides/slides_3.ipynb new file mode 100644 index 0000000000..362a169322 --- /dev/null +++ b/docs/user/next/workshop/slides/slides_3.ipynb @@ -0,0 +1,204 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "b4fe65fe", + "metadata": {}, + "source": [ + "\"cscs\" \"c2sm\"\n", + "\"exclaim\" \"mch\"" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "c1e32292", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "42fa04fd", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "import numpy as np\n", + "import gt4py.next as gtx\n", + "from gt4py.next import where, Dims" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "5f51044d", + "metadata": {}, + "outputs": [], + "source": [ + "Cell = gtx.Dimension(\"Cell\")\n", + "K = gtx.Dimension(\"K\", kind=gtx.DimensionKind.VERTICAL)" + ] + }, + { + "cell_type": "markdown", + "id": "ca021e19", + "metadata": {}, + "source": [ + "## Using conditionals on Fields\n", + "\n", + "To conditionally compose a Field from two inputs, we borrow the `where` function from numpy. \n", + "\n", + "This function takes 3 input arguments:\n", + "- mask: a Field of booleans\n", + "- true branch: a tuple, a Field, or a scalar\n", + "- false branch: a tuple, a Field, of a scalar" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "d8ad52ce", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "mask array: [ True False True True False]\n", + "true_Field array: [11. 12. 13. 14. 15.]\n", + "false_Field array: [21. 22. 23. 24. 25.]\n", + "where return: [11. 22. 13. 14. 25.]\n" + ] + } + ], + "source": [ + "mask = gtx.as_field([Cell], np.asarray([True, False, True, True, False]))\n", + "\n", + "true_Field = gtx.as_field([Cell], np.asarray([11.0, 12.0, 13.0, 14.0, 15.0]))\n", + "false_Field = gtx.as_field([Cell], np.asarray([21.0, 22.0, 23.0, 24.0, 25.0]))\n", + "\n", + "result = gtx.zeros(gtx.domain({Cell:5}))\n", + "\n", + "@gtx.field_operator\n", + "def conditional(mask: gtx.Field[Dims[Cell], bool], true_Field: gtx.Field[Dims[Cell], gtx.float64], false_Field: gtx.Field[Dims[Cell], gtx.float64]\n", + ") -> gtx.Field[Dims[Cell], gtx.float64]:\n", + " return where(mask, true_Field, false_Field)\n", + "\n", + "conditional(mask, true_Field, false_Field, out=result, offset_provider={})\n", + "print(\"mask array: {}\".format(mask.asnumpy()))\n", + "print(\"true_Field array: {}\".format(true_Field.asnumpy()))\n", + "print(\"false_Field array: {}\".format(false_Field.asnumpy()))\n", + "print(\"where return: {}\".format(result.asnumpy()))" + ] + }, + { + "cell_type": "markdown", + "id": "125181f9", + "metadata": {}, + "source": [ + "## Using domain on Fields\n", + "\n", + "By default the whole `out` Field is updated. If only a subset should be updated, we can specify the output domain by passing the `domain` keyword argument when calling the Field operator." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "777d2613", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def add(a: gtx.Field[Dims[Cell, K], gtx.float64],\n", + " b: gtx.Field[Dims[Cell, K], gtx.float64]) -> gtx.Field[Dims[Cell, K], gtx.float64]:\n", + " return a + b # 2.0 + 3.0\n", + "\n", + "@gtx.program\n", + "def run_add_domain(a : gtx.Field[Dims[Cell, K], gtx.float64],\n", + " b : gtx.Field[Dims[Cell, K], gtx.float64],\n", + " result : gtx.Field[Dims[Cell, K], gtx.float64]):\n", + " add(a, b, out=result, domain={Cell: (1, 3), K: (1, 4)}) " + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "a2f9b111", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "result array: \n", + " [[0. 0. 0. 0. 0. 0.]\n", + " [0. 5. 5. 5. 0. 0.]\n", + " [0. 5. 5. 5. 0. 0.]\n", + " [0. 0. 0. 0. 0. 0.]\n", + " [0. 0. 0. 0. 0. 0.]]\n" + ] + } + ], + "source": [ + "domain = gtx.domain({Cell: 5, K: 6})\n", + "\n", + "a = gtx.full(domain, fill_value=2.0, dtype=np.float64)\n", + "b = gtx.full(domain, fill_value=3.0, dtype=np.float64)\n", + "result = gtx.zeros(domain)\n", + "run_add_domain(a, b, result, offset_provider={})\n", + "\n", + "print(\"result array: \\n {}\".format(result.asnumpy()))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2b3a64f3", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/workshop/slides/slides_4.ipynb b/docs/user/next/workshop/slides/slides_4.ipynb new file mode 100644 index 0000000000..324e30fe29 --- /dev/null +++ b/docs/user/next/workshop/slides/slides_4.ipynb @@ -0,0 +1,201 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "ec002a4c", + "metadata": {}, + "source": [ + "\"cscs\" \"c2sm\"\n", + "\"exclaim\" \"mch\"" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "9ae7e945", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "de5c676c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "import numpy as np\n", + "import gt4py.next as gtx\n", + "from gt4py.next import float64, Dims" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "c99ea049", + "metadata": {}, + "outputs": [], + "source": [ + "K = gtx.Dimension(\"K\", kind=gtx.DimensionKind.VERTICAL)" + ] + }, + { + "cell_type": "markdown", + "id": "256fbdfd", + "metadata": {}, + "source": [ + "## Scan algorithm\n", + "\n", + "All operations so far where map operations over the output domain. The only other algorithm that we currently support is _scanning_ of an axis, one example of a scan is the partial sum as illustrated in the following code snippet." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "e0a06aa8", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "input:\n", + " [1. 2. 4. 6. 0. 2. 5.]\n", + "partial sum:\n", + " [ 1. 3. 7. 13. 13. 15. 20.]\n" + ] + } + ], + "source": [ + "x = np.asarray([1.0, 2.0, 4.0, 6.0, 0.0, 2.0, 5.0])\n", + "def partial_sum(x):\n", + " for i in range(len(x)):\n", + " if i > 0:\n", + " x[i] = x[i-1] + x[i]\n", + " return x\n", + "print(f\"input:\\n {x}\") \n", + "print(f\"partial sum:\\n {partial_sum(x)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "8cc7af6a", + "metadata": {}, + "source": [ + "Visually, this is what `partial_sum` is doing: \n", + "\n", + "| ![scan_operator](../images/scan_operator.png) |\n", + "| :---------------------------------: |\n", + "| _Iterative sum over K_ |" + ] + }, + { + "cell_type": "markdown", + "id": "0efcdd5c", + "metadata": {}, + "source": [ + "In GT4Py the a scan pattern is implemented with the so-called `scan_operator` where the return statement expresses the computation at the current position in the scan direction. This value is additionally injected as the first argument to the next position, usually called `state` or `carry`.\n", + "\n", + "The `scan_operator` decorator takes 3 arguments:\n", + "- `axis`: a `Dimension` that specifies the scan axis; note: the `Dimension` has to be of kind `VERTICAL`\n", + "- `forward`: True if order of operations is from bottom to top, False if from top to bottom\n", + "- `init`: value that is injected as the `state` at the start\n", + "\n", + "Note: Unlike a `field_operator`, the `scan_operator` is actually a local, scalar operation. It is applied to all points in the dimensions orthogonal to the scan axis (a form of single-column-abstraction). That might change in the future or be extended with a field version." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "aef8d1ce", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.scan_operator(axis=K, forward=True, init=0.0)\n", + "def add_scan(state: float, k_field: float) -> float:\n", + " return state + k_field" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "d56f1056", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "result array: \n", + " [ 1. 3. 7. 13. 13. 15. 20.]\n" + ] + } + ], + "source": [ + "k_field = gtx.as_field([K], np.asarray([1.0, 2.0, 4.0, 6.0, 0.0, 2.0, 5.0]))\n", + "result = gtx.zeros(domain=gtx.domain({K: 7}))\n", + "\n", + "add_scan(k_field, out=result, offset_provider={})\n", + "\n", + "print(\"result array: \\n {}\".format(result.asnumpy()))" + ] + }, + { + "cell_type": "markdown", + "id": "d3088de4", + "metadata": {}, + "source": [ + "Note: `scan_operators` can be called from `field_operators` and `programs`. Likewise, `field_operators` can be called from `scan_operators`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0c4b8ea2", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tox.ini b/tox.ini index ff67764b05..d4418c4ebc 100644 --- a/tox.ini +++ b/tox.ini @@ -107,7 +107,13 @@ commands = [testenv:notebooks-py{310,311}] description = Run notebooks -commands = python -m pytest --nbmake examples -v -n {env:NUM_PROCESSES:1} +commands_pre = + jupytext docs/user/next/QuickstartGuide.md --to .ipynb +commands = + python -m pytest --nbmake docs/user/next/workshop/slides -v -n {env:NUM_PROCESSES:1} + python -m pytest --nbmake docs/user/next/workshop/exercises -k 'solutions' -v -n {env:NUM_PROCESSES:1} + python -m pytest --nbmake docs/user/next/QuickstartGuide.ipynb -v -n {env:NUM_PROCESSES:1} + python -m pytest --nbmake examples -v -n {env:NUM_PROCESSES:1} # -- Other artefacts -- [testenv:dev-py{38,39,310,311}{-atlas,}] @@ -120,15 +126,6 @@ set_env = {[testenv]set_env} PIP_EXTRA_INDEX_URL = {env:PIP_EXTRA_INDEX_URL:https://test.pypi.org/simple/} -# [testenv:docs] -# usedevelop = true -# commands_pre = -# changedir = docs/user/next -# commands = -# jupytext QuickstartGuide.md --to .py -# python QuickstartGuide.py -# commands_post = - # [testenv:diagrams] # install_command = echo {packages} # skip_install = true