Repository: UXARRAY/uxarray

Issue:
Support Dask Arrays for Grid 
Currently, the `Grid` class is entirely written around NumPy arrays. We should support Dask arrays. 

Discussion:

Patch:
diff --git a/docs/user_api/index.rst b/docs/user_api/index.rst
index ac521c04..ec3278b6 100644
--- a/docs/user_api/index.rst
+++ b/docs/user_api/index.rst
@@ -224,6 +224,7 @@ Methods
    Grid.get_kd_tree
    Grid.copy
    Grid.isel
+   Grid.chunk
 
 
 Dimensions
diff --git a/uxarray/core/api.py b/uxarray/core/api.py
index 28312266..204b10e3 100644
--- a/uxarray/core/api.py
+++ b/uxarray/core/api.py
@@ -165,6 +165,13 @@ def open_dataset(
         grid_filename_or_obj, latlon=latlon, use_dual=use_dual, **grid_kwargs
     )
 
+    if "chunks" in kwargs:
+        # correctly chunk standardized ugrid dimension names
+        source_dims_dict = uxgrid._source_dims_dict
+        for original_grid_dim, ugrid_grid_dim in source_dims_dict.items():
+            if ugrid_grid_dim in kwargs["chunks"]:
+                kwargs["chunks"][original_grid_dim] = kwargs["chunks"][ugrid_grid_dim]
+
     # UxDataset
     ds = xr.open_dataset(filename_or_obj, **kwargs)  # type: ignore
 
@@ -254,6 +261,13 @@ def open_mfdataset(
         grid_filename_or_obj, latlon=latlon, use_dual=use_dual, **grid_kwargs
     )
 
+    if "chunks" in kwargs:
+        # correctly chunk standardized ugrid dimension names
+        source_dims_dict = uxgrid._source_dims_dict
+        for original_grid_dim, ugrid_grid_dim in source_dims_dict.items():
+            if ugrid_grid_dim in kwargs["chunks"]:
+                kwargs["chunks"][original_grid_dim] = kwargs["chunks"][ugrid_grid_dim]
+
     # UxDataset
     ds = xr.open_mfdataset(paths, **kwargs)  # type: ignore
 
diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py
index 8a0bce17..97b9f5d1 100644
--- a/uxarray/grid/grid.py
+++ b/uxarray/grid/grid.py
@@ -496,6 +496,13 @@ class Grid:
 
         return set([conn for conn in CONNECTIVITY_NAMES if conn in self._ds])
 
+    @property
+    def descriptors(self) -> set:
+        """Names of all descriptor variables."""
+        from uxarray.conventions.descriptors import DESCRIPTOR_NAMES
+
+        return set([desc for desc in DESCRIPTOR_NAMES if desc in self._ds])
+
     @property
     def parsed_attrs(self) -> dict:
         """Dictionary of parsed attributes from the source grid."""
@@ -572,6 +579,12 @@ class Grid:
 
         return self._ds["n_nodes_per_face"]
 
+    @n_nodes_per_face.setter
+    def n_nodes_per_face(self, value):
+        """Setter for ``n_nodes_per_face``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["n_nodes_per_face"] = value
+
     @property
     def node_lon(self) -> xr.DataArray:
         """Longitude of each node in degrees.
@@ -583,6 +596,12 @@ class Grid:
             _populate_node_latlon(self)
         return self._ds["node_lon"]
 
+    @node_lon.setter
+    def node_lon(self, value):
+        """Setter for ``node_lon``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["node_lon"] = value
+
     @property
     def node_lat(self) -> xr.DataArray:
         """Latitude of each node in degrees.
@@ -594,6 +613,12 @@ class Grid:
             _populate_node_latlon(self)
         return self._ds["node_lat"]
 
+    @node_lat.setter
+    def node_lat(self, value):
+        """Setter for ``node_lat``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["node_lat"] = value
+
     @property
     def node_x(self) -> xr.DataArray:
         """Cartesian x location of each node in meters.
@@ -605,6 +630,12 @@ class Grid:
 
         return self._ds["node_x"]
 
+    @node_x.setter
+    def node_x(self, value):
+        """Setter for ``node_x``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["node_x"] = value
+
     @property
     def node_y(self) -> xr.DataArray:
         """Cartesian y location of each node in meters.
@@ -615,6 +646,12 @@ class Grid:
             _populate_node_xyz(self)
         return self._ds["node_y"]
 
+    @node_y.setter
+    def node_y(self, value):
+        """Setter for ``node_y``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["node_y"] = value
+
     @property
     def node_z(self) -> xr.DataArray:
         """Cartesian z location of each node in meters.
@@ -625,6 +662,12 @@ class Grid:
             _populate_node_xyz(self)
         return self._ds["node_z"]
 
+    @node_z.setter
+    def node_z(self, value):
+        """Setter for ``node_z``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["node_z"] = value
+
     @property
     def edge_lon(self) -> xr.DataArray:
         """Longitude of the center of each edge in degrees.
@@ -637,6 +680,12 @@ class Grid:
         _set_desired_longitude_range(self._ds)
         return self._ds["edge_lon"]
 
+    @edge_lon.setter
+    def edge_lon(self, value):
+        """Setter for ``edge_lon``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_lon"] = value
+
     @property
     def edge_lat(self) -> xr.DataArray:
         """Latitude of the center of each edge in degrees.
@@ -648,6 +697,12 @@ class Grid:
         _set_desired_longitude_range(self._ds)
         return self._ds["edge_lat"]
 
+    @edge_lat.setter
+    def edge_lat(self, value):
+        """Setter for ``edge_lat``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_lat"] = value
+
     @property
     def edge_x(self) -> xr.DataArray:
         """Cartesian x location of the center of each edge in meters.
@@ -659,6 +714,12 @@ class Grid:
 
         return self._ds["edge_x"]
 
+    @edge_x.setter
+    def edge_x(self, value):
+        """Setter for ``edge_x``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_x"] = value
+
     @property
     def edge_y(self) -> xr.DataArray:
         """Cartesian y location of the center of each edge in meters.
@@ -669,6 +730,12 @@ class Grid:
             _populate_edge_centroids(self)
         return self._ds["edge_y"]
 
+    @edge_y.setter
+    def edge_y(self, value):
+        """Setter for ``edge_y``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_y"] = value
+
     @property
     def edge_z(self) -> xr.DataArray:
         """Cartesian z location of the center of each edge in meters.
@@ -679,6 +746,12 @@ class Grid:
             _populate_edge_centroids(self)
         return self._ds["edge_z"]
 
+    @edge_z.setter
+    def edge_z(self, value):
+        """Setter for ``edge_z``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_z"] = value
+
     @property
     def face_lon(self) -> xr.DataArray:
         """Longitude of the center of each face in degrees.
@@ -690,6 +763,12 @@ class Grid:
             _set_desired_longitude_range(self._ds)
         return self._ds["face_lon"]
 
+    @face_lon.setter
+    def face_lon(self, value):
+        """Setter for ``face_lon``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_lon"] = value
+
     @property
     def face_lat(self) -> xr.DataArray:
         """Latitude of the center of each face in degrees.
@@ -702,6 +781,12 @@ class Grid:
 
         return self._ds["face_lat"]
 
+    @face_lat.setter
+    def face_lat(self, value):
+        """Setter for ``face_lat``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_lat"] = value
+
     @property
     def face_x(self) -> xr.DataArray:
         """Cartesian x location of the center of each face in meters.
@@ -713,6 +798,12 @@ class Grid:
 
         return self._ds["face_x"]
 
+    @face_x.setter
+    def face_x(self, value):
+        """Setter for ``face_x``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_x"] = value
+
     @property
     def face_y(self) -> xr.DataArray:
         """Cartesian y location of the center of each face in meters.
@@ -723,6 +814,12 @@ class Grid:
             _populate_face_centroids(self)
         return self._ds["face_y"]
 
+    @face_y.setter
+    def face_y(self, value):
+        """Setter for ``face_x``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_y"] = value
+
     @property
     def face_z(self) -> xr.DataArray:
         """Cartesian z location of the center of each face in meters.
@@ -733,6 +830,12 @@ class Grid:
             _populate_face_centroids(self)
         return self._ds["face_z"]
 
+    @face_z.setter
+    def face_z(self, value):
+        """Setter for ``face_z``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_z"] = value
+
     @property
     def face_node_connectivity(self) -> xr.DataArray:
         """Indices of the nodes that make up each face.
@@ -755,6 +858,12 @@ class Grid:
 
         return self._ds["face_node_connectivity"]
 
+    @face_node_connectivity.setter
+    def face_node_connectivity(self, value):
+        """Setter for ``face_node_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_node_connectivity"] = value
+
     @property
     def edge_node_connectivity(self) -> xr.DataArray:
         """Indices of the two nodes that make up each edge.
@@ -768,6 +877,12 @@ class Grid:
 
         return self._ds["edge_node_connectivity"]
 
+    @edge_node_connectivity.setter
+    def edge_node_connectivity(self, value):
+        """Setter for ``edge_node_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_node_connectivity"] = value
+
     @property
     def node_node_connectivity(self) -> xr.DataArray:
         """Indices of the nodes that surround each node."""
@@ -777,6 +892,12 @@ class Grid:
             )
         return self._ds["node_node_connectivity"]
 
+    @node_node_connectivity.setter
+    def node_node_connectivity(self, value):
+        """Setter for ``node_node_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["node_node_connectivity"] = value
+
     @property
     def face_edge_connectivity(self) -> xr.DataArray:
         """Indices of the edges that surround each face.
@@ -788,6 +909,12 @@ class Grid:
 
         return self._ds["face_edge_connectivity"]
 
+    @face_edge_connectivity.setter
+    def face_edge_connectivity(self, value):
+        """Setter for ``face_edge_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_edge_connectivity"] = value
+
     @property
     def edge_edge_connectivity(self) -> xr.DataArray:
         """Indices of the edges that surround each edge.
@@ -801,6 +928,12 @@ class Grid:
 
         return self._ds["edge_edge_connectivity"]
 
+    @edge_edge_connectivity.setter
+    def edge_edge_connectivity(self, value):
+        """Setter for ``edge_edge_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_edge_connectivity"] = value
+
     @property
     def node_edge_connectivity(self) -> xr.DataArray:
         """Indices of the edges that surround each node."""
@@ -811,6 +944,12 @@ class Grid:
 
         return self._ds["node_edge_connectivity"]
 
+    @node_edge_connectivity.setter
+    def node_edge_connectivity(self, value):
+        """Setter for ``node_edge_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["node_edge_connectivity"] = value
+
     @property
     def face_face_connectivity(self) -> xr.DataArray:
         """Indices of the faces that surround each face.
@@ -822,6 +961,12 @@ class Grid:
 
         return self._ds["face_face_connectivity"]
 
+    @face_face_connectivity.setter
+    def face_face_connectivity(self, value):
+        """Setter for ``face_face_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_face_connectivity"] = value
+
     @property
     def edge_face_connectivity(self) -> xr.DataArray:
         """Indices of the faces that saddle each edge.
@@ -833,6 +978,12 @@ class Grid:
 
         return self._ds["edge_face_connectivity"]
 
+    @edge_face_connectivity.setter
+    def edge_face_connectivity(self, value):
+        """Setter for ``edge_face_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_face_connectivity"] = value
+
     @property
     def node_face_connectivity(self) -> xr.DataArray:
         """Indices of the faces that surround each node.
@@ -844,6 +995,12 @@ class Grid:
 
         return self._ds["node_face_connectivity"]
 
+    @node_face_connectivity.setter
+    def node_face_connectivity(self, value):
+        """Setter for ``node_face_connectivity``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["node_face_connectivity"] = value
+
     @property
     def edge_node_distances(self):
         """Distances between the two nodes that surround each edge.
@@ -854,6 +1011,12 @@ class Grid:
             _populate_edge_node_distances(self)
         return self._ds["edge_node_distances"]
 
+    @edge_node_distances.setter
+    def edge_node_distances(self, value):
+        """Setter for ``edge_node_distances``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_node_distances"] = value
+
     @property
     def edge_face_distances(self):
         """Distances between the centers of the faces that saddle each edge.
@@ -864,6 +1027,12 @@ class Grid:
             _populate_edge_face_distances(self)
         return self._ds["edge_face_distances"]
 
+    @edge_face_distances.setter
+    def edge_face_distances(self, value):
+        """Setter for ``edge_face_distances``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["edge_face_distances"] = value
+
     @property
     def antimeridian_face_indices(self) -> np.ndarray:
         """Index of each face that crosses the antimeridian."""
@@ -883,6 +1052,12 @@ class Grid:
             )
         return self._ds["face_areas"]
 
+    @face_areas.setter
+    def face_areas(self, value):
+        """Setter for ``face_areas``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["face_areas"] = value
+
     @property
     def bounds(self):
         """Latitude Longitude Bounds for each Face in degrees.
@@ -897,6 +1072,12 @@ class Grid:
 
         return self._ds["bounds"]
 
+    @bounds.setter
+    def bounds(self, value):
+        """Setter for ``bounds``"""
+        assert isinstance(value, xr.DataArray)
+        self._ds["bounds"] = value
+
     @property
     def face_jacobian(self):
         """Declare face_jacobian as a property."""
@@ -904,6 +1085,72 @@ class Grid:
             _ = self.face_areas
         return self._face_jacobian
 
+    def chunk(self, n_node="auto", n_edge="auto", n_face="auto"):
+        """Converts all arrays to dask arrays with given chunks across grid
+        dimensions in-place.
+
+        Non-dask arrays will be converted to dask arrays. Dask arrays will be chunked to the given chunk size.
+
+        Parameters
+        ----------
+        n_node : int, tuple
+            How to chunk node variables. Must be one of the following forms:
+
+            - A blocksize like 1000.
+            - A blockshape like (1000, 1000).
+            - Explicit sizes of all blocks along all dimensions like
+              ((1000, 1000, 500), (400, 400)).
+            - A size in bytes, like "100 MiB" which will choose a uniform
+              block-like shape
+            - The word "auto" which acts like the above, but uses a configuration
+              value ``array.chunk-size`` for the chunk size
+
+            -1 or None as a blocksize indicate the size of the corresponding
+            dimension.
+
+        n_edge : int, tuple
+            How to chunk edge variables. Must be one of the following forms:
+
+            - A blocksize like 1000.
+            - A blockshape like (1000, 1000).
+            - Explicit sizes of all blocks along all dimensions like
+              ((1000, 1000, 500), (400, 400)).
+            - A size in bytes, like "100 MiB" which will choose a uniform
+              block-like shape
+            - The word "auto" which acts like the above, but uses a configuration
+              value ``array.chunk-size`` for the chunk size
+
+            -1 or None as a blocksize indicate the size of the corresponding
+            dimension.
+
+        n_face : int, tuple
+            How to chunk face variables. Must be one of the following forms:
+
+            - A blocksize like 1000.
+            - A blockshape like (1000, 1000).
+            - Explicit sizes of all blocks along all dimensions like
+              ((1000, 1000, 500), (400, 400)).
+            - A size in bytes, like "100 MiB" which will choose a uniform
+              block-like shape
+            - The word "auto" which acts like the above, but uses a configuration
+              value ``array.chunk-size`` for the chunk size
+
+            -1 or None as a blocksize indicate the size of the corresponding
+            dimension.
+        """
+
+        grid_var_names = self.coordinates | self.connectivity | self.descriptors
+
+        for var_name in grid_var_names:
+            grid_var = getattr(self, var_name)
+
+            if "n_node" in grid_var.dims:
+                setattr(self, var_name, grid_var.chunk(chunks={"n_node": n_node}))
+            elif "n_edge" in grid_var.dims:
+                setattr(self, var_name, grid_var.chunk(chunks={"n_edge": n_edge}))
+            elif "n_face" in grid_var.dims:
+                setattr(self, var_name, grid_var.chunk(chunks={"n_face": n_face}))
+
     def get_ball_tree(
         self,
         coordinates: Optional[str] = "nodes",


Test patch:
diff --git a/test/test_dask.py b/test/test_dask.py
new file mode 100644
index 00000000..f3b55586
--- /dev/null
+++ b/test/test_dask.py
@@ -0,0 +1,51 @@
+import uxarray as ux
+import numpy as np
+import dask.array as da
+
+import pytest
+import os
+from pathlib import Path
+
+
+
+current_path = Path(os.path.dirname(os.path.realpath(__file__)))
+
+mpas_grid = current_path / 'meshfiles' / "mpas" / "QU" / 'oQU480.231010.nc'
+
+csne30_grid = current_path / 'meshfiles' / "ugrid" / "outCSne30" / 'outCSne30.ug'
+csne30_data = current_path / 'meshfiles' / "ugrid" / "outCSne30" / 'outCSne30_var2.nc'
+
+
+def test_grid_chunking():
+    """Tests the chunking of an entire grid."""
+    uxgrid = ux.open_grid(mpas_grid)
+
+    for var in uxgrid._ds:
+        # variables should all be np.ndarray
+        assert isinstance(uxgrid._ds[var].data, np.ndarray)
+
+    # chunk every data variable
+    uxgrid.chunk(n_node=1, n_face=2, n_edge=4)
+
+    for var in uxgrid._ds:
+        # variables should all be da.Array
+        assert isinstance(uxgrid._ds[var].data, da.Array)
+
+def test_individual_var_chunking():
+    """Tests the chunking of a single grid variable."""
+    uxgrid = ux.open_grid(mpas_grid)
+
+    # face_node_conn should originally be a numpy array
+    assert isinstance(uxgrid.face_node_connectivity.data, np.ndarray)
+
+    # chunk face_node_connectivity
+    uxgrid.face_node_connectivity = uxgrid.face_node_connectivity.chunk(chunks={"n_face": 16})
+
+    # face_node_conn should now be a dask array
+    assert isinstance(uxgrid.face_node_connectivity.data, da.Array)
+
+
+def test_uxds_chunking():
+    uxds = ux.open_dataset(csne30_grid, csne30_data, chunks={"n_face": 4})
+
+    pass