Support for 4d Data

ndsl/dsl/gt4py_utils.py

@@ -53,11 +53,14 @@ def wrapper(*args, **kwargs) -> Any:
 def _mask_to_dimensions(
     mask: Tuple[bool, ...], shape: Sequence[int]
 ) -> List[Union[str, int]]:
-    assert len(mask) == 3
+    assert len(mask) >= 3
     dimensions: List[Union[str, int]] = []
     for i, axis in enumerate(("I", "J", "K")):
         if mask[i]:
             dimensions.append(axis)
+    if len(mask) > 3:
+        for i in range(3, len(mask)):
+            dimensions.append(str(shape[i]))
     offset = int(sum(mask))
     dimensions.extend(shape[offset:])
     return dimensions
@@ -154,6 +157,8 @@ def make_storage_data(
         data = _make_storage_data_2d(
             data, shape, start, dummy, axis, read_only, backend=backend
         )
+    elif n_dims >= 4:
+        data = _make_storage_data_Nd(data, shape, start, backend=backend)
     else:
         data = _make_storage_data_3d(data, shape, start, backend=backend)
 
@@ -257,6 +262,21 @@ def _make_storage_data_3d(
     return buffer
 
 
+def _make_storage_data_Nd(
+    data: Field,
+    shape: Tuple[int, ...],
+    start: Tuple[int, ...] = None,
+    *,
+    backend: str,
+) -> Field:
+    if start is None:
+        start = tuple([0] * data.ndim)
+    buffer = zeros(shape, backend=backend)
+    idx = tuple([slice(start[i], start[i] + data.shape[i]) for i in range(len(start))])
+    buffer[idx] = asarray(data, type(buffer))
+    return buffer
+
+
 def make_storage_from_shape(
     shape: Tuple[int, ...],
     origin: Tuple[int, ...] = origin,
@@ -310,6 +330,7 @@ def make_storage_dict(
     axis: int = 2,
     *,
     backend: str,
+    dtype: DTypes = Float,
 ) -> Dict[str, "Field"]:
     assert names is not None, "for 4d variable storages, specify a list of names"
     if shape is None:
@@ -324,6 +345,7 @@ def make_storage_dict(
             dummy=dummy,
             axis=axis,
             backend=backend,
+            dtype=dtype,
         )
     return data_dict
 

ndsl/stencils/testing/translate.py

@@ -5,7 +5,7 @@
 
 import ndsl.dsl.gt4py_utils as utils
 from ndsl.dsl.stencil import StencilFactory
-from ndsl.dsl.typing import Field, Float  # noqa: F401
+from ndsl.dsl.typing import Field, Float, Int  # noqa: F401
 from ndsl.quantity import Quantity
 from ndsl.stencils.testing.grid import Grid  # type: ignore
 
@@ -113,6 +113,12 @@ def make_storage_data(
         elif not full_shape and len(array.shape) < 3 and axis == len(array.shape) - 1:
             use_shape[1] = 1
         start = (int(istart), int(jstart), int(kstart))
+        if "float" in str(array.dtype):
+            dtype = Float
+        elif "int" in str(array.dtype):
+            dtype = Int
+        else:
+            dtype = array.dtype
         if names_4d:
             return utils.make_storage_dict(
                 array,
@@ -123,8 +129,12 @@ def make_storage_data(
                 axis=axis,
                 names=names_4d,
                 backend=self.stencil_factory.backend,
+                dtype=dtype,
             )
         else:
+            if len(array.shape) == 4:
+                start = (int(istart), int(jstart), int(kstart), 0)  # type: ignore
+                use_shape.append(array.shape[-1])
             return utils.make_storage_data(
                 array,
                 tuple(use_shape),
@@ -134,6 +144,7 @@ def make_storage_data(
                 axis=axis,
                 read_only=read_only,
                 backend=self.stencil_factory.backend,
+                dtype=dtype,
             )
 
     def storage_vars(self):
@@ -159,7 +170,7 @@ def collect_start_indices(self, datashape, varinfo):
         kstart = self.get_index_from_info(varinfo, "kstart", 0)
         return istart, jstart, kstart
 
-    def make_storage_data_input_vars(self, inputs, storage_vars=None):
+    def make_storage_data_input_vars(self, inputs, storage_vars=None, dict_4d=True):
         inputs_in = {**inputs}
         inputs_out = {}
         if storage_vars is None:
@@ -185,7 +196,7 @@ def make_storage_data_input_vars(self, inputs, storage_vars=None):
             )
 
             names_4d = None
-            if len(inputs_in[serialname].shape) == 4:
+            if (len(inputs_in[serialname].shape) == 4) and dict_4d:
                 names_4d = info.get("names_4d", utils.tracer_variables)
 
             dummy_axes = info.get("dummy_axes", None)