State V0

ndsl/__init__.py

@@ -28,7 +28,7 @@
 from .performance.collector import NullPerformanceCollector, PerformanceCollector
 from .performance.profiler import NullProfiler, Profiler
 from .performance.report import Experiment, Report, TimeReport
-from .quantity import Quantity
+from .quantity import Quantity, State
 from .quantity.field_bundle import FieldBundle, FieldBundleType  # Break circular import
 from .testing.dummy_comm import DummyComm
 from .types import Allocator
@@ -48,7 +48,6 @@
     "FV3CodePath",
     "DaceConfig",
     "DaCeOrchestration",
-    "FrozenCompiledSDFG",
     "orchestrate",
     "orchestrate_function",
     "ArrayReport",
@@ -87,4 +86,5 @@
     "DummyComm",
     "Allocator",
     "MetaEnumStr",
+    "State",
 ]

ndsl/quantity/__init__.py

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

ndsl/quantity/quantity.py

@@ -263,6 +263,9 @@ def data(self) -> Union[np.ndarray, cupy.ndarray]:
     def data(self, inputData):
         if type(inputData) in [np.ndarray, cupy.ndarray]:
             self._data = inputData
+            self._compute_domain_view = BoundedArrayView(
+                self.data, self.dims, self.origin, self.extent
+            )
 
     @property
     def origin(self) -> Tuple[int, ...]:

ndsl/quantity/state.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING, Any, Self
+
+import dacite
+import xarray as xr
+from mpi4py import MPI
+from numpy.typing import ArrayLike
+
+
+if TYPE_CHECKING:
+    from ndsl import QuantityFactory
+
+
+@dataclasses.dataclass
+class State:
+    """Base class for State object in models that bundles a collection
+    of functions to deal with nested dataclasses and common usage of States:
+        - init (zero, from memory, zero copy buffer swap)
+        - IO (save to NetCDF, from NetCDF)
+
+    The State expects Quantities.
+    """
+
+    @classmethod
+    def zeros(cls, quantity_factory: QuantityFactory) -> Self:
+        """Init all quantities to zeros - included nested ones"""
+
+        def _zeros_recursive(cls):
+            initial_quantities = {}
+            for _field in dataclasses.fields(cls):
+                if dataclasses.is_dataclass(_field.type):
+                    initial_quantities[_field.name] = _zeros_recursive(_field.type)
+                else:
+                    if "dims" not in _field.metadata.keys():
+                        raise ValueError(
+                            "Malformed state - no dims to init "
+                            f"Quantity in  {_field.name} of type {_field.type}"
+                        )
+
+                    initial_quantities[_field.name] = quantity_factory.zeros(
+                        _field.metadata["dims"],
+                        _field.metadata["units"],
+                        dtype=_field.metadata["dtype"],
+                        allow_mismatch_float_precision=True,
+                    )
+
+            return initial_quantities
+
+        dict_of_qty = _zeros_recursive(cls)
+        return dacite.from_dict(data_class=cls, data=dict_of_qty)
+
+    def init_from_memory(self, memory_map: dict[str, Any]):
+        """Will copy data from the memory map if it follows the nested
+        naming convention of the dataclass"""
+
+        def _init_from_memory_recursive(dataclss, memory_map: dict[str, Any]):
+            for name, array in memory_map.items():
+                if isinstance(array, dict):
+                    _init_from_memory_recursive(dataclss.__getattribute__(name), array)
+                else:
+                    try:
+                        dataclss.__getattribute__(name).field[:] = array
+                    except ValueError as e:
+                        e.add_note(
+                            f"Error when initializing field {name} on state {type(self)}"
+                        )
+                        raise e
+
+        _init_from_memory_recursive(self, memory_map)
+
+    def init_zero_copy(self, memory_map: dict[str, Any], check: bool = True):
+        """Swap buffers given into the Quantities carried by the state
+        by following dataclass naming convention"""
+
+        def _init_zero_copy_recursive(dataclss, memory_map: dict[str, Any | ArrayLike]):
+            for name, array in memory_map.items():
+                if isinstance(array, dict):
+                    _init_zero_copy_recursive(dataclss.__getattribute__(name), array)
+                else:
+                    qty = dataclss.__getattribute__(name)
+                    if check:
+                        if array.shape != qty.field.shape:
+                            e = ValueError("Shape mismatch on zero copy for")
+                            e.add_note(f"  Error on {name} for {type(dataclss)}")
+                            e.add_note(f"  Shapes: {array.shape} != {qty.field.shape}")
+                            raise e
+                        if array.strides != qty.data.strides:
+                            e = ValueError("Stride mismatch on zero copy for")
+                            e.add_note(f"  Error on {name} for {type(dataclss)}")
+                            e.add_note(
+                                f"  Strides: {array.strides} != {qty.data.strides}"
+                            )
+                            raise e
+
+                    qty.data = array
+
+        _init_zero_copy_recursive(self, memory_map)
+
+    def to_netcdf(self, path: str = "./"):
+        def _save_recursive(datclss: State):
+            local_data = {}
+            for _field in dataclasses.fields(datclss):
+                if dataclasses.is_dataclass(_field.type):
+                    local_data[_field.name] = xr.Dataset(
+                        data_vars=_save_recursive(datclss.__getattribute__(_field.name))
+                    )
+                else:
+                    if "dims" not in _field.metadata.keys():
+                        raise ValueError(
+                            "Malformed state - no dims to init "
+                            f"Quantity in  {_field.name} of type {_field.type}"
+                        )
+
+                    local_data[_field.name] = datclss.__getattribute__(
+                        _field.name
+                    ).field_as_xarray
+
+            return local_data
+
+        datatree = _save_recursive(self)
+
+        # Move top-level into their own dataset in the "/" prefix
+        # to match DataTree expected format
+        top_level = {}
+        for key, value in datatree.items():
+            if not isinstance(value, xr.Dataset):
+                top_level[key] = value
+        for key, value in top_level.items():
+            datatree.pop(key)
+        datatree["/"] = xr.Dataset(data_vars=top_level)
+
+        # Resolve rank-tied postfix if needed
+        rank_postfix = ""
+        if MPI.COMM_WORLD.Get_size() > 1:
+            rank_postfix = f"_rank{MPI.COMM_WORLD.Get_rank()}"
+
+        xr.DataTree.from_dict(datatree).to_netcdf(
+            f"{path}{type(self).__name__}{rank_postfix}.nc4"
+        )
+
+    def from_netcdf(self, path: str):
+        datatree = xr.open_datatree(path)
+        datatree_as_dict = datatree.to_dict()
+
+        # All other cases - recursing downward
+        def _load_recursive(data_tree_as_dict: dict[str, xr.Dataset] | xr.Dataset):
+            local_data_dict = {}
+            for name, data_array in data_tree_as_dict.items():
+                # Case of the top_level "/"
+                if name == "/":
+                    for root_name, root_data_array in datatree_as_dict["/"].items():
+                        local_data_dict[root_name] = root_data_array.to_numpy()
+                else:
+                    # Get the leading `/` out
+                    if isinstance(data_array, xr.Dataset):
+                        local_data_dict[name[1:]] = _load_recursive(data_array)
+                    else:
+                        local_data_dict[name] = data_array.to_numpy()
+
+            return local_data_dict
+
+        data_as_numpy_dict = _load_recursive(datatree_as_dict)
+
+        self.init_from_memory(data_as_numpy_dict)

setup.py

@@ -24,6 +24,7 @@ def local_pkg(name: str, relative_path: str) -> str:
     "matplotlib",  # for plotting in boilerplate
     "cartopy",  # for plotting in ndsl.viz
     "pytest-subtests",  # for translate tests
+    "dacite",  # for state
 ]
 
 

tests/quantity/test_state.py

@@ -0,0 +1,67 @@
+import dataclasses
+
+import numpy as np
+
+from ndsl import Quantity, State
+from ndsl.boilerplate import get_factories_single_tile
+from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Float
+
+
+@dataclasses.dataclass
+class CodeState(State):
+    @dataclasses.dataclass
+    class InnerA:
+        A: Quantity = dataclasses.field(
+            metadata={
+                "name": "A",
+                "dims": [X_DIM, Y_DIM, Z_DIM],
+                "units": "kg kg-1",
+                "intent": "?",
+                "dtype": Float,
+            }
+        )
+
+    @dataclasses.dataclass
+    class InnerB:
+        B: Quantity = dataclasses.field(
+            metadata={
+                "name": "B",
+                "dims": [X_DIM, Y_DIM, Z_DIM],
+                "units": "1",
+                "intent": "?",
+                "dtype": Float,
+            }
+        )
+
+    inner_A: InnerA
+    inner_B: InnerB
+    C: Quantity = dataclasses.field(
+        metadata={
+            "name": "C",
+            "dims": [X_DIM, Y_DIM, Z_DIM],
+            "units": "kg kg-1",
+            "intent": "?",
+            "dtype": Float,
+        }
+    )
+
+
+def test_state():
+    _, qty_factry = get_factories_single_tile(5, 5, 3, 0, backend="dace:cpu_kfirst")
+
+    microphys_state = CodeState.zeros(qty_factry)
+    microphys_state.inner_A.A.field[:] = 42.42
+    microphys_state.to_netcdf()
+    microphys_state2 = CodeState.zeros(qty_factry)
+    microphys_state2.from_netcdf("CodeState.nc4")
+    assert (microphys_state2.inner_A.A.field[:] == 42.42).all()
+    a = np.ones((5, 5, 3))
+    b = np.ones((5, 5, 3))
+    c = np.ones((5, 5, 3))
+    b[:] = 23.23
+    microphys_state2.init_zero_copy(
+        {"inner_A": {"A": a}, "inner_B": {"B": b}, "C": c},
+        check=False,
+    )
+    assert (microphys_state2.inner_A.A.field[:] == 1.0).all()
+    assert (microphys_state2.inner_B.B.field[:] == 23.23).all()