Change gt4py imports in examples and tests

README.md

@@ -18,7 +18,7 @@ NDSL submodules `gt4py` and `dace` to point to vetted versions, use `git clone -
 NDSL is __NOT__ available on `pypi`. Installation of the package has to be local, via `pip install ./NDSL` (`-e` supported). The packages has a few options:
 
 - `ndsl[test]`: installs the test packages (based on `pytest`)
-- `ndsl[demos]`: installs extra requirements to run [NDSL exmpales](./examples/NDSL/)
+- `ndsl[demos]`: installs extra requirements to run [NDSL examples](./examples/NDSL/)
 - `ndsl[docs]`: installs extra requirements to build the docs
 - `ndsl[develop]`: installs tools for development, docs, and tests.
 

examples/NDSL/01_gt4py_basics.ipynb

@@ -40,7 +40,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gt4py.cartesian.gtscript import PARALLEL, computation, interval, stencil\n",
+    "from ndsl.dsl.gt4py import PARALLEL, computation, interval, stencil\n",
     "from ndsl.dsl.typing import FloatField\n",
     "from ndsl.quantity import Quantity\n",
     "import numpy as np"
@@ -109,14 +109,14 @@
     "\n",
     "We see that this stencil does not contain any explicit loops.  As mentioned above in the notebook, GT4Py has a particular computation policy that implicitly executes in parallel within an `IJ` plane and is user defined in the `K` interval.  This execution policy in the `K` interval is dictated by the `computation` and `interval` keywords.  \n",
     "\n",
-    "- `with computation(PARALLEL)` means that there's no order preference to executing the `K` interval.  This also means that the `K` interval can be computed in parallel to potentially gain performace if computational resources are available.\n",
+    "- `with computation(PARALLEL)` means that there's no order preference to executing the `K` interval.  This also means that the `K` interval can be computed in parallel to potentially gain performance if computational resources are available.\n",
     "\n",
     "- `interval(...)` means that the entire `K` interval is executed.  Instead of `(...)`, more specific intervals can be specified using a tuple of two integers.  For example... \n",
     "\n",
     "    - `interval(0,2)` : The interval `K` = 0 to 1 is executed.\n",
     "    - `interval(0,-1)` : The interval `K` = 0 to N-2 (where N is the size of `K`) is executed.\n",
     "\n",
-    "The decorator `@stencil(backend=backend)` (Note: `stencil` comes from the package `gt4py.cartesian.gtscript`) converts `copy_stencil` to use the specified `backend` to \"compile\" the stencil.  `stencil` can also be a function call to create a stencil object."
+    "The decorator `@stencil(backend=backend)` (Note: `stencil` comes from the package `ndsl.dsl.gt4py`) converts `copy_stencil` to use the specified `backend` to \"compile\" the stencil.  `stencil` can also be a function call to create a stencil object."
    ]
   },
   {
@@ -269,7 +269,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gt4py.cartesian.gtscript import FORWARD, BACKWARD\n",
+    "from ndsl.dsl.gt4py import FORWARD, BACKWARD\n",
     "\n",
     "nx = 5\n",
     "ny = 5\n",
@@ -470,7 +470,7 @@
     "\n",
     "GT4Py also has the capability to create functions in order to better organize code.  The main difference between a GT4Py function call and a GT4Py stencil is that a function does not (and cannot) contain the keywords `computation` and `interval`.  However, array index referencing within a GT4py function is the same as in a GT4Py stencil.\n",
     "\n",
-    "GT4Py functions can be created by using the decorator `function` (Note: `function` originates from the package `gt4py.cartesian.gtscript`)."
+    "GT4Py functions can be created by using the decorator `function` (Note: `function` originates from the package `ndsl.dsl.gt4py`)."
    ]
   },
   {
@@ -479,7 +479,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gt4py.cartesian.gtscript import function\n",
+    "from ndsl.dsl.gt4py import function\n",
     "\n",
     "@function\n",
     "def plus_one(field: FloatField):\n",

examples/NDSL/02_NDSL_basics.ipynb

@@ -26,14 +26,14 @@
    "outputs": [],
    "source": [
     "from ndsl import StencilFactory\n",
-    "from ndsl.boilerplate import get_factories_single_tile_numpy\n",
+    "from ndsl.boilerplate import get_factories_single_tile\n",
     "\n",
     "nx = 6\n",
     "ny = 6\n",
     "nz = 1\n",
     "nhalo = 1\n",
     "\n",
-    "stencil_factory, _ = get_factories_single_tile_numpy(nx, ny, nz, nhalo)"
+    "stencil_factory, _ = get_factories_single_tile(nx, ny, nz, nhalo)"
    ]
   },
   {
@@ -59,8 +59,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from ndsl.dsl.gt4py import PARALLEL, computation, interval\n",
     "from ndsl.dsl.typing import FloatField\n",
-    "from gt4py.cartesian.gtscript import PARALLEL, computation, interval\n",
     "\n",
     "def copy_field_stencil(field_in: FloatField, field_out: FloatField):\n",
     "    with computation(PARALLEL), interval(...):\n",
@@ -150,7 +150,7 @@
     "print(\"Plotting qty_in at K = 0\")\n",
     "qty_in.plot_k_level(0)\n",
     "print(\"Plotting qty_out at K = 0\")\n",
-    "qty_out.plot_k_level(0)\n"
+    "qty_out.plot_k_level(0)"
    ]
   },
   {
@@ -189,7 +189,7 @@
     "\n",
     "The next example will create a stencil that takes a `Quantity` as an input, shift the input by 1 in the `-J` direction, and write it to an output `Quantity`.  This stencil is defined in `copy_field_offset_stencil`.\n",
     "\n",
-    "Note that in `copy_field_offset_stencil`, the shift in the J dimension is performed by referencing the `J` object from `gt4py.cartesian.gtscript` for simplicity.  This reference will apply the shift in J to the entire input domain.  Another way to perform the shift without referencing the `J` object is to write `[0,-1,0]` (assuming that the variable being modified is 3-dimensional) instead of `[J-1]`.\n",
+    "Note that in `copy_field_offset_stencil`, the shift in the `J` dimension is performed by referencing the `J` object from `ndsl.dsl.gt4py` for simplicity.  This reference will apply the shift in `J` to the entire input domain.  Another way to perform the shift without referencing the `J` object is to write `[0,-1,0]` (assuming that the variable being modified is 3-dimensional) instead of `[J-1]`.\n",
     "\n",
     "With the stencil in place, a class `CopyFieldOffset` is defined using the `StencilFactory` object and `copy_field_offset_stencil`.  The class is instantiated and demonstrated to shift `qty_in` by 1 in the J-dimension and write to `qty_out`."
    ]
@@ -200,7 +200,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gt4py.cartesian.gtscript import J\n",
+    "from ndsl.dsl.gt4py import J\n",
     "\n",
     "def copy_field_offset_stencil(field_in: FloatField, field_out: FloatField):\n",
     "    with computation(PARALLEL), interval(...):\n",
@@ -230,7 +230,7 @@
     "                   gt4py_backend=backend\n",
     "                   )\n",
     "\n",
-    "print(\"Initialize qty_out to zeros\")\n"
+    "print(\"Initialize qty_out to zeros\")"
    ]
   },
   {
@@ -251,7 +251,7 @@
    "source": [
     "### **Limits to offset : Cannot set offset outside of usable domain**\n",
     "\n",
-    "Note that when the copy offset by -1 in the j-direction is performed, the 'halo' region at J = 8 is copied over due to the `J` shift.  This means that there are limits to the shift amount since choosing a large shift amount may result in accessing a data region that does not exist.  The following example shows this by trying to perform a shift by -2 in the j-direction."
+    "Note that when the copy offset by `-1` in the `j`-direction is performed, the 'halo' region at `J = 8` is copied over due to the `J` shift. This means that there are limits to the shift amount since choosing a large shift amount may result in accessing a data region that does not exist. The following example shows this by trying to perform a shift by `-2` in the `j`-direction."
    ]
   },
   {
@@ -282,7 +282,7 @@
     "        \n",
     "copy_field_offset = CopyFieldOffset(stencil_factory)\n",
     "\n",
-    "copy_field_offset(qty_in, qty_out)\n"
+    "copy_field_offset(qty_in, qty_out)"
    ]
   },
   {
@@ -300,8 +300,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gt4py.cartesian.gtscript import J\n",
-    "\n",
     "def copy_field_offset_output_stencil(field_in: FloatField, field_out: FloatField):\n",
     "    with computation(PARALLEL), interval(...):\n",
     "        field_out[0,1,0] = field_in\n",
@@ -322,14 +320,13 @@
     "    ):\n",
     "        self._copy_field_offset_output(field_in, field_out)\n",
     "        \n",
-    "copy_field_offset_output = CopyFieldOffsetOutput(stencil_factory)\n",
-    "        "
+    "copy_field_offset_output = CopyFieldOffsetOutput(stencil_factory)"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "gt4py_jupyter",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },

examples/NDSL/03_orchestration_basics.ipynb

@@ -24,7 +24,7 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "from gt4py.cartesian.gtscript import (\n",
+    "from ndsl.dsl.gt4py import (\n",
     "    PARALLEL,\n",
     "    computation,\n",
     "    interval,\n",

tests/dsl/test_caches.py

@@ -1,5 +1,4 @@
 import pytest
-from gt4py.cartesian.gtscript import PARALLEL, Field, computation, interval
 from gt4py.storage import empty, ones
 
 from ndsl import (
@@ -12,6 +11,7 @@
 )
 from ndsl.comm.mpi import MPI
 from ndsl.dsl.dace.orchestration import orchestrate
+from ndsl.dsl.gt4py import PARALLEL, Field, computation, interval
 
 
 def _make_storage(

tests/dsl/test_skip_passes.py

@@ -5,7 +5,6 @@
     HorizontalExecutionMerging,
 )
 from gt4py.cartesian.gtc.passes.oir_pipeline import DefaultPipeline
-from gt4py.cartesian.gtscript import PARALLEL, computation, interval
 
 from ndsl import (
     CompilationConfig,
@@ -15,6 +14,7 @@
     StencilFactory,
 )
 from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl.dsl.gt4py import PARALLEL, computation, interval
 from ndsl.dsl.typing import FloatField
 
 

tests/dsl/test_stencil.py

@@ -1,7 +1,7 @@
-from gt4py.cartesian.gtscript import PARALLEL, Field, computation, interval
 from gt4py.storage import empty, ones
 
 from ndsl import CompilationConfig, GridIndexing, StencilConfig, StencilFactory
+from ndsl.dsl.gt4py import PARALLEL, Field, computation, interval
 
 
 def _make_storage(

tests/dsl/test_stencil_factory.py

@@ -1,6 +1,5 @@
 import numpy as np
 import pytest
-from gt4py.cartesian.gtscript import PARALLEL, computation, horizontal, interval, region
 
 from ndsl import (
     CompilationConfig,
@@ -11,6 +10,7 @@
     StencilFactory,
 )
 from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl.dsl.gt4py import PARALLEL, computation, horizontal, interval, region
 from ndsl.dsl.gt4py_utils import make_storage_from_shape
 from ndsl.dsl.stencil import CompareToNumpyStencil, get_stencils_with_varied_bounds
 from ndsl.dsl.typing import FloatField

tests/dsl/test_stencil_wrapper.py

@@ -4,7 +4,6 @@
 import gt4py.cartesian.gtscript
 import numpy as np
 import pytest
-from gt4py.cartesian.gtscript import PARALLEL, computation, interval
 
 from ndsl import (
     CompilationConfig,
@@ -14,6 +13,7 @@
     Quantity,
     StencilConfig,
 )
+from ndsl.dsl.gt4py import PARALLEL, computation, interval
 from ndsl.dsl.gt4py_utils import make_storage_from_shape
 from ndsl.dsl.stencil import _convert_quantities_to_storage
 from ndsl.dsl.typing import Float, FloatField

tests/test_boilerplate.py

@@ -1,8 +1,8 @@
 import numpy as np
-from gt4py.cartesian.gtscript import PARALLEL, computation, interval
 
 from ndsl import QuantityFactory, StencilFactory
 from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl.dsl.gt4py import PARALLEL, computation, interval
 from ndsl.dsl.typing import FloatField