From a05fb784b88467dd7e1b13cf726a64767691906e Mon Sep 17 00:00:00 2001 From: Aaron Zedwick Date: Fri, 6 Dec 2024 15:05:46 -0600 Subject: [PATCH] Initial work on minimal reader support --- test/test_mpas.py | 6 +++ uxarray/core/api.py | 7 ++- uxarray/grid/grid.py | 24 ++++++--- uxarray/io/_esmf.py | 6 ++- uxarray/io/_geos.py | 15 +++++- uxarray/io/_icon.py | 125 ++++++++++++++++++++++++------------------- uxarray/io/_mpas.py | 70 +++++++++++++----------- uxarray/io/_scrip.py | 35 ++++++------ uxarray/io/_ugrid.py | 6 +-- 9 files changed, 173 insertions(+), 121 deletions(-) diff --git a/test/test_mpas.py b/test/test_mpas.py index 5833e0f85..9f3af4de9 100644 --- a/test/test_mpas.py +++ b/test/test_mpas.py @@ -140,3 +140,9 @@ def test_face_area(self): assert "face_areas" in uxgrid_primal._ds assert "face_areas" in uxgrid_dual._ds + + def test_minimal(self): + """Tests the minimal grid reader""" + uxgrid = ux.open_grid(self.mpas_grid_path, minimal=True) + + assert "node_x" not in uxgrid._ds diff --git a/uxarray/core/api.py b/uxarray/core/api.py index 4acb94fdd..4bb57ca7e 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -19,6 +19,7 @@ def open_grid( ], latlon: Optional[bool] = False, use_dual: Optional[bool] = False, + minimal: Optional[bool] = False, **kwargs: Dict[str, Any], ) -> Grid: """Constructs and returns a ``Grid`` from a grid file. @@ -34,11 +35,13 @@ def open_grid( object to define the grid. latlon : bool, optional - Specify if the grid is lat/lon based + Specify if the grid is lat/lon based use_dual: bool, optional Specify whether to use the primal (use_dual=False) or dual (use_dual=True) mesh if the file type is mpas + minimal: bool, optional + Specify whether to read the minimal information (`nodes` and `face_node_connectivity`) needed for a grid **kwargs : Dict[str, Any] Additional arguments passed on to ``xarray.open_dataset``. Refer to the [xarray @@ -84,7 +87,7 @@ def open_grid( try: grid_ds = xr.open_dataset(grid_filename_or_obj, **kwargs) - uxgrid = Grid.from_dataset(grid_ds, use_dual=use_dual) + uxgrid = Grid.from_dataset(grid_ds, use_dual=use_dual, minimal=minimal) except ValueError: raise ValueError("Inputted grid_filename_or_obj not supported.") diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 3453bc6da..b028aa2e6 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -241,7 +241,11 @@ def __init__( @classmethod def from_dataset( - cls, dataset: xr.Dataset, use_dual: Optional[bool] = False, **kwargs + cls, + dataset: xr.Dataset, + use_dual: Optional[bool] = False, + minimal: Optional[bool] = False, + **kwargs, ): """Constructs a ``Grid`` object from an ``xarray.Dataset``. @@ -251,6 +255,8 @@ def from_dataset( ``xarray.Dataset`` containing unstructured grid coordinates and connectivity variables use_dual : bool, default=False When reading in MPAS formatted datasets, indicates whether to use the Dual Mesh + minimal : bool, default=False + Specify whether to read the minimal information (`nodes` and `face_node_connectivity`) needed for a grid """ if not isinstance(dataset, xr.Dataset): raise ValueError("Input must be an xarray.Dataset") @@ -264,17 +270,21 @@ def from_dataset( if source_grid_spec == "Exodus": grid_ds, source_dims_dict = _read_exodus(dataset) elif source_grid_spec == "Scrip": - grid_ds, source_dims_dict = _read_scrip(dataset) + grid_ds, source_dims_dict = _read_scrip(dataset, minimal) elif source_grid_spec == "UGRID": - grid_ds, source_dims_dict = _read_ugrid(dataset) + grid_ds, source_dims_dict = _read_ugrid(dataset, minimal=minimal) elif source_grid_spec == "MPAS": - grid_ds, source_dims_dict = _read_mpas(dataset, use_dual=use_dual) + grid_ds, source_dims_dict = _read_mpas( + dataset, use_dual=use_dual, minimal=minimal + ) elif source_grid_spec == "ESMF": - grid_ds, source_dims_dict = _read_esmf(dataset) + grid_ds, source_dims_dict = _read_esmf(dataset, minimal=minimal) elif source_grid_spec == "GEOS-CS": - grid_ds, source_dims_dict = _read_geos_cs(dataset) + grid_ds, source_dims_dict = _read_geos_cs(dataset, minimal=minimal) elif source_grid_spec == "ICON": - grid_ds, source_dims_dict = _read_icon(dataset, use_dual=use_dual) + grid_ds, source_dims_dict = _read_icon( + dataset, use_dual=use_dual, minimal=minimal + ) elif source_grid_spec == "Structured": grid_ds = _read_structured_grid(dataset[lon_name], dataset[lat_name]) source_dims_dict = {"n_face": (lon_name, lat_name)} diff --git a/uxarray/io/_esmf.py b/uxarray/io/_esmf.py index f16abc657..ead569289 100644 --- a/uxarray/io/_esmf.py +++ b/uxarray/io/_esmf.py @@ -6,7 +6,7 @@ from uxarray.conventions import ugrid -def _read_esmf(in_ds): +def _read_esmf(in_ds, minimal=False): """Reads in an Xarray dataset containing an ESMF formatted Grid dataset and encodes it in the UGRID conventions. @@ -27,6 +27,8 @@ def _read_esmf(in_ds): ---------- in_ds: xr.Dataset ESMF Grid Dataset + minimal : bool, optional + Specify whether to read the minimal information (`nodes` and `face_node_connectivity`) needed for a grid Returns ------- @@ -56,7 +58,7 @@ def _read_esmf(in_ds): node_lat, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LAT_ATTRS ) - if "centerCoords" in in_ds: + if "centerCoords" in in_ds and not minimal: # parse center coords (face centers) if available face_lon = in_ds["centerCoords"].isel(coordDim=0).values out_ds[ugrid.FACE_COORDINATES[0]] = xr.DataArray( diff --git a/uxarray/io/_geos.py b/uxarray/io/_geos.py index af0de78ee..19e70a89e 100644 --- a/uxarray/io/_geos.py +++ b/uxarray/io/_geos.py @@ -5,8 +5,19 @@ from uxarray.conventions import ugrid -def _read_geos_cs(in_ds: xr.Dataset): +def _read_geos_cs(in_ds: xr.Dataset, minimal=False): """Reads and encodes a GEOS Cube-Sphere grid into the UGRID conventions. + Parameters + ---------- + in_ds: xr.Dataset + GEOS_CS Grid Dataset + minimal : bool, optional + Specify whether to read the minimal information (`nodes` and `face_node_connectivity`) needed for a grid + + Returns + ------- + out_ds: xr.Dataset + GEOS_CS Grid encoder in the UGRID conventions https://gmao.gsfc.nasa.gov/gmaoftp/ops/GEOSIT_sample/doc/CS_Description_c180_v1.pdf """ @@ -23,7 +34,7 @@ def _read_geos_cs(in_ds: xr.Dataset): data=node_lat, dims=ugrid.NODE_DIM, attrs=ugrid.NODE_LAT_ATTRS ) - if "lons" in in_ds: + if "lons" in in_ds and not minimal: face_lon = in_ds["lons"].values.ravel() face_lat = in_ds["lats"].values.ravel() diff --git a/uxarray/io/_icon.py b/uxarray/io/_icon.py index 01ed891f2..4539d8097 100644 --- a/uxarray/io/_icon.py +++ b/uxarray/io/_icon.py @@ -3,8 +3,20 @@ import numpy as np -def _primal_to_ugrid(in_ds, out_ds): - """Encodes the Primal Mesh of an ICON Grid into the UGRID conventions.""" +def _primal_to_ugrid(in_ds, out_ds, minimal=False): + """Encodes the Primal Mesh of an ICON Grid into the UGRID conventions. + Parameters + ---------- + in_ds: xr.Dataset + ICON Grid Dataset + minimal : bool, optional + Specify whether to read the minimal information (`nodes` and `face_node_connectivity`) needed for a grid + + Returns + ------- + out_ds: xr.Dataset + ICON Grid encoder in the UGRID conventions + """ source_dims_dict = {"vertex": "n_node", "edge": "n_edge", "cell": "n_face"} # rename dimensions to match ugrid conventions @@ -21,28 +33,6 @@ def _primal_to_ugrid(in_ds, out_ds): data=node_lat, dims=ugrid.NODE_DIM, attrs=ugrid.NODE_LAT_ATTRS ) - # edge coordinates - edge_lon = np.rad2deg(in_ds["elon"]) - edge_lat = np.rad2deg(in_ds["elat"]) - - out_ds["edge_lon"] = xr.DataArray( - data=edge_lon, dims=ugrid.EDGE_DIM, attrs=ugrid.EDGE_LON_ATTRS - ) - out_ds["edge_lat"] = xr.DataArray( - data=edge_lat, dims=ugrid.EDGE_DIM, attrs=ugrid.EDGE_LAT_ATTRS - ) - - # face coordinates - face_lon = np.rad2deg(in_ds["clon"]) - face_lat = np.rad2deg(in_ds["clat"]) - - 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 - ) - face_node_connectivity = in_ds["vertex_of_cell"].T - 1 out_ds["face_node_connectivity"] = xr.DataArray( @@ -51,45 +41,68 @@ def _primal_to_ugrid(in_ds, out_ds): attrs=ugrid.FACE_NODE_CONNECTIVITY_ATTRS, ) - face_edge_connectivity = in_ds["edge_of_cell"].T - 1 - - out_ds["face_edge_connectivity"] = xr.DataArray( - data=face_edge_connectivity, - dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, - attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, - ) - - face_face_connectivity = in_ds["neighbor_cell_index"].T - 1 - - out_ds["face_face_connectivity"] = xr.DataArray( - data=face_face_connectivity, - dims=ugrid.FACE_FACE_CONNECTIVITY_DIMS, - attrs=ugrid.FACE_FACE_CONNECTIVITY_ATTRS, - ) - - edge_face_connectivity = in_ds["adjacent_cell_of_edge"].T - 1 - - out_ds["edge_face_connectivity"] = xr.DataArray( - data=edge_face_connectivity, - dims=ugrid.EDGE_FACE_CONNECTIVITY_DIMS, - attrs=ugrid.EDGE_FACE_CONNECTIVITY_ATTRS, - ) - - edge_node_connectivity = in_ds["edge_vertices"].T - 1 - out_ds["edge_node_connectivity"] = xr.DataArray( - data=edge_node_connectivity, - dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, - attrs=ugrid.EDGE_NODE_CONNECTIVITY_ATTRS, - ) + if not minimal: + # edge coordinates + edge_lon = np.rad2deg(in_ds["elon"]) + edge_lat = np.rad2deg(in_ds["elat"]) + + out_ds["edge_lon"] = xr.DataArray( + data=edge_lon, dims=ugrid.EDGE_DIM, attrs=ugrid.EDGE_LON_ATTRS + ) + out_ds["edge_lat"] = xr.DataArray( + data=edge_lat, dims=ugrid.EDGE_DIM, attrs=ugrid.EDGE_LAT_ATTRS + ) + + # face coordinates + face_lon = np.rad2deg(in_ds["clon"]) + face_lat = np.rad2deg(in_ds["clat"]) + + 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 + ) + + face_edge_connectivity = in_ds["edge_of_cell"].T - 1 + + out_ds["face_edge_connectivity"] = xr.DataArray( + data=face_edge_connectivity, + dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, + ) + + face_face_connectivity = in_ds["neighbor_cell_index"].T - 1 + + out_ds["face_face_connectivity"] = xr.DataArray( + data=face_face_connectivity, + dims=ugrid.FACE_FACE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_FACE_CONNECTIVITY_ATTRS, + ) + + edge_face_connectivity = in_ds["adjacent_cell_of_edge"].T - 1 + + out_ds["edge_face_connectivity"] = xr.DataArray( + data=edge_face_connectivity, + dims=ugrid.EDGE_FACE_CONNECTIVITY_DIMS, + attrs=ugrid.EDGE_FACE_CONNECTIVITY_ATTRS, + ) + + edge_node_connectivity = in_ds["edge_vertices"].T - 1 + out_ds["edge_node_connectivity"] = xr.DataArray( + data=edge_node_connectivity, + dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.EDGE_NODE_CONNECTIVITY_ATTRS, + ) return out_ds, source_dims_dict -def _read_icon(ext_ds, use_dual=False): +def _read_icon(ext_ds, use_dual=False, minimal=False): """Reads and encodes an ICON mesh into the UGRID conventions.""" out_ds = xr.Dataset() if not use_dual: - return _primal_to_ugrid(ext_ds, out_ds) + return _primal_to_ugrid(ext_ds, out_ds, minimal) else: raise ValueError("Conversion of the ICON Dual mesh is not yet supported.") diff --git a/uxarray/io/_mpas.py b/uxarray/io/_mpas.py index 539c596cf..6ece67628 100644 --- a/uxarray/io/_mpas.py +++ b/uxarray/io/_mpas.py @@ -5,7 +5,7 @@ from uxarray.conventions import ugrid, descriptors -def _primal_to_ugrid(in_ds, out_ds): +def _primal_to_ugrid(in_ds, out_ds, minimal=False): """Encodes the MPAS Primal-Mesh in the UGRID conventions. Parameters @@ -15,55 +15,59 @@ def _primal_to_ugrid(in_ds, out_ds): out_ds : xarray.Dataset Output dataset where the MPAS Primal-Mesh is encoded in the UGRID conventions + minimal : bool + Specify whether to read the minimal information (`nodes` and `face_node_connectivity`) needed for a grid """ source_dims_dict = {} + _parse_face_nodes(in_ds, out_ds, mesh_type="primal") + if "lonVertex" in in_ds: _parse_node_latlon_coords(in_ds, out_ds, mesh_type="primal") - if "xVertex" in in_ds: - _parse_node_xyz_coords(in_ds, out_ds, mesh_type="primal") + if not minimal: + if "xVertex" in in_ds: + _parse_node_xyz_coords(in_ds, out_ds, mesh_type="primal") - if "lonCell" in in_ds: - _parse_face_latlon_coords(in_ds, out_ds, mesh_type="primal") + if "lonCell" in in_ds: + _parse_face_latlon_coords(in_ds, out_ds, mesh_type="primal") - if "xCell" in in_ds: - _parse_face_xyz_coords(in_ds, out_ds, mesh_type="primal") + if "xCell" in in_ds: + _parse_face_xyz_coords(in_ds, out_ds, mesh_type="primal") - if "lonEdge" in in_ds: - _parse_edge_latlon_coords(in_ds, out_ds, mesh_type="primal") + if "lonEdge" in in_ds: + _parse_edge_latlon_coords(in_ds, out_ds, mesh_type="primal") - if "xEdge" in in_ds: - _parse_edge_xyz_coords(in_ds, out_ds, mesh_type="primal") + if "xEdge" in in_ds: + _parse_edge_xyz_coords(in_ds, out_ds, mesh_type="primal") - _parse_face_nodes(in_ds, out_ds, mesh_type="primal") - _parse_node_faces(in_ds, out_ds, mesh_type="primal") + _parse_node_faces(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" + if "verticesOnEdge" in in_ds: + _parse_edge_nodes(in_ds, out_ds, "primal") + source_dims_dict[in_ds["verticesOnEdge"].dims[0]] = "n_edge" - if "edgesOnCell" in in_ds: - _parse_face_edges(in_ds, out_ds, mesh_type="primal") + if "edgesOnCell" in in_ds: + _parse_face_edges(in_ds, out_ds, mesh_type="primal") - if "cellsOnEdge" in in_ds: - _parse_edge_faces(in_ds, out_ds, mesh_type="primal") + if "cellsOnEdge" in in_ds: + _parse_edge_faces(in_ds, out_ds, mesh_type="primal") - if "dvEdge" in in_ds: - _parse_edge_node_distances(in_ds, out_ds) + if "dvEdge" in in_ds: + _parse_edge_node_distances(in_ds, out_ds) - if "dcEdge" in in_ds: - _parse_edge_face_distances(in_ds, out_ds) + if "dcEdge" in in_ds: + _parse_edge_face_distances(in_ds, out_ds) - if "cellsOnCell" in in_ds: - _parse_face_faces(in_ds, out_ds) + if "cellsOnCell" in in_ds: + _parse_face_faces(in_ds, out_ds) - if "areaCell" in in_ds: - _parse_face_areas(in_ds, out_ds, mesh_type="primal") + if "areaCell" in in_ds: + _parse_face_areas(in_ds, out_ds, mesh_type="primal") - if "boundaryVertex" in in_ds: - _parse_boundary_node_indices(in_ds, out_ds, mesh_type="primal") + if "boundaryVertex" in in_ds: + _parse_boundary_node_indices(in_ds, out_ds, mesh_type="primal") # set global attributes _parse_global_attrs(in_ds, out_ds) @@ -600,7 +604,7 @@ def _to_zero_index(grid_var): return grid_var -def _read_mpas(ext_ds, use_dual=False): +def _read_mpas(ext_ds, use_dual=False, minimal=False): """Function to read in a MPAS Grid dataset and encode either the Primal or Dual Mesh in the UGRID conventions. @@ -613,6 +617,8 @@ def _read_mpas(ext_ds, use_dual=False): MPAS datafile of interest use_dual : bool, optional Flag to select whether to encode the Dual-Mesh. Defaults to False + minimal : bool, optional + Specify whether to read the minimal information (`nodes` and `face_node_connectivity`) needed for a grid Returns ------- @@ -628,6 +634,6 @@ def _read_mpas(ext_ds, use_dual=False): source_dim_map = _dual_to_ugrid(ext_ds, ds) # convert primal-mesh to UGRID else: - source_dim_map = _primal_to_ugrid(ext_ds, ds) + source_dim_map = _primal_to_ugrid(ext_ds, ds, minimal=minimal) return ds, source_dim_map diff --git a/uxarray/io/_scrip.py b/uxarray/io/_scrip.py index c07824147..6bb07a455 100644 --- a/uxarray/io/_scrip.py +++ b/uxarray/io/_scrip.py @@ -7,7 +7,7 @@ from uxarray.conventions import ugrid -def _to_ugrid(in_ds, out_ds): +def _to_ugrid(in_ds, out_ds, minimal=False): """If input dataset (``in_ds``) file is an unstructured SCRIP file, function will reassign SCRIP variables to UGRID conventions in output file (``out_ds``). @@ -16,10 +16,11 @@ def _to_ugrid(in_ds, out_ds): ---------- in_ds : xarray.Dataset Original scrip dataset of interest being used - out_ds : xarray.Variable file to be returned by ``_populate_scrip_data``, used as an empty placeholder file to store reassigned SCRIP variables in UGRID conventions + minimal : bool, optional + Specify whether to read the minimal information (`nodes` and `face_node_connectivity`) needed for a grid """ source_dims_dict = {} @@ -55,19 +56,19 @@ def _to_ugrid(in_ds, out_ds): out_ds[ugrid.NODE_COORDINATES[1]] = xr.DataArray( unq_lat, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LAT_ATTRS ) - - # Create face_lon & face_lat from grid_center_lat/lon - out_ds[ugrid.FACE_COORDINATES[0]] = xr.DataArray( - in_ds["grid_center_lon"].values, - dims=[ugrid.FACE_DIM], - attrs=ugrid.FACE_LON_ATTRS, - ) - - out_ds[ugrid.FACE_COORDINATES[1]] = xr.DataArray( - in_ds["grid_center_lat"].values, - dims=[ugrid.FACE_DIM], - attrs=ugrid.FACE_LAT_ATTRS, - ) + if not minimal: + # Create face_lon & face_lat from grid_center_lat/lon + out_ds[ugrid.FACE_COORDINATES[0]] = xr.DataArray( + in_ds["grid_center_lon"].values, + dims=[ugrid.FACE_DIM], + attrs=ugrid.FACE_LON_ATTRS, + ) + + out_ds[ugrid.FACE_COORDINATES[1]] = xr.DataArray( + in_ds["grid_center_lat"].values, + dims=[ugrid.FACE_DIM], + attrs=ugrid.FACE_LAT_ATTRS, + ) # standardize fill values and data type face nodes face_nodes = _replace_fill_values( @@ -90,7 +91,7 @@ def _to_ugrid(in_ds, out_ds): return source_dims_dict -def _read_scrip(ext_ds): +def _read_scrip(ext_ds, minimal): """Function to reassign lat/lon variables to node variables. Currently, supports unstructured SCRIP grid files following traditional SCRIP @@ -114,7 +115,7 @@ def _read_scrip(ext_ds): """ ds = xr.Dataset() - source_dims_dict = _to_ugrid(ext_ds, ds) + source_dims_dict = _to_ugrid(ext_ds, ds, minimal=minimal) return ds, source_dims_dict diff --git a/uxarray/io/_ugrid.py b/uxarray/io/_ugrid.py index 3629a843b..ca652c936 100644 --- a/uxarray/io/_ugrid.py +++ b/uxarray/io/_ugrid.py @@ -7,7 +7,7 @@ import uxarray.conventions.ugrid as ugrid -def _read_ugrid(ds): +def _read_ugrid(ds, minimal=False): """Parses an unstructured grid dataset and encodes it in the UGRID conventions.""" @@ -22,13 +22,13 @@ def _read_ugrid(ds): node_lat_name: ugrid.NODE_COORDINATES[1], } - if "edge_coordinates" in ds["grid_topology"].attrs: + if "edge_coordinates" in ds["grid_topology"].attrs and not minimal: # get the names of edge_lon and edge_lat, if they exist edge_lon_name, edge_lat_name = ds["grid_topology"].edge_coordinates.split() coord_dict[edge_lon_name] = ugrid.EDGE_COORDINATES[0] coord_dict[edge_lat_name] = ugrid.EDGE_COORDINATES[1] - if "face_coordinates" in ds["grid_topology"].attrs: + if "face_coordinates" in ds["grid_topology"].attrs and not minimal: # get the names of face_lon and face_lat, if they exist face_lon_name, face_lat_name = ds["grid_topology"].face_coordinates.split() coord_dict[face_lon_name] = ugrid.FACE_COORDINATES[0]