2D Stencil Indexing Method

ndsl/dsl/stencil.py

@@ -797,6 +797,25 @@ def get_origin_domain(
             domain[i] += 2 * n
         return tuple(origin), tuple(domain)
 
+    def get_2d_compute_origin_domain(
+        self,
+        klevel: int = 0,
+    ) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
+        """
+        Get the origin and domain for a computation that occurs on the lowest klevel over a certain grid
+        configuration (given by dims) and a certain number of halo points.
+
+        Args:
+            klevel: the vertical level of the domain, defaults to zero
+
+        Returns:
+            origin: origin of the computation
+            domain: shape of the computation
+        """
+        origin = (self.isc, self.jsc, klevel)
+        domain = (self.iec + 1 - self.isc, self.jec + 1 - self.jsc, 1)
+        return (origin, domain)
+
     def _origin_from_dims(self, dims: Iterable[str]) -> List[int]:
         return_origin = []
         for dim in dims:

tests/dsl/test_stencil_factory.py

@@ -202,6 +202,48 @@ def test_stencil_factory_numpy_comparison_from_origin_domain(
         assert isinstance(stencil, FrozenStencil)
 
 
+@pytest.mark.parametrize("enabled", [True, False])
+@pytest.mark.parametrize("klevel", [0, 1, 30])
+def test_stencil_factory_numpy_comparison_from_origin_domain_2d(
+    enabled: bool, klevel: int
+):
+    backend = "numpy"
+    dace_config = DaceConfig(communicator=None, backend=backend)
+    config = StencilConfig(
+        compilation_config=CompilationConfig(
+            backend=backend,
+            rebuild=False,
+            validate_args=False,
+            format_source=False,
+            device_sync=False,
+        ),
+        compare_to_numpy=enabled,
+        dace_config=dace_config,
+    )
+    indexing = GridIndexing(
+        domain=(12, 12, 79),
+        n_halo=3,
+        south_edge=True,
+        north_edge=True,
+        west_edge=True,
+        east_edge=True,
+    )
+    if klevel > 0:
+        origin, domain = indexing.get_2d_compute_origin_domain(klevel=klevel)
+    else:
+        origin, domain = indexing.get_2d_compute_origin_domain()
+    assert domain[2] == 1
+    assert origin[2] == klevel
+    factory = StencilFactory(config=config, grid_indexing=indexing)
+    stencil = factory.from_origin_domain(
+        func=copy_stencil, origin=origin, domain=domain
+    )
+    if enabled:
+        assert isinstance(stencil, CompareToNumpyStencil)
+    else:
+        assert isinstance(stencil, FrozenStencil)
+
+
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_stencil_factory_numpy_comparison_runs_without_exceptions(backend: str) -> None:
     dace_config = DaceConfig(communicator=None, backend=backend)