GT4Py / NDSL / Serialbox Tutorials

examples/Fortran_porting/02_read_serialized_data_python.ipynb

@@ -0,0 +1,229 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## **Serialbox Tutorial : Incorporating Fortran Serialbox Data into Python**\n",
+    "\n",
+    "In the [previous notebook](./01.ipynb), we covered how to extract data from a Fortran code using Serialbox.  In this notebook, we'll cover how to read and incorporate those files within a Python code."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Notebook Requirements**\n",
+    "\n",
+    "- Python v3.11.x to v3.12.x\n",
+    "- [NOAA/NASA Domain Specific Language Middleware](https://github.com/NOAA-GFDL/NDSL)\n",
+    "- `ipykernel==6.1.0`\n",
+    "- [`ipython_genutils`](https://pypi.org/project/ipython_genutils/)\n",
+    "\n",
+    "This notebook assumes that the code from the [previous notebook](./01.ipynb) was run, and the serialized data from Fortran was written out."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Importing Fortran Serialbox Data From Example 1 into Python** ###\n",
+    "\n",
+    "We'll step through importing Serialbox data [created in Fortran](./01.ipynb#Serialbox-Example-1) into Python to test a Python port of `FILLQ2ZERO1`.  Importing Serialbox data into Python essentially comes from opening a file via a \"serializer\" object denoted by a particular Serialbox initialization prefix (see [Serialbox directive calls in Fortran code](./01.ipynb#Serialbox-directive-calls-in-Fortran-code)) and stepping through the savepoints within the \"serializer\" object to read the data.  This is done by the following Python calls assuming that the imported `serialbox` package is referenced via `ser`.\n",
+    "\n",
+    "- `ser.Serializer(ser.OpenModeKind.Read,\"<Path to Serialbox Data>\", \"<Name of prefix used during Serialbox initialization>\")` : This function call creates a \"serializer\" object that will read Serialbox files within a declared path and reference data from a particular Serialbox initialization prefix.\n",
+    "\n",
+    "- `serializer.savepoint_list()` : Using a \"serializer\" object called `serializer`, this function call creates a list of Serialbox savepoints\n",
+    "\n",
+    "- `serializer.read(\"<Serialbox variable name>\", <Savepoint from savepoint list>)` : Using a \"serializer\" object called `serializer`, this function call will look for the specified Serialbox variable name from the savepoint list and output that variable.\n",
+    "\n",
+    "Below is a Python example that uses these three calls to import the [Example 1](./01.ipynb#Serialbox-Example-1) Fortran data into Python.  You can check to see that the summation of the arrays with Python match closely with the [values presented in Fortran](./01.ipynb#Building-and-Running-Fortran-code-with-Serialbox-library)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Sum of Qin_out =  58.74463748931885\n",
+      "Sum of mass =  62.169867515563965\n",
+      "Sum of fq_out =  0.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "import sys\n",
+    "# Appends the Serialbox python path to PYTHONPATH.  If needed, change to appropriate path containing serialbox installation\n",
+    "sys.path.append('/home/ckung/Documents/Code/SMT-Nebulae/sw_stack_path/install/serialbox/python')\n",
+    "import serialbox as ser\n",
+    "import numpy as np\n",
+    "\n",
+    "# If needed, change the path in second parameter of ser.Serializer to appropriate path that contains Fortran data via Serialbox from 01.ipynb\n",
+    "serializer = ser.Serializer(ser.OpenModeKind.Read,\"./Fortran/sb/\",\"FILLQ2ZERO_InOut\")\n",
+    "\n",
+    "savepoints = serializer.savepoint_list()\n",
+    "\n",
+    "Qin_out = serializer.read(\"q_in\", savepoints[0])\n",
+    "mass    = serializer.read(\"m_in\", savepoints[0])\n",
+    "fq_out  = serializer.read(\"fq_in\", savepoints[0])\n",
+    "\n",
+    "print('Sum of Qin_out = ', sum(sum(sum(Qin_out))))\n",
+    "print('Sum of mass = ', sum(sum(sum(mass))))\n",
+    "print('Sum of fq_out = ', sum(sum(fq_out)))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, we'll create a rudimentary port of `fillq2zero1` and test whether or not it computes properly by comparing the output arrays `Qin_out` and `fq_out` to the corresonding arrays created from Fortran, which are retrieved using `serializer.read()`.  In this example, the comparison between the Fortran and Python data is performed using `np.allclose`; however, note that the proper metric of comparison will depend on the application.  We'll see that `np.allclose()` will report `True` for both the `Qin_out` and `fq_out` array comparisons. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Sum of Qin_out =  58.78519535064697\n",
+      "Sum of fq_out =  0.25218469463288784\n",
+      "True\n",
+      "True\n"
+     ]
+    }
+   ],
+   "source": [
+    "def fillq2zero1(Q, MASS, FILLQ):\n",
+    "    IM = Q.shape[0]\n",
+    "    JM = Q.shape[1]\n",
+    "    LM = Q.shape[2]\n",
+    "\n",
+    "    TPW = np.sum(Q*MASS,2)\n",
+    "    for J in range(JM):\n",
+    "        for I in range(IM):\n",
+    "            NEGTPW = 0.\n",
+    "            for L in range(LM):\n",
+    "                if(Q[I,J,L] < 0.0):\n",
+    "                    NEGTPW = NEGTPW + (Q[I,J,L]*MASS[I,J,L])\n",
+    "                    Q[I,J,L] = 0.0\n",
+    "            for L in range(LM):\n",
+    "                if(Q[I,J,L] >= 0.0):\n",
+    "                    Q[I,J,L] = Q[I,J,L]*(1.0 + NEGTPW/(TPW[I,J]-NEGTPW))\n",
+    "            FILLQ[I,J] = -NEGTPW\n",
+    "            \n",
+    "fillq2zero1(Qin_out,mass,fq_out)\n",
+    "\n",
+    "print('Sum of Qin_out = ', sum(sum(sum(Qin_out))))\n",
+    "print('Sum of fq_out = ', sum(sum(fq_out)))\n",
+    "\n",
+    "Qin_out_ref = serializer.read(\"q_out\", savepoints[0])\n",
+    "mass_ref    = serializer.read(\"m_out\", savepoints[0])\n",
+    "fq_out_ref  = serializer.read(\"fq_out\", savepoints[0])\n",
+    "\n",
+    "print(np.allclose(Qin_out,Qin_out_ref))\n",
+    "print(np.allclose(fq_out,fq_out_ref))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### **Importing Fortran Data from Example 2 into Python : Looping Regions** ###\n",
+    "\n",
+    "In [Example 2](./01.ipynb#Serialbox-Example-2), Serialbox was set up to record data within a looping region.  This results in a larger list of savepoints that we can step through in Python to recreating the looping process done in Fortran.  The code below replicates the looping of `FILLQ2ZERO1` and reads multiple savepoints to intialize the data and compare outputs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Current savepoint =  sp1 {\"timestep\": 1}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 2}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 3}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 4}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 5}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 6}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 7}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 8}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 9}\n",
+      "True\n",
+      "True\n",
+      "Current savepoint =  sp1 {\"timestep\": 10}\n",
+      "True\n",
+      "True\n"
+     ]
+    }
+   ],
+   "source": [
+    "# If needed, change the path in second parameter of ser.Serializer to appropriate path that contains Fortran data via Serialbox from 01.ipynb\n",
+    "serializer = ser.Serializer(ser.OpenModeKind.Read,\"./Fortran_ts/sb/\",\"FILLQ2ZERO_InOut\")\n",
+    "\n",
+    "savepoints = serializer.savepoint_list()\n",
+    "\n",
+    "for currentSavepoint in savepoints:\n",
+    "    Qin_out = serializer.read(\"q_in\", currentSavepoint)\n",
+    "    mass    = serializer.read(\"m_in\", currentSavepoint)\n",
+    "    fq_out  = serializer.read(\"fq_in\", currentSavepoint)\n",
+    "\n",
+    "    fillq2zero1(Qin_out,mass,fq_out)\n",
+    "\n",
+    "    Qin_out_ref = serializer.read(\"q_out\", currentSavepoint)\n",
+    "    mass_ref    = serializer.read(\"m_out\", currentSavepoint)\n",
+    "    fq_out_ref  = serializer.read(\"fq_out\", currentSavepoint)\n",
+    "\n",
+    "    print('Current savepoint = ', currentSavepoint)\n",
+    "    print(np.allclose(Qin_out,Qin_out_ref))\n",
+    "    print(np.allclose(fq_out,fq_out_ref))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "gt4py_jupyter",
+   "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.11.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/Fortran_porting/03_translate_test.ipynb

@@ -0,0 +1,51 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Make a translate test:\n",
+    "- Make the savepoint by following the naming convention `XYZ-In` and `XYZ-Out`\n",
+    "- Write the Translate class, following the naming convention `TranslateXYZ`\n",
+    "    - Inherit from the base class (default `TranslateFortranData2Py`)\n",
+    "    - `compute_func` is called, override that with runtime of the ported code\n",
+    "    - definte `self.in_vars` and `self.out_vars` dimensions, names and serialize override to \n",
+    "    make the bridge between the saved data and the calling signature"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We made an example for `FILLQZERO1`. You can run it with"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "vscode": {
+     "languageId": "shellscript"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "%%bash\n",
+    "\n",
+    "python -m pytest \\ \n",
+    "       --data_path=./test_data/example_data/Fortran/FILLQ2ZERO \\\n",
+    "       --grid=default \\\n",
+    "       --backend=numpy \\\n",
+    "       --layout=tile \\\n",
+    "       ./examples/Fortran_porting/tests/"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/Fortran_porting/tests/conftest.py

@@ -0,0 +1,7 @@
+import ndsl.stencils.testing.conftest
+from ndsl.stencils.testing.conftest import *  # noqa: F403,F401
+
+import savepoint
+
+# Point to an __init__.py where all the TestX are improted
+ndsl.stencils.testing.conftest.translate = savepoint  # type: ignore

examples/Fortran_porting/tests/savepoint/__init__.py

@@ -0,0 +1 @@
+from .translate_fillqzero import TranslateFILLQ2ZERO1

examples/Fortran_porting/tests/savepoint/translate_fillqzero.py

@@ -0,0 +1,87 @@
+import numpy as np
+from ndsl import StencilFactory, QuantityFactory, orchestrate, DaceConfig
+from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ
+from ndsl.constants import X_DIM, Y_DIM, Z_DIM, X_INTERFACE_DIM, Y_INTERFACE_DIM
+from ndsl import Namelist, StencilFactory
+from ndsl.stencils.testing.translate import TranslateFortranData2Py
+from gt4py.cartesian.gtscript import computation, FORWARD, PARALLEL, interval
+from typing import Tuple
+
+def _make_range(offset: int, domain: int):        
+    return range(offset, offset+domain)
+
+def fillq2zero1_plain_python(
+    domain: Tuple[int, ...],
+    offset: Tuple[int, ...],
+    q: FloatField,
+    mass:FloatField,
+    fillq:FloatField
+):
+    tpw = np.sum(q*mass,2)
+    for J in _make_range(offset[1], domain[1]):
+        for I in _make_range(offset[0], domain[0]):
+            neg_tpw = 0.
+            for L in _make_range(offset[2], domain[2]):
+                if(q[I,J,L] < 0.0):
+                    neg_tpw = neg_tpw + (q[I,J,L]*mass[I,J,L])
+                    q[I,J,L] = 0.0
+            for L in _make_range(offset[2], domain[2]):
+                if(q[I,J,L] >= 0.0):
+                    q[I,J,L] = q[I,J,L]*(1.0 + neg_tpw/(tpw[I,J]-neg_tpw))
+            fillq[I,J] = -neg_tpw
+
+class FillQZero:
+    def __init__(self, stencil_factory: StencilFactory, quantity_factory: QuantityFactory):
+        orchestrate(
+            obj=self,
+            config=(
+                stencil_factory.config.dace_config or
+                DaceConfig(communicator=None, backend=stencil_factory.backend))
+        )
+        self._domain = quantity_factory.sizer.get_extent([X_DIM, Y_DIM, Z_DIM])
+        self._offset = quantity_factory.sizer.get_origin([X_DIM, Y_DIM, Z_DIM])
+
+    def __call__(
+        self,
+        q: FloatField,
+        mass: FloatField,
+        fillq: FloatField
+    ):
+        fillq2zero1_plain_python(
+            domain=self._domain,
+            offset=self._offset,
+            q=q,
+            mass=mass,
+            fillq=fillq)
+
+
+class TranslateFILLQ2ZERO1(TranslateFortranData2Py):
+    def __init__(
+        self,
+        grid,
+        namelist: Namelist,
+        stencil_factory: StencilFactory,
+    ):
+        super().__init__(grid, stencil_factory)
+        self.max_error=1e-7
+        self.compute_func = FillQZero(   # type: ignore
+            self.stencil_factory,
+            self.grid.quantity_factory,
+        )
+
+        fillq_info = self.grid.compute_dict()
+        fillq_info["serialname"] = "fq"
+        self.in_vars["data_vars"] = {
+            "mass": self.grid.compute_dict(),
+            "q": self.grid.compute_dict(),
+            "fillq": fillq_info,
+        }
+        self.out_vars = {
+            "fillq": fillq_info,
+            "q": self.grid.compute_dict(),
+        }
+
+
+    def compute_from_storage(self, inputs):
+        self.compute_func(**inputs)
+        return inputs

examples/Fortran_porting/tests/test_translate.py

@@ -0,0 +1 @@
+from ndsl.stencils.testing.test_translate import *  # noqa: F403,F401