From 01b2bf893f307a99514c82dd81d512a0a2bee9a7 Mon Sep 17 00:00:00 2001 From: nfarabullini Date: Tue, 12 Mar 2024 11:27:22 +0100 Subject: [PATCH] edits to quickstart --- docs/user/next/QuickstartGuide.md | 62 ++++++++++++++++--------------- 1 file changed, 32 insertions(+), 30 deletions(-) diff --git a/docs/user/next/QuickstartGuide.md b/docs/user/next/QuickstartGuide.md index dc70f804fd..2a80301bc3 100644 --- a/docs/user/next/QuickstartGuide.md +++ b/docs/user/next/QuickstartGuide.md @@ -47,6 +47,7 @@ import numpy as np import gt4py.next as gtx from gt4py.next import float64, neighbor_sum, where +from gt4py.next.common import Dims ``` #### Fields @@ -91,11 +92,12 @@ 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[[CellDim, KDim], float64], - b: gtx.Field[[CellDim, KDim], float64]) -> gtx.Field[[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 ``` +\_Note: for now `Dims` is not mandatory, hence this type hint is also accepted: `gtx.Field[[CellDim, KDim], float64]` You can call field operators from [programs](#Programs), other field operators, or directly. The code snippet below shows a direct call, in which case you have to supply two additional arguments: `out`, which is a field to write the return value to, and `offset_provider`, which is left empty for now. The result of the field operator is a field with all entries equal to 5, but for brevity, only the average and the standard deviation of the entries are printed: ```{code-cell} ipython3 @@ -115,9 +117,9 @@ This example program below calls the above elementwise addition field operator t ```{code-cell} ipython3 @gtx.program -def run_add(a : gtx.Field[[CellDim, KDim], float64], - b : gtx.Field[[CellDim, KDim], float64], - result : gtx.Field[[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) ``` @@ -247,11 +249,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[[CellDim], float64]) -> gtx.Field[[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[[CellDim], float64], out : gtx.Field[[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}) @@ -273,12 +275,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[[CellDim], float64]) -> gtx.Field[[EdgeDim], float64]: - # type of cells(E2C) is gtx.Field[[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[CellDim, E2CDim], float64] return neighbor_sum(cells(E2C), axis=E2CDim) @gtx.program -def run_sum_adjacent_cells(cells : gtx.Field[[CellDim], float64], out : gtx.Field[[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}) @@ -302,7 +304,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[[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: @@ -312,8 +314,8 @@ result_where = gtx.as_field([CellDim, KDim], np.zeros(shape=grid_shape)) b = 6.0 @gtx.field_operator -def conditional(mask: gtx.Field[[CellDim, KDim], bool], a: gtx.Field[[CellDim, KDim], float64], b: float -) -> gtx.Field[[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={}) @@ -329,13 +331,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[[CellDim, KDim], bool], a: gtx.Field[[CellDim, KDim], float64], b: float -) -> tuple[gtx.Field[[CellDim, KDim], float64], gtx.Field[[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[[CellDim, KDim], bool], a: gtx.Field[[CellDim, KDim], float64], b: float, -result_1: gtx.Field[[CellDim, KDim], float64], result_2: gtx.Field[[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)) @@ -360,17 +362,17 @@ result_2 = gtx.as_field([CellDim, KDim], np.zeros(shape=grid_shape)) @gtx.field_operator def _conditional_tuple_nested( - mask: gtx.Field[[CellDim, KDim], bool], a: gtx.Field[[CellDim, KDim], float64], b: gtx.Field[[CellDim, KDim], float64], c: gtx.Field[[CellDim, KDim], float64], d: gtx.Field[[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[[CellDim, KDim], float64], gtx.Field[[CellDim, KDim], float64]], - tuple[gtx.Field[[CellDim, KDim], float64], gtx.Field[[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[[CellDim, KDim], bool], a: gtx.Field[[CellDim, KDim], float64], b: gtx.Field[[CellDim, KDim], float64], c: gtx.Field[[CellDim, KDim], float64], d: gtx.Field[[CellDim, KDim], float64], - result_1: gtx.Field[[CellDim, KDim], float64], result_2: gtx.Field[[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))) @@ -425,9 +427,9 @@ The second lines first creates a temporary field using `edge_differences(C2E)`, ```{code-cell} ipython3 @gtx.field_operator -def pseudo_lap(cells : gtx.Field[[CellDim], float64], - edge_weights : gtx.Field[[CellDim, C2EDim], float64]) -> gtx.Field[[CellDim], float64]: - edge_differences = cells(E2C[0]) - cells(E2C[1]) # type: gtx.Field[[EdgeDim], float64] +def pseudo_lap(cells : gtx.Field[Dims[CellDim], float64], + edge_weights : gtx.Field[Dims[CellDim, C2EDim], float64]) -> gtx.Field[Dims[CellDim], float64]: + edge_differences = cells(E2C[0]) - cells(E2C[1]) # type: gtx.Field[Dims[EdgeDim], float64] return neighbor_sum(edge_differences(C2E) * edge_weights, axis=C2EDim) ``` @@ -435,9 +437,9 @@ The program itself is just a shallow wrapper over the `pseudo_lap` field operato ```{code-cell} ipython3 @gtx.program -def run_pseudo_laplacian(cells : gtx.Field[[CellDim], float64], - edge_weights : gtx.Field[[CellDim, C2EDim], float64], - out : gtx.Field[[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,))) @@ -454,7 +456,7 @@ 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[[CellDim], float64], - edge_weights : gtx.Field[[CellDim, C2EDim], float64]) -> gtx.Field[[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) ```