MPAS Reader Updates (#1310)

* update mpas reader

* update docstrings and organize code

* update test name

---------

Co-authored-by: Orhan Eroglu <[email protected]>

Author: philipc2

test/test_mpas.py

@@ -18,9 +18,20 @@
 # Fill value
 fv = INT_FILL_VALUE
 
-def test_read_mpas():
-    """Tests execution of _read_mpas()"""
+
+
+def test_read_primal():
     mpas_primal_ugrid, _ = _read_mpas(mpas_xr_ds, use_dual=False)
+
+    assert "face_edge_connectivity" in mpas_primal_ugrid
+    assert "face_node_connectivity" in mpas_primal_ugrid
+    assert "node_edge_connectivity" in mpas_primal_ugrid
+    assert "edge_face_connectivity" in mpas_primal_ugrid
+    assert "node_face_connectivity" in mpas_primal_ugrid
+    assert "edge_node_connectivity" in mpas_primal_ugrid
+    assert "face_face_connectivity" in mpas_primal_ugrid
+
+def test_read_dual():
     mpas_dual_ugrid, _ = _read_mpas(mpas_xr_ds, use_dual=True)
 
 def test_mpas_to_grid():

uxarray/io/_mpas.py

@@ -51,6 +51,9 @@ def _primal_to_ugrid(in_ds, out_ds):
     _parse_face_nodes(in_ds, out_ds, mesh_type="primal")
     _parse_node_faces(in_ds, out_ds, mesh_type="primal")
 
+    if "edgesOnVertex" in in_ds:
+        _parse_node_edges(in_ds, out_ds, mesh_type="primal")
+
     if "verticesOnEdge" in in_ds:
         _parse_edge_nodes(in_ds, out_ds, "primal")
         source_dims_dict[in_ds["verticesOnEdge"].dims[0]] = "n_edge"
@@ -68,7 +71,7 @@ def _primal_to_ugrid(in_ds, out_ds):
         _parse_edge_face_distances(in_ds, out_ds)
 
     if "cellsOnCell" in in_ds:
-        _parse_face_faces(in_ds, out_ds)
+        _parse_face_faces(in_ds, out_ds, mesh_type="primal")
 
     if "areaCell" in in_ds:
         _parse_face_areas(in_ds, out_ds, mesh_type="primal")
@@ -112,6 +115,9 @@ def _dual_to_ugrid(in_ds, out_ds):
     _parse_face_nodes(in_ds, out_ds, mesh_type="dual")
     _parse_node_faces(in_ds, out_ds, mesh_type="dual")
 
+    if "edgesOnCell" in in_ds:
+        _parse_node_edges(in_ds, out_ds, mesh_type="dual")
+
     if "cellsOnEdge" in in_ds:
         _parse_edge_nodes(in_ds, out_ds, mesh_type="dual")
         source_dims_dict[in_ds["cellsOnEdge"].dims[0]] = "n_edge"
@@ -266,141 +272,112 @@ def _parse_edge_xyz_coords(in_ds, out_ds, mesh_type):
 
 
 def _parse_face_nodes(in_ds, out_ds, mesh_type):
-    """Parses face node connectivity for either the Primal or Dual Mesh."""
-    if mesh_type == "primal":
-        verticesOnCell = in_ds["verticesOnCell"].astype(INT_DTYPE)
-        nEdgesOnCell = in_ds["nEdgesOnCell"].astype(INT_DTYPE)
-
-        # Replace padded values with fill values
-        verticesOnCell = _replace_padding(verticesOnCell, nEdgesOnCell)
-
-        # Replace missing/zero values with fill values
-        verticesOnCell = _replace_zeros(verticesOnCell)
-
-        # Convert to zero-indexed
-        face_nodes = _to_zero_index(verticesOnCell)
+    """Parses face node connectivity for either the primal or dual mesh."""
+    key = "verticesOnCell" if mesh_type == "primal" else "cellsOnVertex"
+    arr = in_ds[key].astype(INT_DTYPE)
 
-    else:
-        cellsOnVertex = in_ds["cellsOnVertex"].astype(INT_DTYPE)
-
-        # Replace missing/zero values with fill values
-        cellsOnVertex = _replace_zeros(cellsOnVertex)
+    if mesh_type == "primal":
+        count = in_ds["nEdgesOnCell"].astype(INT_DTYPE)
+        arr = _replace_padding(arr, count)
 
-        # Convert to zero-indexed
-        face_nodes = _to_zero_index(cellsOnVertex)
+    arr = _replace_zeros(arr)
+    face_nodes = _to_zero_index(arr)
 
     out_ds["face_node_connectivity"] = face_nodes.assign_attrs(
         ugrid.FACE_NODE_CONNECTIVITY_ATTRS
     ).rename(dict(zip(face_nodes.dims, ugrid.FACE_NODE_CONNECTIVITY_DIMS)))
 
 
 def _parse_edge_nodes(in_ds, out_ds, mesh_type):
-    """Parses edge node connectivity for either the Primal or Dual Mesh."""
-    if mesh_type == "primal":
-        verticesOnEdge = in_ds["verticesOnEdge"].astype(INT_DTYPE)
-
-        # Replace missing/zero values with fill values
-        verticesOnEdge = _replace_zeros(verticesOnEdge)
-
-        # Convert to zero-indexed
-        edge_nodes = _to_zero_index(verticesOnEdge)
-
-    else:
-        cellsOnEdge = in_ds["cellsOnEdge"].astype(INT_DTYPE)
+    """Parses edge node connectivity for either the primal or dual mesh."""
+    key = "verticesOnEdge" if mesh_type == "primal" else "cellsOnEdge"
+    arr = in_ds[key].astype(INT_DTYPE)
 
-        # Replace missing/zero values with fill values
-        cellsOnEdge = _replace_zeros(cellsOnEdge)
-
-        # Convert to zero-indexed
-        edge_nodes = _to_zero_index(cellsOnEdge)
+    arr = _replace_zeros(arr)
+    edge_nodes = _to_zero_index(arr)
 
     out_ds["edge_node_connectivity"] = edge_nodes.assign_attrs(
         ugrid.EDGE_NODE_CONNECTIVITY_ATTRS
     ).rename(dict(zip(edge_nodes.dims, ugrid.EDGE_NODE_CONNECTIVITY_DIMS)))
 
 
 def _parse_node_faces(in_ds, out_ds, mesh_type):
-    """Parses node face connectivity for either the Primal or Dual Mesh."""
-    if mesh_type == "primal":
-        cellsOnVertex = in_ds["cellsOnVertex"].astype(INT_DTYPE)
-
-        # Replace missing/zero values with fill values
-        cellsOnVertex = _replace_zeros(cellsOnVertex)
-
-        # Convert to zero-indexed
-        node_faces = _to_zero_index(cellsOnVertex)
-    else:
-        verticesOnCell = in_ds["verticesOnCell"].astype(INT_DTYPE)
-        nEdgesOnCell = in_ds["nEdgesOnCell"].astype(INT_DTYPE)
-
-        # Replace padded values with fill values
-        verticesOnCell = _replace_padding(verticesOnCell, nEdgesOnCell)
+    """Parses node face connectivity for either the primal or dual mesh."""
+    key = "cellsOnVertex" if mesh_type == "primal" else "verticesOnCell"
+    arr = in_ds[key].astype(INT_DTYPE)
 
-        # Replace missing/zero values with fill values
-        verticesOnCell = _replace_zeros(verticesOnCell)
+    if mesh_type != "primal":
+        count = in_ds["nEdgesOnCell"].astype(INT_DTYPE)
+        arr = _replace_padding(arr, count)
 
-        # Convert to zero-indexed
-        node_faces = _to_zero_index(verticesOnCell)
+    arr = _replace_zeros(arr)
+    node_faces = _to_zero_index(arr)
 
     out_ds["node_face_connectivity"] = node_faces.assign_attrs(
         ugrid.NODE_FACE_CONNECTIVITY_ATTRS
     ).rename(dict(zip(node_faces.dims, ugrid.NODE_FACE_CONNECTIVITY_DIMS)))
 
 
-def _parse_face_edges(in_ds, out_ds, mesh_type):
-    """Parses face edge connectivity for either the Primal or Dual Mesh."""
-    if mesh_type == "primal":
-        edgesOnCell = in_ds["edgesOnCell"].astype(INT_DTYPE)
-        nEdgesOnCell = in_ds["nEdgesOnCell"].astype(INT_DTYPE)
+def _parse_node_edges(in_ds, out_ds, mesh_type):
+    """Parses node edge connectivity for either the Primal or Dual Mesh."""
+    key = "edgesOnVertex" if mesh_type == "primal" else "edgesOnCell"
+    arr = in_ds[key].astype(INT_DTYPE)
+    arr = _replace_zeros(arr)
+    node_edges = _to_zero_index(arr)
 
-        # Replace padded values with fill values
-        edgesOnCell = _replace_padding(edgesOnCell, nEdgesOnCell)
+    out_ds["node_edge_connectivity"] = node_edges.assign_attrs(
+        ugrid.NODE_EDGE_CONNECTIVITY_ATTRS
+    ).rename(dict(zip(node_edges.dims, ugrid.NODE_EDGE_CONNECTIVITY_DIMS)))
 
-        # Replace missing/zero values with fill values
-        edgesOnCell = _replace_zeros(edgesOnCell)
 
-        # Convert to zero-indexed
-        face_edges = _to_zero_index(edgesOnCell)
-
-    else:
-        edgesOnVertex = in_ds["edgesOnVertex"].astype(INT_DTYPE)
+def _parse_face_edges(in_ds, out_ds, mesh_type):
+    """Parses face edge connectivity for either the primal or dual mesh."""
+    # choose the right source array
+    key = "edgesOnCell" if mesh_type == "primal" else "edgesOnVertex"
+    arr = in_ds[key].astype(INT_DTYPE)
 
-        # Replace missing/zero values with fill values
-        edgesOnVertex = _replace_zeros(edgesOnVertex)
+    if mesh_type == "primal":
+        count = in_ds["nEdgesOnCell"].astype(INT_DTYPE)
+        arr = _replace_padding(arr, count)
 
-        # Convert to zero-indexed
-        face_edges = _to_zero_index(edgesOnVertex)
+    arr = _replace_zeros(arr)
+    face_edges = _to_zero_index(arr)
 
     out_ds["face_edge_connectivity"] = face_edges.assign_attrs(
         ugrid.FACE_EDGE_CONNECTIVITY_ATTRS
     ).rename(dict(zip(face_edges.dims, ugrid.FACE_EDGE_CONNECTIVITY_DIMS)))
 
 
 def _parse_edge_faces(in_ds, out_ds, mesh_type):
-    """Parses edge face connectivity for either the Primal or Dual Mesh."""
-    if mesh_type == "primal":
-        cellsOnEdge = in_ds["cellsOnEdge"].astype(INT_DTYPE)
-
-        # Replace missing/zero values with fill values
-        cellsOnEdge = _replace_zeros(cellsOnEdge)
-
-        # Convert to zero-indexed
-        edge_faces = _to_zero_index(cellsOnEdge)
+    """Parses edge face connectivity for either the primal or dual mesh."""
 
-    else:
-        verticesOnEdge = in_ds["verticesOnEdge"].astype(INT_DTYPE)
-
-        # Replace missing/zero values with fill values
-        verticesOnEdge = _replace_zeros(verticesOnEdge)
+    key = "cellsOnEdge" if mesh_type == "primal" else "verticesOnEdge"
+    arr = in_ds[key].astype(INT_DTYPE)
 
-        # Convert to zero-indexed
-        edge_faces = _to_zero_index(verticesOnEdge)
+    arr = _replace_zeros(arr)
+    edge_faces = _to_zero_index(arr)
 
     out_ds["edge_face_connectivity"] = edge_faces.assign_attrs(
         ugrid.EDGE_FACE_CONNECTIVITY_ATTRS
     ).rename(dict(zip(edge_faces.dims, ugrid.EDGE_FACE_CONNECTIVITY_DIMS)))
 
 
+def _parse_face_faces(in_ds, out_ds, mesh_type):
+    """Parses face face connectivity for the primal mesh."""
+    cellsOnCell = in_ds["cellsOnCell"].astype(INT_DTYPE)
+    nEdgesOnCell = in_ds["nEdgesOnCell"].astype(INT_DTYPE)
+
+    cellsOnCell = _replace_padding(cellsOnCell, nEdgesOnCell)
+
+    cellsOnCell = _replace_zeros(cellsOnCell)
+
+    face_face_connectivity = _to_zero_index(cellsOnCell)
+
+    out_ds["face_face_connectivity"] = face_face_connectivity.assign_attrs(
+        ugrid.FACE_FACE_CONNECTIVITY_ATTRS
+    ).rename(dict(zip(face_face_connectivity.dims, ugrid.FACE_FACE_CONNECTIVITY_DIMS)))
+
+
 def _parse_edge_node_distances(in_ds, out_ds):
     """Parses ``edge_node_distances``"""
     edge_node_distances = in_ds["dvEdge"] / in_ds.attrs["sphere_radius"]
@@ -424,25 +401,6 @@ def _parse_global_attrs(in_ds, out_ds):
     out_ds.attrs = in_ds.attrs
 
 
-def _parse_face_faces(in_ds, out_ds):
-    """Parses face-face connectivity for Primal Mesh."""
-    cellsOnCell = in_ds["cellsOnCell"].astype(INT_DTYPE)
-    nEdgesOnCell = in_ds["nEdgesOnCell"].astype(INT_DTYPE)
-
-    # Replace padded values with fill values
-    cellsOnCell = _replace_padding(cellsOnCell, nEdgesOnCell)
-
-    # Replace missing/zero values with fill values
-    cellsOnCell = _replace_zeros(cellsOnCell)
-
-    # Convert to zero-indexed
-    face_face_connectivity = _to_zero_index(cellsOnCell)
-
-    out_ds["face_face_connectivity"] = face_face_connectivity.assign_attrs(
-        ugrid.FACE_FACE_CONNECTIVITY_ATTRS
-    ).rename(dict(zip(face_face_connectivity.dims, ugrid.FACE_FACE_CONNECTIVITY_DIMS)))
-
-
 def _parse_face_areas(in_ds, out_ds, mesh_type):
     """Parses the face area for either a primal or dual grid."""
     if mesh_type == "primal":
@@ -457,17 +415,16 @@ def _parse_face_areas(in_ds, out_ds, mesh_type):
 
 def _parse_boundary_node_indices(in_ds, out_ds, mesh_type):
     """Parses the boundary node indices."""
+
     boundary_node_mask = in_ds["boundaryVertex"]
     boundary_node_indices = boundary_node_mask.where(boundary_node_mask).dropna(
         dim=boundary_node_mask.dims[0]
     )
 
-    # Convert to integer indices
     boundary_node_indices = boundary_node_indices.coords[
         boundary_node_indices.dims[0]
     ].astype(INT_DTYPE)
 
-    # Ensure zero-indexed
     boundary_node_indices = boundary_node_indices - 1
 
     out_ds["boundary_node_indices"] = boundary_node_indices.rename(