From d65c1295e2390ba29876846e5e58c6c0714fb5ae Mon Sep 17 00:00:00 2001 From: Enrique Gonzalez Paredes Date: Thu, 23 Nov 2023 08:48:31 +0100 Subject: [PATCH 1/4] Expand slide 2 --- docs/user/next/slides/slides_2.md | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/docs/user/next/slides/slides_2.md b/docs/user/next/slides/slides_2.md index 3566f7a857..4901a6eb3f 100644 --- a/docs/user/next/slides/slides_2.md +++ b/docs/user/next/slides/slides_2.md @@ -173,12 +173,21 @@ Running the above snippet results in the following edge field: +++ -Another example: E2V +### Another example: E2V + +Creating fields on edges from fields on vertices using an **E2V** connectivity: |
| | :-----------------------------------------: | |
| +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 +|
| +| :-----------------------------------------: | +|
| + + + +++ ### Using reductions on connected mesh elements From 2a8c90eddd55e471ab1e339d981888dd55e08de5 Mon Sep 17 00:00:00 2001 From: Enrique Gonzalez Paredes Date: Thu, 23 Nov 2023 08:52:11 +0100 Subject: [PATCH 2/4] Add ipynb versions of the slides --- docs/user/next/slides/slides_1.ipynb | 289 +++++++++++++++++++++ docs/user/next/slides/slides_2.ipynb | 358 +++++++++++++++++++++++++++ docs/user/next/slides/slides_3.ipynb | 154 ++++++++++++ docs/user/next/slides/slides_4.ipynb | 155 ++++++++++++ 4 files changed, 956 insertions(+) create mode 100644 docs/user/next/slides/slides_1.ipynb create mode 100644 docs/user/next/slides/slides_2.ipynb create mode 100644 docs/user/next/slides/slides_3.ipynb create mode 100644 docs/user/next/slides/slides_4.ipynb diff --git a/docs/user/next/slides/slides_1.ipynb b/docs/user/next/slides/slides_1.ipynb new file mode 100644 index 0000000000..f749fb4528 --- /dev/null +++ b/docs/user/next/slides/slides_1.ipynb @@ -0,0 +1,289 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "4598486a", + "metadata": {}, + "source": [ + "\"cscs\" \"c2sm\"\n", + "\"exclaim\" \"mch\"" + ] + }, + { + "cell_type": "markdown", + "id": "c5abdaa5", + "metadata": {}, + "source": [ + "# GT4Py workshop" + ] + }, + { + "cell_type": "markdown", + "id": "64c0b236", + "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", + "- `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": "d84571cc", + "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": null, + "id": "d55089eb", + "metadata": {}, + "outputs": [], + "source": [ + "#! pip install $HOME/gt4py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e25b300a", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4ac0ca57", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import gt4py.next as gtx\n", + "from gt4py.next import float64, neighbor_sum, where" + ] + }, + { + "cell_type": "markdown", + "id": "219a8afd", + "metadata": {}, + "source": [ + "## Key concepts and application structure\n", + "\n", + "- [Fields](#Fields),\n", + "- [Field operators](#Field-operators), and\n", + "- [Programs](#Programs)." + ] + }, + { + "cell_type": "markdown", + "id": "332533a3", + "metadata": {}, + "source": [ + "### Fields\n", + "Fields are **multi-dimensional array** defined over a set of dimensions and a dtype: `gtx.Field[[dimensions], dtype]`.\n", + "\n", + "|
|\n", + "| :-----------------------------------------: |\n", + "|
|" + ] + }, + { + "cell_type": "markdown", + "id": "437cad42", + "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", + "- `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": null, + "id": "f4be5a76", + "metadata": {}, + "outputs": [], + "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": "f7875b85", + "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": null, + "id": "7092af27", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def add(a: gtx.Field[[Cell, K], float64],\n", + " b: gtx.Field[[Cell, K], float64]) -> gtx.Field[[Cell, K], float64]:\n", + " return a + b" + ] + }, + { + "cell_type": "markdown", + "id": "ba88d740", + "metadata": {}, + "source": [ + "Direct calls to field operators require two additional arguments: \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": null, + "id": "e93f8d29", + "metadata": {}, + "outputs": [], + "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": "b482c77c", + "metadata": {}, + "source": [ + "### Programs" + ] + }, + { + "cell_type": "markdown", + "id": "2b5dbfe8", + "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": null, + "id": "7ff60fda", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.program\n", + "def run_add(a : gtx.Field[[Cell, K], float64],\n", + " b : gtx.Field[[Cell, K], float64],\n", + " result : gtx.Field[[Cell, K], float64]):\n", + " add(a, b, out=result)\n", + " add(b, result, out=result)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "08713dc1", + "metadata": {}, + "outputs": [], + "source": [ + "result = gtx.zeros(domain, dtype=float64)\n", + "run_add(a, b, result, offset_provider={})\n", + "\n", + "print(\"result array: \\n {}\".format(result.asnumpy()))" + ] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/slides/slides_2.ipynb b/docs/user/next/slides/slides_2.ipynb new file mode 100644 index 0000000000..24830dfbff --- /dev/null +++ b/docs/user/next/slides/slides_2.ipynb @@ -0,0 +1,358 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "fa87d9c9", + "metadata": {}, + "source": [ + "\"cscs\" \"c2sm\"\n", + "\"exclaim\" \"mch\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9bc2e917", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dac85554", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import gt4py.next as gtx\n", + "from gt4py.next import float64, neighbor_sum, where" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "44b71d2c", + "metadata": {}, + "outputs": [], + "source": [ + "Cell = gtx.Dimension(\"Cell\")\n", + "K = gtx.Dimension(\"K\", kind=gtx.DimensionKind.VERTICAL)" + ] + }, + { + "cell_type": "markdown", + "id": "87ea8657", + "metadata": {}, + "source": [ + "## Offsets\n", + "Fields can be shifted with a (Cartesian) offset.\n", + "\n", + "Take the following array:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a5512dbb", + "metadata": {}, + "outputs": [], + "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": "06ff59eb", + "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": null, + "id": "2edd3e62", + "metadata": {}, + "outputs": [], + "source": [ + "Koff = gtx.FieldOffset(\"Koff\", source=K, target=(K,))\n", + "\n", + "@gtx.field_operator\n", + "def a_offset(a_off: gtx.Field[[K], float64]) -> gtx.Field[[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": "7db7006e", + "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": null, + "id": "a12b054c", + "metadata": {}, + "outputs": [], + "source": [ + "Cell = gtx.Dimension(\"Cell\")\n", + "Edge = gtx.Dimension(\"Edge\")" + ] + }, + { + "cell_type": "markdown", + "id": "b8b521de", + "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": null, + "id": "97de911e", + "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": "6f008a18", + "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": null, + "id": "7eedf9d7", + "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": "fdf9ec2e", + "metadata": {}, + "source": [ + "| ![cell_values](../images/connectivity_cell_field.svg) |\n", + "| :-----------------------------------------: |\n", + "| _Cell values_ |" + ] + }, + { + "cell_type": "markdown", + "id": "deacca9e", + "metadata": {}, + "source": [ + "`field_offset` is used to remap fields over one domain to another domain, e.g. cells -> edges." + ] + }, + { + "cell_type": "markdown", + "id": "1e1c4462", + "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": null, + "id": "78cb7a4c", + "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": null, + "id": "0fc11ed6", + "metadata": {}, + "outputs": [], + "source": [ + "E2C_offset_provider = gtx.NeighborTableOffsetProvider(e2c_table, Edge, Cell, 2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b34802f1", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def nearest_cell_to_edge(cell_field: gtx.Field[[Cell], float64]) -> gtx.Field[[Edge], float64]:\n", + " return cell_field(E2C[0]) # 0th index to isolate edge dimension\n", + "\n", + "@gtx.program(backend=gtx.gtfn_cpu) # uses skip_values, therefore we cannot use embedded\n", + "def run_nearest_cell_to_edge(cell_field: gtx.Field[[Cell], float64], edge_field: gtx.Field[[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": "61eea28a", + "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": "0cea8256", + "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": "1d289b73", + "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": null, + "id": "dcd7b2fd", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def sum_adjacent_cells(cell_field : gtx.Field[[Cell], float64]) -> gtx.Field[[Edge], float64]:\n", + " return neighbor_sum(cell_field(E2C), axis=E2CDim)\n", + "\n", + "@gtx.program(backend=gtx.gtfn_cpu) # uses skip_values, therefore we cannot use embedded\n", + "def run_sum_adjacent_cells(cell_field : gtx.Field[[Cell], float64], edge_field: gtx.Field[[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": "9986df3c", + "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_ |" + ] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/slides/slides_3.ipynb b/docs/user/next/slides/slides_3.ipynb new file mode 100644 index 0000000000..7114a8a772 --- /dev/null +++ b/docs/user/next/slides/slides_3.ipynb @@ -0,0 +1,154 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "7cb0a15c", + "metadata": {}, + "source": [ + "\"cscs\" \"c2sm\"\n", + "\"exclaim\" \"mch\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "09384b9b", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "038e6cbf", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import gt4py.next as gtx\n", + "from gt4py.next import neighbor_sum, where" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2ebfe773", + "metadata": {}, + "outputs": [], + "source": [ + "Cell = gtx.Dimension(\"Cell\")\n", + "K = gtx.Dimension(\"K\", kind=gtx.DimensionKind.VERTICAL)" + ] + }, + { + "cell_type": "markdown", + "id": "8305625b", + "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": null, + "id": "5fdcb013", + "metadata": {}, + "outputs": [], + "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[[Cell], bool], true_field: gtx.Field[[Cell], gtx.float64], false_field: gtx.Field[[Cell], gtx.float64]\n", + ") -> gtx.Field[[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": "c9de5ab3", + "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": null, + "id": "58e30dae", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def add(a: gtx.Field[[Cell, K], gtx.float64],\n", + " b: gtx.Field[[Cell, K], gtx.float64]) -> gtx.Field[[Cell, K], gtx.float64]:\n", + " return a + b # 2.0 + 3.0\n", + "\n", + "@gtx.program\n", + "def run_add_domain(a : gtx.Field[[Cell, K], gtx.float64],\n", + " b : gtx.Field[[Cell, K], gtx.float64],\n", + " result : gtx.Field[[Cell, K], gtx.float64]):\n", + " add(a, b, out=result, domain={Cell: (1, 3), K: (1, 4)}) " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5f905447", + "metadata": {}, + "outputs": [], + "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": "2e2052af", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/user/next/slides/slides_4.ipynb b/docs/user/next/slides/slides_4.ipynb new file mode 100644 index 0000000000..b9d5f0531c --- /dev/null +++ b/docs/user/next/slides/slides_4.ipynb @@ -0,0 +1,155 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0a0ab013", + "metadata": {}, + "source": [ + "\"cscs\" \"c2sm\"\n", + "\"exclaim\" \"mch\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "02e63233", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings('ignore')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2892c9e2", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import gt4py.next as gtx\n", + "from gt4py.next import float64, neighbor_sum, where" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3486fbec", + "metadata": {}, + "outputs": [], + "source": [ + "K = gtx.Dimension(\"K\", kind=gtx.DimensionKind.VERTICAL)" + ] + }, + { + "cell_type": "markdown", + "id": "cca4ec91", + "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": null, + "id": "ffc6cc17", + "metadata": {}, + "outputs": [], + "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": "3e48a4c3", + "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": "2c9bfe6f", + "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": null, + "id": "691bce93", + "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": null, + "id": "c8844650", + "metadata": {}, + "outputs": [], + "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": "80924e30", + "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": "f6b488d5", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "jupytext": { + "formats": "ipynb,md:myst" + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 626b1bdbb36853c811b4768ce92a7938093a10f0 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 23 Nov 2023 07:52:17 +0000 Subject: [PATCH 3/4] exercise 1 --- .../next/exercises/1_simple_addition.ipynb | 139 ++++++++++++++++++ 1 file changed, 139 insertions(+) create mode 100644 docs/user/next/exercises/1_simple_addition.ipynb diff --git a/docs/user/next/exercises/1_simple_addition.ipynb b/docs/user/next/exercises/1_simple_addition.ipynb new file mode 100644 index 0000000000..918e72b084 --- /dev/null +++ b/docs/user/next/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 +} From 56e9d729c4283847ed11752d602f43d241a6f9de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Thu, 23 Nov 2023 08:55:14 +0100 Subject: [PATCH 4/4] fix path to images folder --- .../2_divergence_exercise_solution.ipynb | 466 +++++++++--------- .../3_gradient_exercise_solution.ipynb | 322 ++++++------ .../exercises/4_curl_exercise_solution.ipynb | 336 +++++++------ 3 files changed, 580 insertions(+), 544 deletions(-) diff --git a/docs/user/next/exercises/2_divergence_exercise_solution.ipynb b/docs/user/next/exercises/2_divergence_exercise_solution.ipynb index 8422b8dabd..c41bf6c282 100644 --- a/docs/user/next/exercises/2_divergence_exercise_solution.ipynb +++ b/docs/user/next/exercises/2_divergence_exercise_solution.ipynb @@ -1,230 +1,242 @@ { - "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", - "![](../divergence_picture.png \"Divergence\")\n", - "\n", - "And the equation:\n", - "\n", - "![](../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", - "![](../edge_orientation.png \"Edge Orientation\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4eba62c1", - "metadata": {}, - "outputs": [], - "source": [ - "from helpers import *" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "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": null, - "id": "8fc6416f", - "metadata": {}, - "outputs": [], - "source": [ - "@gtx.field_operator\n", - "def divergence(\n", - " u: gtx.Field[[E], float],\n", - " v: gtx.Field[[E], float],\n", - " nx: gtx.Field[[E], float],\n", - " ny: gtx.Field[[E], float],\n", - " L: gtx.Field[[E], float],\n", - " A: gtx.Field[[C], float],\n", - " edge_orientation: gtx.Field[[C, C2EDim], float],\n", - ") -> gtx.Field[[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": null, - "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": null, - "id": "bbcb9bf5", - "metadata": {}, - "outputs": [], - "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", - "```" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - } + "cells": [ + { + "cell_type": "markdown", + "id": "c841c53b", + "metadata": {}, + "source": [ + "# 3. Divergence" + ] }, - "nbformat": 4, - "nbformat_minor": 5 + { + "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": null, + "id": "4eba62c1", + "metadata": {}, + "outputs": [], + "source": [ + "from helpers import *" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": null, + "id": "8fc6416f", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def divergence(\n", + " u: gtx.Field[[E], float],\n", + " v: gtx.Field[[E], float],\n", + " nx: gtx.Field[[E], float],\n", + " ny: gtx.Field[[E], float],\n", + " L: gtx.Field[[E], float],\n", + " A: gtx.Field[[C], float],\n", + " edge_orientation: gtx.Field[[C, C2EDim], float],\n", + ") -> gtx.Field[[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": null, + "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": null, + "id": "bbcb9bf5", + "metadata": {}, + "outputs": [], + "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", + "```" + ] + } + ], + "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" + } + }, + "nbformat": 4, + "nbformat_minor": 5 } diff --git a/docs/user/next/exercises/3_gradient_exercise_solution.ipynb b/docs/user/next/exercises/3_gradient_exercise_solution.ipynb index 69e587a1d6..2b651b2297 100644 --- a/docs/user/next/exercises/3_gradient_exercise_solution.ipynb +++ b/docs/user/next/exercises/3_gradient_exercise_solution.ipynb @@ -1,158 +1,170 @@ { - "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", - "![](../gradient_picture.png \"Divergence\")\n", - "\n", - "![](../gradient_formula.png \"Divergence\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1bddcf59", - "metadata": {}, - "outputs": [], - "source": [ - "from helpers import *\n", - "\n", - "\n", - "import gt4py.next as gtx" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "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": null, - "id": "afa80e49", - "metadata": {}, - "outputs": [], - "source": [ - "@gtx.field_operator\n", - "def gradient(\n", - " f: gtx.Field[[E], float],\n", - " nx: gtx.Field[[E], float],\n", - " ny: gtx.Field[[E], float],\n", - " L: gtx.Field[[E], float],\n", - " A: gtx.Field[[C], float],\n", - " edge_orientation: gtx.Field[[C, C2EDim], float],\n", - ") -> gtx.tuple[gtx.Field[[C], float], gtx.Field[[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": null, - "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": null, - "id": "9267fe99", - "metadata": {}, - "outputs": [], - "source": [ - "test_gradient()\n", - "print(\"Test successful\")" - ] - } - ], - "metadata": { - "jupytext": { - "formats": "ipynb,md:myst" - }, - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - } + "cells": [ + { + "cell_type": "markdown", + "id": "8cfd1e20", + "metadata": {}, + "source": [ + "# 4. Gradient" + ] }, - "nbformat": 4, - "nbformat_minor": 5 + { + "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": null, + "id": "1bddcf59", + "metadata": {}, + "outputs": [], + "source": [ + "from helpers import *\n", + "\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": null, + "id": "afa80e49", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def gradient(\n", + " f: gtx.Field[[E], float],\n", + " nx: gtx.Field[[E], float],\n", + " ny: gtx.Field[[E], float],\n", + " L: gtx.Field[[E], float],\n", + " A: gtx.Field[[C], float],\n", + " edge_orientation: gtx.Field[[C, C2EDim], float],\n", + ") -> gtx.tuple[gtx.Field[[C], float], gtx.Field[[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": null, + "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": null, + "id": "9267fe99", + "metadata": {}, + "outputs": [], + "source": [ + "test_gradient()\n", + "print(\"Test successful\")" + ] + } + ], + "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.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 } diff --git a/docs/user/next/exercises/4_curl_exercise_solution.ipynb b/docs/user/next/exercises/4_curl_exercise_solution.ipynb index 1211ba6a94..48abb4a6ff 100644 --- a/docs/user/next/exercises/4_curl_exercise_solution.ipynb +++ b/docs/user/next/exercises/4_curl_exercise_solution.ipynb @@ -1,165 +1,177 @@ { - "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", - "![](../curl_picture.png \"Divergence\")\n", - "\n", - "\n", - "![](../curl_formula.png \"Divergence\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1c1af88b", - "metadata": {}, - "outputs": [], - "source": [ - "from helpers import *\n", - "\n", - "import gt4py.next as gtx" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "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": null, - "id": "0925bdb0", - "metadata": {}, - "outputs": [], - "source": [ - "@gtx.field_operator\n", - "def curl(\n", - " u: gtx.Field[[E], float],\n", - " v: gtx.Field[[E], float],\n", - " nx: gtx.Field[[E], float],\n", - " ny: gtx.Field[[E], float],\n", - " dualL: gtx.Field[[E], float],\n", - " dualA: gtx.Field[[V], float],\n", - " edge_orientation: gtx.Field[[V, V2EDim], float],\n", - ") -> gtx.Field[[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": null, - "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": null, - "id": "ae651445", - "metadata": {}, - "outputs": [], - "source": [ - "test_curl()\n", - "print(\"Test successful\")" - ] - } - ], - "metadata": { - "jupytext": { - "formats": "ipynb,md:myst" - }, - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - } + "cells": [ + { + "cell_type": "markdown", + "id": "9e5abeda", + "metadata": {}, + "source": [ + "# 5. Curl" + ] }, - "nbformat": 4, - "nbformat_minor": 5 + { + "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": null, + "id": "1c1af88b", + "metadata": {}, + "outputs": [], + "source": [ + "from helpers import *\n", + "\n", + "import gt4py.next as gtx" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": null, + "id": "0925bdb0", + "metadata": {}, + "outputs": [], + "source": [ + "@gtx.field_operator\n", + "def curl(\n", + " u: gtx.Field[[E], float],\n", + " v: gtx.Field[[E], float],\n", + " nx: gtx.Field[[E], float],\n", + " ny: gtx.Field[[E], float],\n", + " dualL: gtx.Field[[E], float],\n", + " dualA: gtx.Field[[V], float],\n", + " edge_orientation: gtx.Field[[V, V2EDim], float],\n", + ") -> gtx.Field[[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": null, + "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": null, + "id": "ae651445", + "metadata": {}, + "outputs": [], + "source": [ + "test_curl()\n", + "print(\"Test successful\")" + ] + } + ], + "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.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 }