Merge pull request #78 from FlorianDeconinck/feature/K_axis_bounds_in…

…_externals

[feature] Expose stencil vertical domain to Frozen Stencil

ndsl/dsl/stencil.py

@@ -737,6 +737,9 @@ def axis_offsets(
  "local_js": gtscript.J[0] + self.jsc - origin[1],
  "j_end": j_end,
  "local_je": gtscript.J[-1] + self.jec - origin[1] - domain[1] + 1,
+ "k_start": origin[2] if len(origin) > 2 else 0,
+ "k_end": (origin[2] if len(origin) > 2 else 0)
+ + (domain[2] - 1 if len(domain) > 2 else 0),
  }
 
  def get_origin_domain(

tests/dsl/test_stencil_factory.py

@@ -107,6 +107,23 @@ def test_get_stencils_with_varied_bounds_and_regions(backend: str):
  np.testing.assert_array_equal(q_orig.data, q_ref.data)
 
 
+def test_stencil_vertical_bounds(backend: str):
+ factory = get_stencil_factory(backend)
+ origins = [(3, 3, 0), (2, 2, 1)]
+ domains = [(1, 1, 3), (2, 2, 4)]
+ stencils = get_stencils_with_varied_bounds(
+ add_1_in_region_stencil,
+ origins,
+ domains,
+ stencil_factory=factory,
+ )
+
+ assert "k_start" in stencils[0].externals and stencils[0].externals["k_start"] == 0
+ assert "k_end" in stencils[0].externals and stencils[0].externals["k_end"] == 2
+ assert "k_start" in stencils[1].externals and stencils[1].externals["k_start"] == 1
+ assert "k_end" in stencils[1].externals and stencils[1].externals["k_end"] == 4
+
+
 @pytest.mark.parametrize("enabled", [True, False])
 def test_stencil_factory_numpy_comparison_from_dims_halo(enabled: bool):
  backend = "numpy"