Improved docs for NDSL examples

examples/NDSL/03_orchestration_basics.ipynb

@@ -19,29 +19,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<script src=\"https://spcl.github.io/dace-webclient/dist/sdfv.js\"></script>\n"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "2024-05-29 13:16:33|INFO|rank 0|ndsl.logging:Constant selected: ConstantVersions.GFS\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import numpy as np\n",
     "from gt4py.cartesian.gtscript import (\n",
@@ -57,8 +37,7 @@
     ")\n",
     "from ndsl.constants import X_DIM, Y_DIM, Z_DIM\n",
     "from ndsl.dsl.typing import FloatField, Float\n",
-    "\n",
-    "from orch_boilerplate import get_one_tile_factory_orchestrated"
+    "from ndsl.boilerplate import get_factories_single_tile_orchestrated_cpu"
    ]
   },
   {
@@ -70,14 +49,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "def localsum_stencil(\n",
-    "    field: FloatField,  # type: ignore\n",
+    "    field: FloatField,   # type: ignore\n",
     "    result: FloatField,  # type: ignore\n",
-    "    weight: Float,  # type: ignore\n",
+    "    weight: Float,       # type: ignore\n",
     "):\n",
     "    with computation(PARALLEL), interval(...):\n",
     "        result = weight * (\n",
@@ -94,7 +73,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -119,8 +98,8 @@
     "        self._n_halo = quantity_factory.sizer.n_halo\n",
     "\n",
     "    def __call__(self, in_field: FloatField, out_result: FloatField) -> None:\n",
-    "        self._local_sum(in_field, out_result, 2.0) # GT4Py Stencil\n",
-    "        tmp_field = out_result[:, :, :] + 2        # Regular Python code\n",
+    "        self._local_sum(in_field, out_result, 2.0)  # GT4Py Stencil\n",
+    "        tmp_field = out_result[:, :, :] + 2         # Regular Python code\n",
     "        self._local_sum(tmp_field, out_result, 2.0) # GT4Py Stencil"
    ]
   },
@@ -133,51 +112,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "2024-05-29 13:16:33|INFO|rank 0|ndsl.logging:[DaCeOrchestration.BuildAndRun] Rank 0 reading/writing cache .gt_cache_FV3_A\n",
-      "2024-05-29 13:16:33|INFO|rank 0|ndsl.logging:Building DaCe orchestration\n",
-      "Inlined 2 SDFGs.\n",
-      "Fused 4 states.\n",
-      "Inferred 2 optional arrays.\n",
-      "SDFG 0: Eliminated 1 arrays: {'out_result_0'}.\n",
-      "Fused 2 states.\n",
-      "Inferred 4 optional arrays.\n",
-      "Inlined 2 SDFGs.\n",
-      "2024-05-29 13:16:34|INFO|rank 0|ndsl.logging:[DaCeOrchestration.BuildAndRun] LocalSum___call__:\n",
-      "StorageType.Default:\n",
-      "  Alloc ref 0.01 mb\n",
-      "  Alloc unref 0.00 mb\n",
-      "  Pooled 0.00 mb\n",
-      "  Top lvl alloc: 0.01mb\n",
-      "\n",
-      "[DaCe Config] Rank 0 loading SDFG /home/ckung/Documents/Code/SMT-Nebulae-Tutorial/tutorial/NDSL/.gt_cache_FV3_A/dacecache/LocalSum___call__\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# ----- Driver ----- #\n",
     "\n",
     "if __name__ == \"__main__\":\n",
     "    # Settings\n",
-    "    backend = \"dace:cpu\"\n",
     "    dtype = np.float64\n",
     "    origin = (0, 0, 0)\n",
     "    rebuild = True\n",
     "    tile_size = (3, 3, 3)\n",
     "\n",
     "    # Setup\n",
-    "    stencil_factory, qty_factory = get_one_tile_factory_orchestrated(\n",
+    "    stencil_factory, qty_factory = get_factories_single_tile_orchestrated_cpu(\n",
     "        nx=tile_size[0],\n",
     "        ny=tile_size[1],\n",
     "        nz=tile_size[2],\n",
     "        nhalo=2,\n",
-    "        backend=backend,\n",
     "    )\n",
     "    local_sum = LocalSum(stencil_factory, qty_factory)\n",
     "\n",
@@ -206,7 +159,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

examples/NDSL/README.md

@@ -0,0 +1,14 @@
+# NDSL examples
+
+This folder contains a couple of Jupyter notebooks with the following examples:
+
+- Getting started with GT4Py: [GT4Py basics](./01_gt4py_basics.ipynb)
+- Getting started with NDSL middleware: [NDSL basics](./02_NDSL_basics.ipynb)
+- Combining stencil and non-stencil code with DaCe: [Orchestration basics](./03_orchestration_basics.ipynb)
+
+## Quickstart
+
+1. Make sure you fulfill the [requirements of NDSL](../../README.md#quickstart), e.g. python version.
+2. Create a virtual environment `python -m venv .venv`, and activate it `source .venv/bin/activate`.
+3. Install NDSL into your environment `pip install ../../[demos]`.
+4. In VSCode, install the Jupyter extension, select your virtual environment as kernel and run the notebooks.

ndsl/boilerplate.py

@@ -82,7 +82,7 @@ def _get_factories(
 def get_factories_single_tile_orchestrated_cpu(
     nx, ny, nz, nhalo
 ) -> Tuple[StencilFactory, QuantityFactory]:
-    """Build a Stencil & Quantity factory for orchestrated CPU, on a single tile toplogy."""
+    """Build a Stencil & Quantity factory for orchestrated CPU, on a single tile topology."""
     return _get_factories(
         nx=nx,
         ny=ny,
@@ -97,7 +97,7 @@ def get_factories_single_tile_orchestrated_cpu(
 def get_factories_single_tile_numpy(
     nx, ny, nz, nhalo
 ) -> Tuple[StencilFactory, QuantityFactory]:
-    """Build a Stencil & Quantity factory for Numpy, on a single tile toplogy."""
+    """Build a Stencil & Quantity factory for Numpy, on a single tile topology."""
     return _get_factories(
         nx=nx,
         ny=ny,

ndsl/quantity.py

@@ -2,6 +2,7 @@
 import warnings
 from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union, cast
 
+import matplotlib.pyplot as plt
 import numpy as np
 
 import ndsl.constants as constants
@@ -593,6 +594,22 @@ def transpose(
         transposed._attrs = self._attrs
         return transposed
 
+    def plot_k_level(self, k_index=0):
+        field = self.data
+        print(
+            "Min and max values:",
+            field[:, :, k_index].min(),
+            field[:, :, k_index].max(),
+        )
+        plt.xlabel("I")
+        plt.ylabel("J")
+
+        im = plt.imshow(field[:, :, k_index].transpose(), origin="lower")
+
+        plt.colorbar(im)
+        plt.title("Plot at K = " + str(k_index))
+        plt.show()
+
 
 def transpose_sequence(sequence, order):
     return sequence.__class__(sequence[i] for i in order)

setup.py

@@ -13,8 +13,13 @@ def local_pkg(name: str, relative_path: str) -> str:
 
 test_requirements = ["pytest", "pytest-subtests", "coverage"]
 develop_requirements = test_requirements + ["pre-commit"]
+demos_requirements = ["ipython", "ipykernel"]
 
-extras_requires = {"test": test_requirements, "develop": develop_requirements}
+extras_requires = {
+    "test": test_requirements,
+    "develop": develop_requirements,
+    "demos": demos_requirements,
+}
 
 requirements: List[str] = [
     local_pkg("gt4py", "external/gt4py"),
@@ -29,6 +34,7 @@ def local_pkg(name: str, relative_path: str) -> str:
     "h5netcdf",  # for xarray
     "dask",  # for xarray
     "numpy==1.26.4",
+    "matplotlib",  # for plotting in boilerplate
 ]