Field Bundle: Functional code for 4D array

ndsl/__init__.py

@@ -29,6 +29,7 @@
 from .performance.profiler import NullProfiler, Profiler
 from .performance.report import Experiment, Report, TimeReport
 from .quantity import Quantity
+from .quantity.field_bundle import FieldBundle, FieldBundleType  # Break circular import
 from .testing.dummy_comm import DummyComm
 from .types import Allocator
 from .utils import MetaEnumStr

ndsl/dsl/stencil.py

@@ -33,6 +33,7 @@
 from ndsl.initialization.sizer import GridSizer, SubtileGridSizer
 from ndsl.logging import ndsl_log
 from ndsl.quantity import Quantity
+from ndsl.quantity.field_bundle import FieldBundleType, MarkupFieldBundleType
 from ndsl.testing.comparison import LegacyMetric
 
 
@@ -335,6 +336,14 @@ def __init__(
             ):
                 block_waiting_for_compilation(MPI.COMM_WORLD, compilation_config)
 
+            # Field Bundle might have dropped a placeholder type that we know
+            # have to resolve to the proper type
+            for name, types in func.__annotations__.items():
+                if isinstance(types, MarkupFieldBundleType):
+                    func.__annotations__[name] = FieldBundleType.T(
+                        types.name, do_markup=False
+                    )
+
             self.stencil_object = gtscript.stencil(
                 definition=func,
                 externals=externals,

ndsl/quantity/__init__.py

@@ -1,9 +1,11 @@
-from ndsl.quantity.metadata import QuantityHaloSpec, QuantityMetadata
-from ndsl.quantity.quantity import Quantity
+from .metadata import QuantityHaloSpec, QuantityMetadata
+from .quantity import Quantity
 
 
 __all__ = [
     "Quantity",
     "QuantityMetadata",
     "QuantityHaloSpec",
+    "FieldBundle",
+    "FieldBundleType",
 ]

ndsl/quantity/field_bundle.py

@@ -0,0 +1,143 @@
+import copy
+from dataclasses import dataclass
+from typing import Any, Dict, Tuple
+
+from ndsl.dsl.typing import Float
+from ndsl.initialization.allocator import QuantityFactory
+from ndsl.quantity.quantity import Quantity
+
+
+import gt4py.cartesian.gtscript as gtscript  # isort: skip
+
+
+# ToDo: This is 4th dimensions restricted. We need a concept
+#       of data dimensions index here to be able to extend to N dimensions
+_DataDimensionIndex = int
+_FieldBundleIndexer = Dict[str, _DataDimensionIndex]
+
+
+class FieldBundle:
+    """Field Bundle wraps a nD array (3D + n Data Dimensions) into a complex
+    indexing scheme that allows a dual name-based and index-based
+    access to the underlying memory. It is paired with the `FieldBundleType`
+    which provides a way to
+
+    WARNING: The present implementation only allows for 4D array.
+    """
+
+    _quantity: Quantity
+    _indexer: _FieldBundleIndexer = {}
+
+    def __init__(
+        self,
+        bundle_name: str,
+        quantity: Quantity,
+        mapping: _FieldBundleIndexer = {},
+        do_register_type: bool = False,
+    ):
+        if len(quantity.shape) != 4:
+            raise NotImplementedError("FieldBundle implementation restricted to 4D")
+
+        self.name = bundle_name
+        self._quantity = quantity
+        self._indexer = mapping
+        if do_register_type:
+            FieldBundleType.register(bundle_name, quantity.shape[3:])
+
+    def map(self, index: _DataDimensionIndex, name: str):
+        self._indexer[name] = index
+
+    @property
+    def quantity(self) -> Quantity:
+        return self._quantity
+
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        return self._quantity.shape
+
+    def groupby(self, name: str) -> Quantity:
+        raise NotImplementedError
+
+    def __getattr__(self, name: str) -> Quantity:
+        if name not in self._indexer.keys():
+            # This replicates as close possible the default behavior of getattr
+            # without breaking orchestration
+            return None  # type: ignore
+        # ToDo: extend the dims below to work with more than 4 dims
+        assert len(self._quantity.data.shape) == 4
+        return Quantity(
+            data=self._quantity.data[:, :, :, self._indexer[name]],
+            dims=self._quantity.dims[:-1],
+            units=self._quantity.units,
+            origin=self._quantity.origin[:-1],
+            extent=self._quantity.extent[:-1],
+        )
+
+    @property
+    def __array_interface__(self):
+        return self._quantity.__array_interface__
+
+    @property
+    def __cuda_array_interface__(self):
+        return self._quantity.__array_interface__
+
+    def __descriptor__(self) -> Any:
+        return self._quantity.__descriptor__()
+
+    @staticmethod
+    def extend_3D_quantity_factory(
+        quantity_factory: QuantityFactory,
+        extra_dims: dict[str, int],
+    ) -> QuantityFactory:
+        new_factory = copy.copy(quantity_factory)
+        new_factory.set_extra_dim_lengths(
+            **{
+                **extra_dims,
+            }
+        )
+        return new_factory
+
+    def index(self, name: str) -> int:
+        return self._indexer[name]
+
+
+@dataclass
+class MarkupFieldBundleType:
+    """Markup a field bundle to delay specialization
+
+    Properties:
+        name: name of the future type to look into the registrat
+    """
+
+    name: str
+
+
+class FieldBundleType:
+    """Field Bundle Types to help with static sizing of Data Dimensions
+
+    Methods:
+        register: Register a type by sizing it's data dimensions
+        T: access any registered types for type hinting
+    """
+
+    _field_type_registrar: dict[str, gtscript._FieldDescriptor] = {}
+
+    @classmethod
+    def register(cls, name: str, data_dims: tuple[int]) -> gtscript._FieldDescriptor:
+        if name in cls._field_type_registrar.keys():
+            raise RuntimeError(f"Registering {name} a second time!")
+        cls._field_type_registrar[name] = gtscript.Field[
+            gtscript.IJK, (Float, (data_dims))
+        ]
+        return cls._field_type_registrar[name]
+
+    @classmethod
+    def T(
+        cls, name: str, do_markup: bool = True
+    ) -> gtscript._FieldDescriptor | MarkupFieldBundleType:
+        if name not in cls._field_type_registrar:
+            if do_markup:
+                return MarkupFieldBundleType(name)
+            else:
+                raise RuntimeError(f"FieldBundle type {name} as not been registered!")
+        return cls._field_type_registrar[name]

tests/quantity/test_fieldbundle.py

@@ -0,0 +1,77 @@
+from ndsl.boilerplate import get_factories_single_tile
+from ndsl.constants import X_DIM, Y_DIM, Z_DIM
+from ndsl.dsl.gt4py import PARALLEL, computation, interval
+from ndsl.dsl.typing import FloatField
+from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType
+
+
+def assign_4d_field_stcl(field_4d: FieldBundleType.T("Tracers")):  # type: ignore # noqa
+    with computation(PARALLEL), interval(...):
+        field_4d[0, 0, 0][1] = 63.63
+        field_4d[0, 0, 0][3] = 63.63
+
+
+def assign_3d_field_stcl(field_3d: FloatField):
+    with computation(PARALLEL), interval(...):
+        field_3d = 121.121
+
+
+def test_field_bundle():
+    # Grid & Factories
+    NX = 2
+    NY = 2
+    NZ = 2
+    N4th = 5
+    stencil_factory, quantity_factory = get_factories_single_tile(NX, NY, NZ, 1)
+
+    # Type register
+    FieldBundleType.register("Tracers", (N4th,))
+
+    # Make stencils
+    assign_4d_field = stencil_factory.from_dims_halo(
+        func=assign_4d_field_stcl,
+        compute_dims=[X_DIM, Y_DIM, Z_DIM],
+    )
+    assign_3d_field = stencil_factory.from_dims_halo(
+        func=assign_3d_field_stcl,
+        compute_dims=[X_DIM, Y_DIM, Z_DIM],
+    )
+
+    # "Input" data
+    new_quantity_factory = FieldBundle.extend_3D_quantity_factory(
+        quantity_factory, {"tracers": N4th}
+    )
+    data = new_quantity_factory.ones([X_DIM, Y_DIM, Z_DIM, "tracers"], units="kg/g")
+
+    # Build Bundle
+    tracers = FieldBundle(
+        bundle_name="tracers",
+        quantity=data,
+        mapping={"vapor": 0, "cloud": 2},
+    )
+
+    # Test
+    tracers.quantity.field[:, :, :, :] = 48.4
+    tracers.quantity.field[:, :, :, 2] = 21.21
+
+    assign_4d_field(tracers.quantity)
+
+    assert (tracers.quantity.field[:, :, :, 0] == 48.4).all()
+    assert (tracers.quantity.field[:, :, :, 1] == 63.63).all()
+    assert (tracers.quantity.field[:, :, :, 2] == 21.21).all()
+    assert (tracers.quantity.field[:, :, :, 3] == 63.63).all()
+    assert (tracers.quantity.field[:, :, :, 4] == 48.4).all()
+
+    tracers.vapor.field[:] = 1000.1000
+
+    assert (tracers.quantity.field[:, :, :, 0] == 1000.1000).all()
+    assert (tracers.quantity.field[:, :, :, 1] == 63.63).all()
+
+    assign_3d_field(tracers.cloud)
+
+    assert (tracers.cloud.field[:] == 121.121).all()
+    assert (tracers.quantity.field[:, :, :, 2] == tracers.cloud.field[:]).all()
+
+
+if __name__ == "__main__":
+    test_field_bundle()