Support GEOS Cube-Sphere Grids (#802)

* start geos reader

* add geos parser

* update user guide

* update geos reader

* update docs and add docstrings

* update testsS

* add comment about geos dimensions

---------

Co-authored-by: Aaron Zedwick <[email protected]>

docs/user-guide/grid-formats.rst

@@ -17,6 +17,7 @@ NetCDF file format. As of the most recent release, the following grid formats ar
 * SCRIP
 * EXODUS
 * ESMF
+* GEOS CS
 
 While each of these formats can be encoded in the UGRID conventions, the amount of information that is parsed from them
 varies. The following sections describes how each format is converted into the UGRID conventions and what variables
@@ -111,6 +112,17 @@ References
 * https://earthsystemmodeling.org/about/
 * https://earthsystemmodeling.org/docs/release/ESMF_8_0_1/ESMF_refdoc/node3.html#SECTION03028200000000000000
 
+GEOS CS
+=======
+
+The Goddard Earth Observing System (GEOS) Cube Sphere (CS) grid format is equidistant gnomonic cubed-sphere grid with
+6 identical faces that wrap a sphere, with some number of grid cells per face. For example, a C720 GEOS-CS grid has
+6 faces, each with 720x720 elements.
+
+References
+----------
+* https://gmao.gsfc.nasa.gov/gmaoftp/ops/GEOSIT_sample/doc/CS_Description_c180_v1.pdf
+
 Parsed Variables
 ================
 
@@ -157,6 +169,7 @@ Coordinates
        <th>SCRIP</th>
        <th>EXODUS</th>
        <th>ESMF</th>
+       <th>GEOS-CS</th>
      </tr>
      <tr>
        <td>node_latlon</td>
@@ -165,6 +178,7 @@ Coordinates
        <td class="yes-cell">Yes</td>
        <td class="no-cell">No</td>
        <td class="yes-cell">Yes</td>
+       <td class="yes-cell">Yes</td>
      </tr>
      <tr>
        <td>edge_latlon</td>
@@ -173,6 +187,8 @@ Coordinates
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
+
      </tr>
      <tr>
        <td>face_latlon</td>
@@ -181,6 +197,7 @@ Coordinates
        <td class="yes-cell">Yes</td>
        <td class="no-cell">No</td>
        <td class="yes-cell">Yes</td>
+       <td class="yes-cell">Yes</td>
      </tr>
      <tr>
        <td>node_xyz</td>
@@ -189,6 +206,7 @@ Coordinates
        <td class="no-cell">No</td>
        <td class="yes-cell">Yes</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>edge_xyz</td>
@@ -197,6 +215,7 @@ Coordinates
        <td class="no-cell">No</td>
        <td class="yes-cell">Yes</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>face_xyz</td>
@@ -205,6 +224,7 @@ Coordinates
        <td class="no-cell">No</td>
        <td class="yes-cell">Yes</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
    </table>
 
@@ -242,6 +262,7 @@ Connectivity
        <th>SCRIP</th>
        <th>EXODUS</th>
        <th>ESMF</th>
+       <th>GEOS-CS</th>
      </tr>
      <tr>
        <td>face_node</td>
@@ -250,6 +271,7 @@ Connectivity
        <td class="yes-cell">Yes</td>
        <td class="yes-cell">Yes</td>
        <td class="yes-cell">Yes</td>
+       <td class="yes-cell">Yes</td>
      </tr>
      <tr>
        <td>face_edge</td>
@@ -258,6 +280,7 @@ Connectivity
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>face_face</td>
@@ -266,6 +289,7 @@ Connectivity
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>edge_node</td>
@@ -274,6 +298,7 @@ Connectivity
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>edge_edge</td>
@@ -282,6 +307,7 @@ Connectivity
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>edge_face</td>
@@ -290,6 +316,7 @@ Connectivity
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>node_node</td>
@@ -298,6 +325,7 @@ Connectivity
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>node_edge</td>
@@ -306,6 +334,7 @@ Connectivity
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
      <tr>
        <td>node_face</td>
@@ -314,5 +343,6 @@ Connectivity
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
        <td class="no-cell">No</td>
+       <td class="no-cell">No</td>
      </tr>
    </table>

test/test_geos.py

@@ -0,0 +1,31 @@
+import uxarray as ux
+
+import os
+from pathlib import Path
+
+current_path = Path(os.path.dirname(os.path.realpath(__file__)))
+
+gridfile_geos_cs = current_path / "meshfiles" / "geos-cs" / "c12" / "test-c12.native.nc4"
+
+
+
+def test_read_geos_cs_grid():
+    """Tests the conversion of a CS12 GEOS-CS Grid to the UGRID conventions.
+
+    A CS12 grid has 6 faces, each with 12x12 faces and 13x13 nodes each.
+    """
+
+    uxgrid = ux.open_grid(gridfile_geos_cs)
+
+    n_face = 6 * 12 * 12
+    n_node = 6 * 13 * 13
+
+    assert uxgrid.n_face == n_face
+    assert uxgrid.n_node == n_node
+
+
+def test_read_geos_cs_uxds():
+    """Tests the creating of a UxDataset from a CS12 GEOS-CS Grid."""
+    uxds = ux.open_dataset(gridfile_geos_cs, gridfile_geos_cs)
+
+    assert uxds['T'].shape[-1] == uxds.uxgrid.n_face

uxarray/core/utils.py

@@ -7,28 +7,36 @@ def _map_dims_to_ugrid(
     remaps the original dimension name to match the UGRID conventions (i.e.
     "nCell": "n_face")"""
 
-    keys_to_drop = []
-    for key in _source_dims_dict.keys():
-        # obtain all dimensions not present in the original dataset
-        if key not in ds.dims:
-            keys_to_drop.append(key)
+    if grid.source_grid_spec == "GEOS-CS":
+        # stack dimensions to flatten them to map to nodes or faces
+        for var_name in list(ds.coords) + list(ds.data_vars):
+            if all(key in ds[var_name].sizes for key in ["nf", "Ydim", "Xdim"]):
+                ds[var_name] = ds[var_name].stack(n_face=["nf", "Ydim", "Xdim"])
+            if all(key in ds[var_name].sizes for key in ["nf", "YCdim", "XCdim"]):
+                ds[var_name] = ds[var_name].stack(n_node=["nf", "YCdim", "XCdim"])
+    else:
+        keys_to_drop = []
+        for key in _source_dims_dict.keys():
+            # obtain all dimensions not present in the original dataset
+            if key not in ds.dims:
+                keys_to_drop.append(key)
 
-    for key in keys_to_drop:
-        # drop dimensions not present in the original dataset
-        _source_dims_dict.pop(key)
+        for key in keys_to_drop:
+            # drop dimensions not present in the original dataset
+            _source_dims_dict.pop(key)
 
-    for dim in set(ds.dims) ^ _source_dims_dict.keys():
-        # obtain dimensions that were not parsed source_dims_dict and attempt to match to a grid element
-        if ds.sizes[dim] == grid.n_face:
-            _source_dims_dict[dim] = "n_face"
-        elif ds.sizes[dim] == grid.n_node:
-            _source_dims_dict[dim] = "n_node"
-        elif ds.sizes[dim] == grid.n_edge:
-            _source_dims_dict[dim] = "n_edge"
+        for dim in set(ds.dims) ^ _source_dims_dict.keys():
+            # obtain dimensions that were not parsed source_dims_dict and attempt to match to a grid element
+            if ds.sizes[dim] == grid.n_face:
+                _source_dims_dict[dim] = "n_face"
+            elif ds.sizes[dim] == grid.n_node:
+                _source_dims_dict[dim] = "n_node"
+            elif ds.sizes[dim] == grid.n_edge:
+                _source_dims_dict[dim] = "n_edge"
 
-    # Possible Issue: https://github.com/UXARRAY/uxarray/issues/610
+        # Possible Issue: https://github.com/UXARRAY/uxarray/issues/610
 
-    # rename dimensions to follow the UGRID conventions
-    ds = ds.swap_dims(_source_dims_dict)
+        # rename dimensions to follow the UGRID conventions
+        ds = ds.swap_dims(_source_dims_dict)
 
     return ds

uxarray/grid/grid.py

@@ -21,6 +21,7 @@
 from uxarray.io._esmf import _read_esmf
 from uxarray.io._vertices import _read_face_vertices
 from uxarray.io._topology import _read_topology
+from uxarray.io._geos import _read_geos_cs
 
 from uxarray.io.utils import _parse_grid_type
 from uxarray.grid.area import get_all_face_area_from_coords
@@ -198,6 +199,8 @@ def from_dataset(
                 grid_ds, source_dims_dict = _read_mpas(dataset, use_dual=use_dual)
             elif source_grid_spec == "ESMF":
                 grid_ds, source_dims_dict = _read_esmf(dataset)
+            elif source_grid_spec == "GEOS-CS":
+                grid_ds, source_dims_dict = _read_geos_cs(dataset)
             elif source_grid_spec == "Shapefile":
                 raise ValueError("Shapefiles not yet supported")
             else:

uxarray/io/_geos.py

@@ -0,0 +1,59 @@
+import xarray as xr
+import numpy as np
+
+from uxarray.constants import INT_DTYPE
+from uxarray.conventions import ugrid
+
+
+def _read_geos_cs(in_ds: xr.Dataset):
+    """Reads and encodes a GEOS Cube-Sphere grid into the UGRID conventions.
+
+    https://gmao.gsfc.nasa.gov/gmaoftp/ops/GEOSIT_sample/doc/CS_Description_c180_v1.pdf
+    """
+    out_ds = xr.Dataset()
+
+    node_lon = in_ds["corner_lons"].values.ravel()
+    node_lat = in_ds["corner_lats"].values.ravel()
+
+    out_ds["node_lon"] = xr.DataArray(
+        data=node_lon, dims=ugrid.NODE_DIM, attrs=ugrid.NODE_LON_ATTRS
+    )
+
+    out_ds["node_lat"] = xr.DataArray(
+        data=node_lat, dims=ugrid.NODE_DIM, attrs=ugrid.NODE_LAT_ATTRS
+    )
+
+    if "lons" in in_ds:
+        face_lon = in_ds["lons"].values.ravel()
+        face_lat = in_ds["lats"].values.ravel()
+
+        out_ds["face_lon"] = xr.DataArray(
+            data=face_lon, dims=ugrid.FACE_DIM, attrs=ugrid.FACE_LON_ATTRS
+        )
+
+        out_ds["face_lat"] = xr.DataArray(
+            data=face_lat, dims=ugrid.FACE_DIM, attrs=ugrid.FACE_LAT_ATTRS
+        )
+
+    nf, nx, ny = in_ds["corner_lons"].shape
+
+    # generate indices for all corner nodes
+    idx = np.arange(nx * ny * nf, dtype=INT_DTYPE).reshape(nf, nx, ny)
+
+    # calculate indices of corner nodes for each face
+    tl = idx[:, :-1, :-1].reshape(-1)
+    tr = idx[:, :-1, 1:].reshape(-1)
+    bl = idx[:, 1:, :-1].reshape(-1)
+    br = idx[:, 1:, 1:].reshape(-1)
+
+    # Concatenate corner node indices for all faces
+    face_node_connectivity = np.column_stack((br, bl, tl, tr))
+
+    out_ds["face_node_connectivity"] = xr.DataArray(
+        data=face_node_connectivity,
+        dims=ugrid.FACE_NODE_CONNECTIVITY_DIMS,
+        attrs=ugrid.FACE_NODE_CONNECTIVITY_ATTRS,
+    )
+
+    # GEOS-CS does not return a source_dims_dict
+    return out_ds, None

uxarray/io/utils.py

@@ -37,6 +37,9 @@ def _parse_grid_type(dataset):
         mesh_type = "MPAS"
     elif "maxNodePElement" in dataset.dims:
         mesh_type = "ESMF"
+    elif all(key in dataset.sizes for key in ["nf", "YCdim", "XCdim"]):
+        # expected dimensions for a GEOS cube sphere grid
+        mesh_type = "GEOS-CS"
     else:
         raise RuntimeError("Could not recognize dataset format.")
     return mesh_type