Add domain-size checks on FrozenStencil

ndsl/dsl/stencil.py

@@ -391,6 +391,8 @@ def __call__(self, *args, **kwargs) -> None:
         if self.stencil_config.verbose:
             ndsl_log.debug(f"Running {self._func_name}")
 
+        self._validate_quantity_sizes(args, kwargs)
+
         # Marshal arguments
         args_list = list(args)
         _convert_quantities_to_storage(args_list, kwargs)
@@ -538,6 +540,26 @@ def closure_resolver(self, constant_args, given_args, parent_closure=None):
             constant_args, given_args, parent_closure=parent_closure
         )
 
+    def _validate_quantity_sizes(self, *args, **kwargs):
+        """Checks that the sizes of quantities are compatible with the domain of the stencil.
+
+        Raises:
+            ValueError: If the quantity's size is a mismatch to the domain.
+        """
+        all_args_as_kwargs = dict(zip(self._argument_names, args)) | kwargs
+        domain_sizes = dict(zip(("x", "y", "z"), self.domain))
+
+        for name, argument in all_args_as_kwargs.items():
+            if not isinstance(argument, Quantity):
+                continue
+
+            for axis, quantity_size in zip(argument.dims, argument.extent):
+                if quantity_size < domain_sizes[axis[0]]:
+                    raise ValueError(
+                        f"Quantity `{name}` is too small for the targeted "
+                        f"domain in axis {axis[0]}: {quantity_size} < {domain_sizes[axis[0]]}."
+                    )
+
 
 def _convert_quantities_to_storage(args, kwargs):
     for i, arg in enumerate(args):

tests/dsl/test_stencil.py

@@ -1,8 +1,21 @@
+from contextlib import nullcontext as does_not_raise
+from typing import Callable
+from unittest.mock import MagicMock
+
+import numpy as np
 import pytest
 from gt4py.storage import empty, ones
 
-from ndsl import CompilationConfig, GridIndexing, StencilConfig, StencilFactory
+from ndsl import (
+    CompilationConfig,
+    FrozenStencil,
+    GridIndexing,
+    StencilConfig,
+    StencilFactory,
+)
 from ndsl.dsl.gt4py import PARALLEL, Field, computation, interval
+from ndsl.dsl.typing import FloatField
+from ndsl.quantity import Quantity
 from tests.dsl import utils
 
 
@@ -61,3 +74,39 @@ def test_grid_indexing_get_2d_compute_origin_domain(
 
     assert origin[2] == expected_origin_k
     assert domain[2] == 1
+
+
+def copy_stencil(q_in: FloatField, q_out: FloatField):  # type: ignore
+    with computation(PARALLEL), interval(...):
+        q_out[0, 0, 0] = q_in
+
+
+@pytest.mark.parametrize(
+    "extent,dimensions,domain,expectation",
+    [
+        ((20, 20, 30), ["x", "y", "z"], (20, 20, 20), does_not_raise()),
+        ((20, 20), ["x", "y"], (20, 20, 30), does_not_raise()),
+        ((20, 20), ["x_interface", "y"], (20, 20, 30), does_not_raise()),
+        ((20, 20), ["x", "y_interface"], (20, 20, 30), does_not_raise()),
+        ((20,), ["z"], (20, 20, 10), does_not_raise()),
+        ((20,), ["z_interface"], (20, 20, 10), does_not_raise()),
+        ((15, 20, 30), ["x", "y", "z"], (20, 20, 30), pytest.raises(ValueError)),
+        ((20, 15, 30), ["x", "y", "z"], (20, 20, 30), pytest.raises(ValueError)),
+        ((20, 20, 15), ["x", "y", "z"], (20, 20, 30), pytest.raises(ValueError)),
+    ],
+)
+def test_domain_size_comparison(
+    extent: tuple[int],
+    dimensions: list[str],
+    domain: tuple[int],
+    expectation: Callable,
+):
+    quantity = Quantity(np.zeros(extent), dimensions, "n/a", extent=extent)
+    stencil = FrozenStencil(
+        copy_stencil,
+        origin=(0, 0, 0),
+        domain=domain,
+        stencil_config=MagicMock(spec=StencilConfig),
+    )
+    with expectation:
+        stencil._validate_quantity_sizes(quantity)