From 545a18da03bfeabf3b6981727f129f8c6454b230 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Date: Mon, 1 Apr 2024 17:17:53 -0500 Subject: [PATCH] UGRID Conformance, Grid Construction from Minimal Grid Variables (#628) * initial work * work on convetnions * topology reader * work on topo, getitem for grid * work on conventions * run pre-commit * fix from face vertices * add sizes, dims, coordinates, and connectivity properties * reader * work on ugrid reader * use ugrid conventions module in esmf reader * update scrip reader * add test for from topo * update topology reader * update mpas reader * update exodus reader * update ugrid reader * fix tests after refactor * fix tests * update connectivity.py with refactor * update connectivity.py with refactor * investigate failing test * investigate failing test * investigate failing test * fix failing connectivity tests * testing * update grid.py * update user api * update face areas and add grid.to_netcdf * update ValueError * update internall api * add docstring to __getitem__ * add edge dimension to grid topology * docstring for Grid.to_netcdf * remove Grid.to_netcdf in favor of Grid.to_xarray * update repr * fix failing test * update docstrings * update repr * coments * fix failing test * remove leftover method in api * specify Cartesiain in xyz coordinates * add cartesiain coordinate names * update docstrings * update docstrings * update conventions --------- Co-authored-by: Hongyu Chen Co-authored-by: Orhan Eroglu <32553057+erogluorhan@users.noreply.github.com> --- docs/user_api/index.rst | 14 +- test/constants.py | 2 +- test/test_api.py | 2 +- test/test_from_topology.py | 87 +++++ test/test_grid.py | 32 +- uxarray/conventions/__init__.py | 0 uxarray/conventions/descriptors.py | 18 + uxarray/conventions/ugrid.py | 329 ++++++++++++++++++ uxarray/core/api.py | 25 +- uxarray/grid/connectivity.py | 44 +-- uxarray/grid/coordinates.py | 86 ++--- uxarray/grid/grid.py | 515 +++++++++++++++++------------ uxarray/io/_esmf.py | 70 ++-- uxarray/io/_exodus.py | 76 +---- uxarray/io/_mpas.py | 173 +++------- uxarray/io/_scrip.py | 45 ++- uxarray/io/_topology.py | 66 ++++ uxarray/io/_ugrid.py | 166 ++++++---- 18 files changed, 1073 insertions(+), 677 deletions(-) create mode 100644 test/test_from_topology.py create mode 100644 uxarray/conventions/__init__.py create mode 100644 uxarray/conventions/descriptors.py create mode 100644 uxarray/conventions/ugrid.py create mode 100644 uxarray/io/_topology.py diff --git a/docs/user_api/index.rst b/docs/user_api/index.rst index 9128262c5..a84e856bd 100644 --- a/docs/user_api/index.rst +++ b/docs/user_api/index.rst @@ -220,6 +220,11 @@ Dimensions Grid.n_face Grid.n_max_face_nodes Grid.n_max_face_edges + Grid.n_max_face_faces + Grid.n_max_edge_edges + Grid.n_max_node_faces + Grid.n_max_node_edges + Grid.n_max_node_nodes Grid.n_nodes_per_face Spherical Coordinates @@ -255,13 +260,12 @@ Connectivity :toctree: generated/ Grid.face_node_connectivity - Grid.edge_node_connectivity - Grid.node_node_connectivity Grid.face_edge_connectivity - Grid.edge_edge_connectivity - Grid.node_edge_connectivity Grid.face_face_connectivity + Grid.edge_node_connectivity + Grid.edge_edge_connectivity Grid.edge_face_connectivity + Grid.node_edge_connectivity Grid.node_face_connectivity Grid Descriptors @@ -279,7 +283,7 @@ Attributes :toctree: generated/ Grid.grid_spec - Grid.parsed_attrs + Grid.attrs Plotting diff --git a/test/constants.py b/test/constants.py index 9132ead37..4a54e1b1a 100644 --- a/test/constants.py +++ b/test/constants.py @@ -5,7 +5,7 @@ NNODES_outCSne8 = 386 NNODES_outCSne30 = 5402 NNODES_outRLL1deg = 64442 -DATAVARS_outCSne30 = 2 +DATAVARS_outCSne30 = 4 TRI_AREA = 1.047 # 4*Pi is 12.56 MESH30_AREA = 12.566 diff --git a/test/test_api.py b/test/test_api.py index 6fc7f0d9c..61679f1f7 100644 --- a/test/test_api.py +++ b/test/test_api.py @@ -131,7 +131,7 @@ def test_open_dataset_grid_kwargs(self): """Drops ``Mesh2_face_nodes`` from the inputted grid file using ``grid_kwargs``""" - with self.assertRaises(KeyError): + with self.assertRaises(ValueError): # attempt to open a dataset after dropping face nodes should raise a KeyError uxds = ux.open_dataset( self.gridfile_ne30, diff --git a/test/test_from_topology.py b/test/test_from_topology.py new file mode 100644 index 000000000..878b94b42 --- /dev/null +++ b/test/test_from_topology.py @@ -0,0 +1,87 @@ +import uxarray as ux + +from uxarray.constants import INT_FILL_VALUE +import numpy.testing as nt +import os + +import pytest + +from pathlib import Path + +current_path = Path(os.path.dirname(os.path.realpath(__file__))) + +GRID_PATHS = [ + current_path / 'meshfiles' / "mpas" / "QU" / 'oQU480.231010.nc', + current_path / "meshfiles" / "ugrid" / "geoflow-small" / "grid.nc", + current_path / "meshfiles" / "ugrid" / "outCSne30" / "outCSne30.ug" +] + + + + +def test_minimal_class_method(): + """Tests the minimal required variables for constructing a grid using the + from topology class method.""" + + for grid_path in GRID_PATHS: + uxgrid = ux.open_grid(grid_path) + + uxgrid_ft = ux.Grid.from_topology(node_lon=uxgrid.node_lon.values, + node_lat=uxgrid.node_lat.values, + face_node_connectivity=uxgrid.face_node_connectivity.values, + fill_value=INT_FILL_VALUE, + start_index=0) + + nt.assert_array_equal(uxgrid.node_lon.values, uxgrid_ft.node_lon.values) + nt.assert_array_equal(uxgrid.node_lat.values, uxgrid_ft.node_lat.values) + nt.assert_array_equal(uxgrid.face_node_connectivity.values, uxgrid_ft.face_node_connectivity.values) + + +def test_minimal_api(): + """Tests the minimal required variables for constructing a grid using the + ``ux.open_dataset`` method.""" + + for grid_path in GRID_PATHS: + uxgrid = ux.open_grid(grid_path) + + uxgrid_ft = ux.Grid.from_topology(node_lon=uxgrid.node_lon.values, + node_lat=uxgrid.node_lat.values, + face_node_connectivity=uxgrid.face_node_connectivity.values, + fill_value=INT_FILL_VALUE, + start_index=0) + + grid_topology = {'node_lon': uxgrid.node_lon.values, + 'node_lat': uxgrid.node_lat.values, + 'face_node_connectivity': uxgrid.face_node_connectivity.values, + 'fill_value': INT_FILL_VALUE, + 'start_index': 0} + + uxgrid_ft = ux.open_grid(grid_topology) + + nt.assert_array_equal(uxgrid.node_lon.values, uxgrid_ft.node_lon.values) + nt.assert_array_equal(uxgrid.node_lat.values, uxgrid_ft.node_lat.values) + nt.assert_array_equal(uxgrid.face_node_connectivity.values, uxgrid_ft.face_node_connectivity.values) + + +def test_dataset(): + uxds = ux.open_dataset(GRID_PATHS[0], GRID_PATHS[0]) + + grid_topology = {'node_lon': uxds.uxgrid.node_lon.values, + 'node_lat': uxds.uxgrid.node_lat.values, + 'face_node_connectivity': uxds.uxgrid.face_node_connectivity.values, + 'fill_value': INT_FILL_VALUE, + 'start_index': 0, + "dims_dict" : {"nVertices": "n_node"}} + + + uxds_ft = ux.open_grid(grid_topology, GRID_PATHS[1]) + + uxgrid = uxds.uxgrid + uxgrid_ft = uxds_ft + + + nt.assert_array_equal(uxgrid.node_lon.values, uxgrid_ft.node_lon.values) + nt.assert_array_equal(uxgrid.node_lat.values, uxgrid_ft.node_lat.values) + nt.assert_array_equal(uxgrid.face_node_connectivity.values, uxgrid_ft.face_node_connectivity.values) + + assert uxds_ft.dims == {'n_face', 'n_node', 'n_max_face_nodes'} diff --git a/test/test_grid.py b/test/test_grid.py index b1e654d16..730969383 100644 --- a/test/test_grid.py +++ b/test/test_grid.py @@ -643,6 +643,8 @@ def test_build_face_edges_connectivity_mpas(self): """Tests the construction of (``Mesh2_edge_nodes``) on an MPAS grid with known edge nodes.""" + from uxarray.grid.connectivity import _build_edge_node_connectivity + # grid with known edge node connectivity mpas_grid_ux = ux.open_grid(self.mpas_filepath) edge_nodes_expected = mpas_grid_ux._ds['edge_node_connectivity'].values @@ -652,8 +654,12 @@ def test_build_face_edges_connectivity_mpas(self): edge_nodes_expected = np.unique(edge_nodes_expected, axis=0) # construct edge nodes - _build_edge_node_connectivity(mpas_grid_ux, repopulate=True) - edge_nodes_output = mpas_grid_ux._ds['edge_node_connectivity'].values + edge_nodes_output, _, _ = _build_edge_node_connectivity(mpas_grid_ux.face_node_connectivity.values, + mpas_grid_ux.n_face, + mpas_grid_ux.n_max_face_nodes) + + # _populate_face_edge_connectivity(mpas_grid_ux) + # edge_nodes_output = mpas_grid_ux._ds['edge_node_connectivity'].values self.assertTrue(np.array_equal(edge_nodes_expected, edge_nodes_output)) @@ -705,28 +711,6 @@ def test_build_face_edges_connectivity(self): np.array_equal(reverted_mesh2_edge_nodes[i], original_face_nodes_connectivity[i])) - def test_build_face_edges_connectivity_mpas(self): - tgrid = ux.open_grid(self.mpas_filepath) - - face_node_connectivity = tgrid._ds["face_node_connectivity"] - - _populate_face_edge_connectivity(tgrid) - mesh2_face_edges = tgrid._ds.face_edge_connectivity - mesh2_edge_nodes = tgrid._ds.edge_node_connectivity - - # Assert if the mesh2_face_edges sizes are correct. - self.assertEqual(mesh2_face_edges.sizes["n_face"], - face_node_connectivity.sizes["n_face"]) - self.assertEqual(mesh2_face_edges.sizes["n_max_face_edges"], - face_node_connectivity.sizes["n_max_face_nodes"]) - - # Assert if the mesh2_edge_nodes sizes are correct. - # Euler formular for determining the edge numbers: n_face = n_edges - n_nodes + 2 - num_edges = mesh2_face_edges.sizes["n_face"] + tgrid._ds[ - "node_lon"].sizes["n_node"] - 2 - size = mesh2_edge_nodes.sizes["n_edge"] - self.assertEqual(mesh2_edge_nodes.sizes["n_edge"], num_edges) - def test_build_face_edges_connectivity_fillvalues(self): verts = [ self.f0_deg, self.f1_deg, self.f2_deg, self.f3_deg, self.f4_deg, diff --git a/uxarray/conventions/__init__.py b/uxarray/conventions/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/uxarray/conventions/descriptors.py b/uxarray/conventions/descriptors.py new file mode 100644 index 000000000..14cc2a727 --- /dev/null +++ b/uxarray/conventions/descriptors.py @@ -0,0 +1,18 @@ +DESCRIPTOR_NAMES = ["face_areas", "edge_face_distances", "edge_node_distances"] + + +FACE_AREAS_DIMS = ["n_face"] + +FACE_AREAS_ATTRS = {"cf_role": "face_areas"} + +EDGE_FACE_DISTANCES_DIMS = ["n_edge"] +EDGE_FACE_DISTANCES_ATTRS = { + "cf_role": "edge_face_distances", + "long_name": "Distances between the face centers that " "saddle each edge", +} + +EDGE_NODE_DISTANCES_DIMS = ["n_edge"] +EDGE_NODE_DISTANCES_ATTRS = { + "cf_role": "edge_node_distances", + "long_name": "Distances between the nodes that make up " "each edge.", +} diff --git a/uxarray/conventions/ugrid.py b/uxarray/conventions/ugrid.py new file mode 100644 index 000000000..0353dbe49 --- /dev/null +++ b/uxarray/conventions/ugrid.py @@ -0,0 +1,329 @@ +from uxarray.constants import INT_FILL_VALUE, INT_DTYPE + +CONVENTIONS_ATTR = "UGRID-v1.0" + +# minimum grid topology, additions are made depending on what is present in a grid +BASE_GRID_TOPOLOGY_ATTRS = { + "cf_role": "mesh_topology", + "topology_dimension": 2, + # dimensions + "face_dimension": "n_face", + "node_dimension": "n_node", + # coordinates + "node_coordinates": "node_lon node_lat", + # connectivity + "face_node_connectivity": "face_node_connectivity", +} + +DIM_NAMES = [ + "n_node", + "n_edge", + "n_face", + "n_max_face_nodes", + "n_max_face_edges", + "n_max_face_faces", + "n_max_edge_nodes", + "n_max_edge_edges", + "n_max_edge_faces", + "n_max_node_faces", + "n_max_node_edges", + "two", +] + +NODE_DIM = "n_node" +EDGE_DIM = "n_edge" +FACE_DIM = "n_face" +N_MAX_FACE_NODES_DIM = "n_max_face_nodes" + +NODE_COORDINATES = ["node_lon", "node_lat"] + +# coordinates (spherical) +NODE_LON_ATTRS = { + "standard_name": "longitude", + "long name": "Longitude of the corner nodes of each face", + "units": "degrees_east", +} + +NODE_LAT_ATTRS = { + "standard_name": "latitude", + "long name": "Latitude of the corner nodes of each face", + "units": "degrees_north", +} + +EDGE_COORDINATES = ["edge_lon", "edge_lat"] + +EDGE_LON_ATTRS = { + "standard_name": "longitude", + "long name": "Longitude of the center of each edge", + "units": "degrees_east", +} + +EDGE_LAT_ATTRS = { + "standard_name": "latitude", + "long name": "Latitude of the center of each edge", + "units": "degrees_north", +} + +FACE_COORDINATES = ["face_lon", "face_lat"] + +FACE_LON_ATTRS = { + "standard_name": "longitude", + "long name": "Longitude of the center of each face", + "units": "degrees_east", +} + +FACE_LAT_ATTRS = { + "standard_name": "latitude", + "long name": "Latitude of the center of each face", + "units": "degrees_north", +} + +CARTESIAN_NODE_COORDINATES = ["node_x", "node_y", "node_z"] + +NODE_X_ATTRS = { + "standard_name": "x", + "long name": "Cartesian x location of the corner nodes of each face", + "units": "meters", +} + +NODE_Y_ATTRS = { + "standard_name": "y", + "long name": "Cartesian y location of the corner nodes of each face", + "units": "meters", +} + +NODE_Z_ATTRS = { + "standard_name": "z", + "long name": "Cartesian z location of the corner nodes of each face", + "units": "meters", +} + +CARTESIAN_EDGE_COORDINATES = ["edge_x", "edge_y", "edge_z"] + +EDGE_X_ATTRS = { + "standard_name": "x", + "long name": "Cartesian x location of the center of each edge", + "units": "meters", +} + +EDGE_Y_ATTRS = { + "standard_name": "y", + "long name": "Cartesian y location of the center of each edge", + "units": "meters", +} + +EDGE_Z_ATTRS = { + "standard_name": "z", + "long name": "Cartesian z location of the center of each edge", + "units": "meters", +} + +CARTESIAN_FACE_COORDINATES = ["face_x", "face_y", "face_z"] + +FACE_X_ATTRS = { + "standard_name": "x", + "long name": "Cartesian x location of the center of each face", + "units": "meters", +} + +FACE_Y_ATTRS = { + "standard_name": "y", + "long name": "Cartesian y location of the center of each face", + "units": "meters", +} + +FACE_Z_ATTRS = { + "standard_name": "z", + "long name": "Cartesian z location of the center of each face", + "units": "meters", +} + +# connectivity (face_) +FACE_NODE_CONNECTIVITY_ATTRS = { + "cf_role": "face_node_connectivity", + "long name": "Maps every face to its corner nodes.", + "start_index": 0, + "_FillValue": INT_FILL_VALUE, + "dtype": INT_DTYPE, +} +FACE_NODE_CONNECTIVITY_DIMS = ["n_face", "n_max_face_nodes"] + +FACE_EDGE_CONNECTIVITY_ATTRS = { + "cf_role": "face_edge_connectivity", + "long name": "Maps every face to its edges.", + "start_index": 0, + "_FillValue": INT_FILL_VALUE, + "dtype": INT_DTYPE, +} +FACE_EDGE_CONNECTIVITY_DIMS = [ + "n_face", + "n_max_face_edges", +] # n_max_face_edges equiv to n_max_face_nodes + +FACE_FACE_CONNECTIVITY_ATTRS = { + "cf_role": "face_face_connectivity", + "long name": "Faces that neighbor each face.", + "start_index": 0, + "_FillValue": INT_FILL_VALUE, + "dtype": INT_DTYPE, +} +FACE_FACE_CONNECTIVITY_DIMS = [ + "n_face", + "n_max_face_faces", +] + +# connectivity (edge_) +EDGE_NODE_CONNECTIVITY_ATTRS = { + "cf_role": "edge_node_connectivity", + "long name": "Maps every edge to the two nodes that it connects.", + "start_index": 0, + "dtype": INT_DTYPE, +} +EDGE_NODE_CONNECTIVITY_DIMS = ["n_edge", "two"] + + +# edge_edge_connectivity not yet supported +# EDGE_EDGE_CONNECTIVITY_ATTRS = { +# "cf_role": "edge_edge_connectivity", +# "long name": "Edges that neighbor each edge", +# "start_index": 0, +# "_FillValue": INT_FILL_VALUE, +# "dtype": INT_DTYPE, +# } +# +# EDGE_EDGE_CONNECTIVITY_DIMS = ["n_edge", "n_max_edge_edges"] + +EDGE_FACE_CONNECTIVITY_ATTRS = { + "cf_role": "edge_face_connectivity", + "long name": "Faces that neighbor each edge", + "start_index": 0, + "_FillValue": INT_FILL_VALUE, + "dtype": INT_DTYPE, +} +EDGE_FACE_CONNECTIVITY_DIMS = ["n_edge", "two"] + + +NODE_EDGE_CONNECTIVITY_ATTRS = { + "cf_role": "node_edge_connectivity", + "long name": "Edges that neighbor each node", + "start_index": 0, + "_FillValue": INT_FILL_VALUE, + "dtype": INT_DTYPE, +} + +NODE_EDGE_CONNECTIVITY_DIMS = ["n_node", "n_max_node_edges"] + +NODE_FACE_CONNECTIVITY_ATTRS = { + "cf_role": "node_face_connectivity", + "long name": "Faces that neighbor each node", + "start_index": 0, + "_FillValue": INT_FILL_VALUE, + "dtype": INT_DTYPE, +} + +NODE_FACE_CONNECTIVITY_DIMS = ["n_node", "n_max_node_faces"] + + +N_NODES_PER_FACE_ATTRS = { + "cf_role": "n_nodes_per_face", + "long name": "Number of nodes per face", + "dtype": INT_DTYPE, +} + +N_NODES_PER_FACE_DIMS = ["n_face"] + + +CONNECTIVITY_NAMES = [ + "face_node_connectivity", + "face_edge_connectivity", + "face_face_connectivity", + "edge_node_connectivity", + "edge_face_connectivity", + "node_edge_connectivity", + "node_face_connectivity", +] + +# as of UGRID v1.0 +UGRID_COMPLIANT_CONNECTIVITY_NAMES = [ + "edge_node_connectivity", + "face_node_connectivity", + "face_edge_connectivity", + "edge_face_connectivity", + "face_face_connectivity", +] +CONNECTIVITY = { + "face_node_connectivity": { + "dims": FACE_NODE_CONNECTIVITY_DIMS, + "attrs": FACE_NODE_CONNECTIVITY_ATTRS, + }, + "face_edge_connectivity": { + "dims": FACE_EDGE_CONNECTIVITY_DIMS, + "attrs": FACE_EDGE_CONNECTIVITY_ATTRS, + }, + "face_face_connectivity": { + "dims": FACE_FACE_CONNECTIVITY_DIMS, + "attrs": FACE_FACE_CONNECTIVITY_ATTRS, + }, + "edge_node_connectivity": { + "dims": EDGE_NODE_CONNECTIVITY_DIMS, + "attrs": EDGE_NODE_CONNECTIVITY_ATTRS, + }, + # "edge_edge_connectivity": { + # "dims": EDGE_EDGE_CONNECTIVITY_DIMS, + # "attrs": EDGE_EDGE_CONNECTIVITY_ATTRS, + # }, + "edge_face_connectivity": { + "dims": EDGE_FACE_CONNECTIVITY_DIMS, + "attrs": EDGE_FACE_CONNECTIVITY_ATTRS, + }, + "node_edge_connectivity": { + "dims": NODE_EDGE_CONNECTIVITY_DIMS, + "attrs": NODE_EDGE_CONNECTIVITY_ATTRS, + }, + "node_face_connectivity": { + "dims": NODE_FACE_CONNECTIVITY_DIMS, + "attrs": NODE_FACE_CONNECTIVITY_ATTRS, + }, +} + +SPHERICAL_COORD_NAMES = [ + "node_lon", + "node_lat", + "edge_lon", + "edge_lat", + "face_lon", + "face_lat", +] + +SPHERICAL_COORDS = { + "node_lon": {"dims": [NODE_DIM], "attrs": NODE_LON_ATTRS}, + "node_lat": {"dims": [NODE_DIM], "attrs": NODE_LAT_ATTRS}, + "edge_lon": {"dims": [EDGE_DIM], "attrs": EDGE_LON_ATTRS}, + "edge_lat": {"dims": [EDGE_DIM], "attrs": EDGE_LAT_ATTRS}, + "face_lon": {"dims": [FACE_DIM], "attrs": FACE_LON_ATTRS}, + "face_lat": {"dims": [FACE_DIM], "attrs": FACE_LAT_ATTRS}, +} + +CARTESIAN_COORD_NAMES = [ + "node_x", + "node_y", + "node_z", + "edge_x", + "edge_y", + "edge_z", + "face_x", + "face_y", + "face_z", +] + +CARTESIAN_COORDS = { + "node_x": {"dims": [NODE_DIM], "attrs": NODE_X_ATTRS}, + "node_y": {"dims": [NODE_DIM], "attrs": NODE_Y_ATTRS}, + "node_z": {"dims": [NODE_DIM], "attrs": NODE_Z_ATTRS}, + "edge_x": {"dims": [EDGE_DIM], "attrs": EDGE_X_ATTRS}, + "edge_y": {"dims": [EDGE_DIM], "attrs": EDGE_Y_ATTRS}, + "edge_z": {"dims": [EDGE_DIM], "attrs": EDGE_Z_ATTRS}, + "face_x": {"dims": [FACE_DIM], "attrs": FACE_X_ATTRS}, + "face_y": {"dims": [FACE_DIM], "attrs": FACE_Y_ATTRS}, + "face_z": {"dims": [FACE_DIM], "attrs": FACE_Z_ATTRS}, +} diff --git a/uxarray/core/api.py b/uxarray/core/api.py index 884929443..95bbba5c1 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -15,18 +15,19 @@ def open_grid( grid_filename_or_obj: Union[ - str, os.PathLike, xr.DataArray, np.ndarray, list, tuple + str, os.PathLike, xr.DataArray, np.ndarray, list, tuple, dict ], latlon: Optional[bool] = False, use_dual: Optional[bool] = False, **kwargs: Dict[str, Any], ) -> Grid: - """Creates a ``uxarray.Grid`` object from a grid topology definition. + """Constructs and returns an ``uxarray.Grid`` object from a grid topology + definition. Parameters ---------- - grid_filename_or_obj : string, xarray.Dataset, ndarray, list, tuple, required + grid_filename_or_obj : string, xarray.Dataset, ndarray, list, tuple, dict, required String or Path object as a path to a netCDF file or an OpenDAP URL that stores the unstructured grid topology/definition. It is read similar to ``filename_or_obj`` in ``xarray.open_dataset``. Otherwise, either @@ -56,7 +57,7 @@ def open_grid( Open dataset with a grid topology file - >>> import uxarray as ux + >>> Import uxarray as ux >>> uxgrid = ux.open_grid("grid_filename.g") """ @@ -67,12 +68,16 @@ def open_grid( stacklevel=2, ) - # construct Grid from dataset if isinstance(grid_filename_or_obj, xr.Dataset): + # construct a grid from a dataset file uxgrid = Grid.from_dataset(grid_filename_or_obj, use_dual=use_dual) - # construct Grid from face vertices + elif isinstance(grid_filename_or_obj, dict): + # unpack the dictionary and construct a grid from topology + uxgrid = Grid.from_topology(**grid_filename_or_obj) + elif isinstance(grid_filename_or_obj, (list, tuple, np.ndarray, xr.DataArray)): + # construct Grid from face vertices uxgrid = Grid.from_face_vertices(grid_filename_or_obj, latlon=latlon) # attempt to use Xarray directly for remaining input types @@ -90,7 +95,9 @@ def open_grid( def open_dataset( - grid_filename_or_obj: str, + grid_filename_or_obj: Union[ + str, os.PathLike, xr.DataArray, np.ndarray, list, tuple, dict + ], filename_or_obj: str, latlon: Optional[bool] = False, use_dual: Optional[bool] = False, @@ -172,7 +179,9 @@ def open_dataset( def open_mfdataset( - grid_filename_or_obj: str, + grid_filename_or_obj: Union[ + str, os.PathLike, xr.DataArray, np.ndarray, list, tuple, dict + ], paths: Union[str, os.PathLike], latlon: Optional[bool] = False, use_dual: Optional[bool] = False, diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index c57dc60e6..ee3654e47 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -4,6 +4,7 @@ from scipy import sparse from uxarray.constants import INT_DTYPE, INT_FILL_VALUE +from uxarray.conventions import ugrid from numba import njit @@ -137,8 +138,8 @@ def _populate_n_nodes_per_face(grid): # add to internal dataset grid._ds["n_nodes_per_face"] = xr.DataArray( data=n_nodes_per_face, - dims=["n_face"], - attrs={"long_name": "number of non-fill value nodes for each face"}, + dims=ugrid.N_NODES_PER_FACE_DIMS, + attrs=ugrid.N_NODES_PER_FACE_ATTRS, ) @@ -165,18 +166,13 @@ def _populate_edge_node_connectivity(grid): grid.face_node_connectivity.values, grid.n_face, grid.n_max_face_nodes ) + edge_node_attrs = ugrid.EDGE_NODE_CONNECTIVITY_ATTRS + edge_node_attrs["inverse_indices"] = inverse_indices + edge_node_attrs["fill_value_mask"] = fill_value_mask + # add edge_node_connectivity to internal dataset grid._ds["edge_node_connectivity"] = xr.DataArray( - edge_nodes, - dims=["n_edge", "Two"], - attrs={ - "cf_role": "edge_node_connectivity", - "_FillValue": INT_FILL_VALUE, - "long_name": "Maps every edge to the two nodes that it connects", - "start_index": INT_DTYPE(0), - "inverse_indices": inverse_indices, - "fill_value_mask": fill_value_mask, - }, + edge_nodes, dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, attrs=edge_node_attrs ) @@ -250,12 +246,8 @@ def _populate_edge_face_connectivity(grid): grid._ds["edge_face_connectivity"] = xr.DataArray( data=edge_faces, - dims=["n_edge", "Two"], - attrs={ - "cf_role": "edge_face_connectivity", - "start_index": INT_DTYPE(0), - "long_name": "Maps the faces that saddle a given edge", - }, + dims=ugrid.EDGE_FACE_CONNECTIVITY_DIMS, + attrs=ugrid.EDGE_FACE_CONNECTIVITY_ATTRS, ) @@ -297,12 +289,8 @@ def _populate_face_edge_connectivity(grid): grid._ds["face_edge_connectivity"] = xr.DataArray( data=face_edges, - dims=["n_face", "n_max_face_edges"], - attrs={ - "cf_role": "face_edges_connectivity", - "start_index": INT_DTYPE(0), - "long_name": "Maps every edge to the two nodes that it connects", - }, + dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, ) @@ -323,12 +311,8 @@ def _populate_node_face_connectivity(grid): grid._ds["node_face_connectivity"] = xr.DataArray( node_faces, - dims=["n_node", "n_max_node_faces"], # TODO - attrs={ - "long_name": "Maps every node to the faces that " "it connects", - "nMaxNumFacesPerNode": n_max_faces_per_node, # todo, this possibly duplicates nNodes_per_face - "_FillValue": INT_FILL_VALUE, - }, + dims=ugrid.NODE_FACE_CONNECTIVITY_DIMS, + attrs=ugrid.NODE_FACE_CONNECTIVITY_ATTRS, ) diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py index 23a8ab7b1..2aba490ea 100644 --- a/uxarray/grid/coordinates.py +++ b/uxarray/grid/coordinates.py @@ -7,6 +7,7 @@ from numba import njit, config from uxarray.constants import ENABLE_JIT_CACHE, ENABLE_JIT, ERROR_TOLERANCE +from uxarray.conventions import ugrid config.DISABLE_JIT = not ENABLE_JIT @@ -150,28 +151,13 @@ def _populate_cartesian_xyz_coord(grid): x, y, z = _get_xyz_from_lonlat(grid.node_lon.values, grid.node_lat.values) grid._ds["node_x"] = xr.DataArray( - data=x, - dims=["n_node"], - attrs={ - "standard_name": "cartesian x", - "units": "m", - }, + data=x, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_X_ATTRS ) grid._ds["node_y"] = xr.DataArray( - data=y, - dims=["n_node"], - attrs={ - "standard_name": "cartesian y", - "units": "m", - }, + data=y, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_Y_ATTRS ) grid._ds["node_z"] = xr.DataArray( - data=z, - dims=["n_node"], - attrs={ - "standard_name": "cartesian z", - "units": "m", - }, + data=z, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_Z_ATTRS ) @@ -214,20 +200,10 @@ def _populate_lonlat_coord(grid): # populate dataset grid._ds["node_lon"] = xr.DataArray( - data=lon, - dims=["n_node"], - attrs={ - "long_name": "longitude of corner nodes", - "units": "degrees_east", - }, + data=lon, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LON_ATTRS ) grid._ds["node_lat"] = xr.DataArray( - data=lat, - dims=["n_node"], - attrs={ - "long_name": "latitude of corner nodes", - "units": "degrees_north", - }, + data=lat, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LAT_ATTRS ) @@ -275,23 +251,23 @@ def _populate_face_centroids(grid, repopulate=False): # Populate the centroids if "face_lon" not in grid._ds or repopulate: grid._ds["face_lon"] = xr.DataArray( - centroid_lon, dims=["n_face"], attrs={"units": "degrees_east"} + centroid_lon, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_LON_ATTRS ) grid._ds["face_lat"] = xr.DataArray( - centroid_lat, dims=["n_face"], attrs={"units": "degrees_north"} + centroid_lat, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_LAT_ATTRS ) if "face_x" not in grid._ds or repopulate: grid._ds["face_x"] = xr.DataArray( - centroid_x, dims=["n_face"], attrs={"units": "meters"} + centroid_x, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_X_ATTRS ) grid._ds["face_y"] = xr.DataArray( - centroid_y, dims=["n_face"], attrs={"units": "meters"} + centroid_y, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_Y_ATTRS ) grid._ds["face_z"] = xr.DataArray( - centroid_z, dims=["n_face"], attrs={"units": "meters"} + centroid_z, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_Z_ATTRS ) @@ -361,53 +337,31 @@ def _populate_edge_centroids(grid, repopulate=False): # Populate the centroids if "edge_lon" not in grid._ds or repopulate: grid._ds["edge_lon"] = xr.DataArray( - centroid_lon, - dims=["n_edge"], - attrs={ - "standard_name": "longitude", - "long name": "Longitude of the center of each edge", - "units": "degrees_east", - }, + centroid_lon, dims=[ugrid.EDGE_DIM], attrs=ugrid.EDGE_LON_ATTRS ) grid._ds["edge_lat"] = xr.DataArray( centroid_lat, - dims=["n_edge"], - attrs={ - "standard_name": "latitude", - "long name": "Latitude of the center of each edge", - "units": "degrees_north", - }, + dims=[ugrid.EDGE_DIM], + attrs=ugrid.EDGE_LAT_ATTRS, ) if "edge_x" not in grid._ds or repopulate: grid._ds["edge_x"] = xr.DataArray( centroid_x, - dims=["n_edge"], - attrs={ - "standard_name": "x", - "long name": "x coordinate of the center of each edge", - "units": "meters", - }, + dims=[ugrid.EDGE_DIM], + attrs=ugrid.EDGE_X_ATTRS, ) grid._ds["edge_y"] = xr.DataArray( centroid_y, - dims=["n_edge"], - attrs={ - "standard_name": "y", - "long name": "y coordinate of the center of each edge", - "units": "meters", - }, + dims=[ugrid.EDGE_DIM], + attrs=ugrid.EDGE_Y_ATTRS, ) grid._ds["edge_z"] = xr.DataArray( centroid_z, - dims=["n_edge"], - attrs={ - "standard_name": "z", - "long name": "z coordinate of the center of each edge", - "units": "meters", - }, + dims=[ugrid.EDGE_DIM], + attrs=ugrid.EDGE_Z_ATTRS, ) diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 4830d5a97..1d5b1a52d 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -2,15 +2,24 @@ import xarray as xr import numpy as np -from typing import Optional, Union +from typing import ( + Optional, + Union, +) + # reader and writer imports from uxarray.io._exodus import _read_exodus, _encode_exodus from uxarray.io._mpas import _read_mpas -from uxarray.io._ugrid import _read_ugrid, _encode_ugrid, _validate_minimum_ugrid +from uxarray.io._ugrid import ( + _read_ugrid, + _encode_ugrid, + _validate_minimum_ugrid, +) from uxarray.io._scrip import _read_scrip, _encode_scrip 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.utils import _parse_grid_type from uxarray.grid.area import get_all_face_area_from_coords @@ -110,10 +119,10 @@ def __init__( # check if inputted dataset is a minimum representable 2D UGRID unstructured grid if not _validate_minimum_ugrid(grid_ds): raise ValueError( - "Direct use of Grid constructor requires grid_ds to follow the internal unstructured grid definition, " - "including variable and dimension names. This grid_ds does not satisfy those requirements. If you are " - "not sure about how to do that, using ux.open_grid() or ux.from_dataset() is suggested." - ) # TODO: elaborate once we have a formal definition + "Grid unable to be represented in the UGRID conventions. Representing an unstructured grid requires " + "at least the following variables: ['node_lon'," + "'node_lat', and 'face_node_connectivity']" + ) # grid spec not provided, check if grid_ds is a minimum representable UGRID dataset if source_grid_spec is None: @@ -136,7 +145,6 @@ def __init__( # initialize attributes self._antimeridian_face_indices = None - self._face_areas = None # initialize cached data structures (visualization) self._gdf = None @@ -202,6 +210,61 @@ def from_dataset( return cls(grid_ds, source_grid_spec, source_dims_dict) + @classmethod + def from_topology( + cls, + node_lon: np.ndarray, + node_lat: np.ndarray, + face_node_connectivity: np.ndarray, + fill_value: Optional = None, + start_index: Optional[int] = 0, + dims_dict: Optional[dict] = None, + **kwargs, + ): + """Constructs a ``Grid`` object from user-defined topology variables + provided in the UGRID conventions. + + Note + ---- + To construct a UGRID-complient grid, the user must provide at least ``node_lon``, ``node_lat`` and ``face_node_connectivity`` + + Parameters + ---------- + node_lon : np.ndarray + Longitude of node coordinates + node_lat : np.ndarray + Latitude of node coordinates + face_node_connectivity : np.ndarray + Face node connectivity, mapping each face to the nodes that surround them + fill_value: Optional + Value used for padding connectivity variables when the maximum number of elements in a row is less than the maximum. + start_index: Optional, default=0 + Start index (typically 0 or 1) + dims_dict : Optional, dict + Dictionary of dimension names mapped to the ugrid conventions (i.e. {"nVertices": "n_node}) + **kwargs : + + Usage + ----- + >>> import uxarray as ux + >>> node_lon, node_lat, face_node_connectivity, fill_value = ... + >>> uxgrid = ux.Grid.from_ugrid(node_lon, node_lat, face_node_connectivity, fill_value) + """ + + if dims_dict is None: + dims_dict = {} + + grid_ds = _read_topology( + node_lon, + node_lat, + face_node_connectivity, + fill_value, + start_index, + **kwargs, + ) + grid_spec = "User Defined Topology" + return cls(grid_ds, grid_spec, dims_dict) + @classmethod def from_face_vertices( cls, @@ -265,88 +328,48 @@ def validate(self): def __repr__(self): """Constructs a string representation of the contents of a ``Grid``.""" + + from uxarray.conventions import ugrid, descriptors + prefix = "\n" original_grid_str = f"Original Grid Type: {self.source_grid_spec}\n" dims_heading = "Grid Dimensions:\n" dims_str = "" - for key, value in zip(self._ds.sizes.keys(), self._ds.sizes.values()): - dims_str += f" * {key}: {value}\n" + for dim_name in ugrid.DIM_NAMES: + if dim_name in self._ds.sizes: + dims_str += f" * {dim_name}: {self._ds.sizes[dim_name]}\n" dims_str += f" * n_nodes_per_face: {self.n_nodes_per_face.shape}\n" coord_heading = "Grid Coordinates (Spherical):\n" coords_str = "" - if "node_lon" in self._ds: - coords_str += f" * node_lon: {self.node_lon.shape}\n" - coords_str += f" * node_lat: {self.node_lat.shape}\n" - if "edge_lon" in self._ds: - coords_str += f" * edge_lon: {self.edge_lon.shape}\n" - coords_str += f" * edge_lat: {self.edge_lat.shape}\n" - if "face_lon" in self._ds: - coords_str += f" * face_lon: {self.face_lon.shape}\n" - coords_str += f" * face_lat: {self.face_lat.shape}\n" + for coord_name in ugrid.SPHERICAL_COORD_NAMES: + if coord_name in self._ds: + coords_str += f" * {coord_name}: {getattr(self, coord_name).shape}\n" coords_str += "Grid Coordinates (Cartesian):\n" - if "node_x" in self._ds: - coords_str += f" * node_x: {self.node_x.shape}\n" - coords_str += f" * node_y: {self.node_y.shape}\n" - coords_str += f" * node_z: {self.node_z.shape}\n" - if "edge_x" in self._ds: - coords_str += f" * edge_x: {self.edge_x.shape}\n" - coords_str += f" * edge_y: {self.edge_y.shape}\n" - coords_str += f" * edge_z: {self.edge_z.shape}\n" - if "face_x" in self._ds: - coords_str += f" * face_x: {self.face_x.shape}\n" - coords_str += f" * face_y: {self.face_y.shape}\n" - coords_str += f" * face_z: {self.face_z.shape}\n" + for coord_name in ugrid.CARTESIAN_COORD_NAMES: + if coord_name in self._ds: + coords_str += f" * {coord_name}: {getattr(self, coord_name).shape}\n" connectivity_heading = "Grid Connectivity Variables:\n" connectivity_str = "" - if "face_node_connectivity" in self._ds: - connectivity_str += ( - f" * face_node_connectivity: {self.face_node_connectivity.shape}\n" - ) - - if "edge_node_connectivity" in self._ds: - connectivity_str += ( - f" * edge_node_connectivity: {self.edge_node_connectivity.shape}\n" - ) - if "node_node_connectivity" in self._ds: - connectivity_str += ( - f" * node_node_connectivity: {self.node_node_connectivity.shape}\n" - ) - - if "face_edge_connectivity" in self._ds: - connectivity_str += ( - f" * face_edge_connectivity: {self.face_edge_connectivity.shape}\n" - ) - - if "edge_edge_connectivity" in self._ds: - connectivity_str += ( - f" * edge_edge_connectivity: {self.edge_edge_connectivity.shape}\n" - ) - - if "node_edge_connectivity" in self._ds: - connectivity_str += ( - f" * node_edge_connectivity: {self.node_edge_connectivity.shape}\n" - ) - - if "face_face_connectivity" in self._ds: - connectivity_str += ( - f" * face_face_connectivity: {self.face_face_connectivity.shape}\n" - ) + for conn_name in ugrid.CONNECTIVITY_NAMES: + if conn_name in self._ds: + connectivity_str += ( + f" * {conn_name}: {getattr(self, conn_name).shape}\n" + ) - if "edge_face_connectivity" in self._ds: - connectivity_str += ( - f" * edge_face_connectivity: {self.edge_face_connectivity.shape}\n" - ) + descriptors_heading = "Grid Descriptor Variables:\n" + descriptors_str = "" - if "node_face_connectivity" in self._ds: - connectivity_str += ( - f" * node_face_connectivity: {self.node_face_connectivity.shape}\n" - ) + for descriptor_name in descriptors.DESCRIPTOR_NAMES: + if descriptor_name in self._ds: + descriptors_str += ( + f" * {descriptor_name}: {getattr(self, descriptor_name).shape}\n" + ) return ( prefix @@ -357,8 +380,20 @@ def __repr__(self): + coords_str + connectivity_heading + connectivity_str + + descriptors_heading + + descriptors_str ) + def __getitem__(self, item): + """Implementation of getitem operator for indexing a grid to obtain + variables. + + Usage + ----- + >>> uxgrid['face_node_connectivity'] + """ + return getattr(self, item) + def __eq__(self, other) -> bool: """Two grids are equal if they have matching grid topology variables, coordinates, and dims all of which are equal. @@ -404,82 +439,118 @@ def __ne__(self, other) -> bool: """ return not self.__eq__(other) + @property + def dims(self) -> set: + """Names of all unstructured grid dimensions.""" + from uxarray.conventions.ugrid import DIM_NAMES + + return set([dim for dim in DIM_NAMES if dim in self._ds.dims]) + + @property + def sizes(self) -> dict: + """Names and values of all unstructured grid dimensions.""" + return {dim: self._ds.dims[dim] for dim in self.dims} + + @property + def coordinates(self) -> set: + """Names of all coordinate variables.""" + from uxarray.conventions.ugrid import ( + SPHERICAL_COORD_NAMES, + CARTESIAN_COORD_NAMES, + ) + + return set( + [coord for coord in SPHERICAL_COORD_NAMES if coord in self._ds] + ).union(set([coord for coord in CARTESIAN_COORD_NAMES if coord in self._ds])) + + @property + def connectivity(self) -> set: + """Names of all connectivity variables.""" + from uxarray.conventions.ugrid import CONNECTIVITY_NAMES + + return set([conn for conn in CONNECTIVITY_NAMES if conn in self._ds]) + @property def parsed_attrs(self) -> dict: """Dictionary of parsed attributes from the source grid.""" + warn( + "Grid.parsed_attrs will be deprecated in a future release. Please use Grid.attrs instead.", + DeprecationWarning, + ) return self._ds.attrs @property - def Mesh2(self) -> xr.DataArray: - """UGRID Attribute ``Mesh2``, which indicates the topology data of a 2D - unstructured mesh. - - Internal use has been deprecated. - """ - return self._ds["Mesh2"] + def attrs(self) -> dict: + """Dictionary of parsed attributes from the source grid.""" + return self._ds.attrs - # ================================================================================================================== - # Grid Dimensions @property def n_node(self) -> int: - """Dimension ``n_node``, which represents the total number of unique - corner nodes.""" + """Total number of nodes.""" return self._ds.sizes["n_node"] - @property - def n_face(self) -> int: - """Dimension ``n_face``, which represents the total number of unique - faces.""" - return self._ds.sizes["n_face"] - @property def n_edge(self) -> int: - """Dimension ``n_edge``, which represents the total number of unique - edges.""" + """Total number of edges.""" if "edge_node_connectivity" not in self._ds: _populate_edge_node_connectivity(self) return self._ds.sizes["n_edge"] - # ================================================================================================================== + @property + def n_face(self) -> int: + """Total number of faces.""" + return self._ds.sizes["n_face"] + @property def n_max_face_nodes(self) -> int: - """Dimension ``n_max_face_nodes``, which represents the maximum number - of nodes that a face may contain.""" + """The maximum number of nodes that can make up a single face.""" return self.face_node_connectivity.shape[1] @property - def n_max_face_edges(self) -> xr.DataArray: - """Dimension ``n_max_face_edges``, which represents the maximum number - of edges per face. + def n_max_face_edges(self) -> int: + """The maximum number of edges that surround a single face. Equivalent to ``n_max_face_nodes`` """ - if "n_max_face_edges" not in self._ds: - _populate_face_edge_connectivity(self) + return self.face_edge_connectivity.shape[1] + + @property + def n_max_face_faces(self) -> int: + """The maximum number of faces that surround a single face.""" + return self.face_face_connectivity.shape[1] + + @property + def n_max_edge_edges(self) -> int: + """The maximum number of edges that surround a single edge.""" + return self.edge_edge_connectivity.shape[1] + + @property + def n_max_node_faces(self) -> int: + """The maximum number of faces that surround a single node.""" + return self.node_face_connectivity.shape[1] - return self._ds["face_edge_connectivity"].shape[1] + @property + def n_max_node_edges(self) -> int: + """The maximum number of edges that surround a single node.""" + return self.node_edge_connectivity.shape[1] @property def n_nodes_per_face(self) -> xr.DataArray: - """Dimension Variable ``n_nodes_per_face``, which contains the number - of non-fill-value nodes per face. + """The number of nodes that make up each face. - Dimensions (``n_node``) and DataType ``INT_DTYPE``. + Dimensions: ``(n_node, )`` """ if "n_nodes_per_face" not in self._ds: _populate_n_nodes_per_face(self) return self._ds["n_nodes_per_face"] - # ================================================================================================================== - # Spherical Node Coordinates @property def node_lon(self) -> xr.DataArray: - """Coordinate ``node_lon``, which contains the longitude of each node - in degrees. + """Longitude of each node in degrees. - Dimensions (``n_node``) + Dimensions: ``(n_node, )`` """ if "node_lon" not in self._ds: _set_desired_longitude_range(self._ds) @@ -488,24 +559,20 @@ def node_lon(self) -> xr.DataArray: @property def node_lat(self) -> xr.DataArray: - """Coordinate ``node_lat``, which contains the latitude of each node in - degrees. + """Latitude of each node in degrees. - Dimensions (``n_node``) + Dimensions: ``(n_node, )`` """ if "node_lat" not in self._ds: _set_desired_longitude_range(self._ds) _populate_lonlat_coord(self) return self._ds["node_lat"] - # ================================================================================================================== - # Cartesian Node Coordinates @property def node_x(self) -> xr.DataArray: - """Coordinate ``node_x``, which contains the Cartesian x location of - each node in meters. + """Cartesian x location of each node in meters. - Dimensions (``n_node``) + Dimensions: ``(n_node, )`` """ if "node_x" not in self._ds: _populate_cartesian_xyz_coord(self) @@ -514,10 +581,9 @@ def node_x(self) -> xr.DataArray: @property def node_y(self) -> xr.DataArray: - """Coordinate ``node_y``, which contains the Cartesian y location of - each node in meters. + """Cartesian y location of each node in meters. - Dimensions (``n_node``) + Dimensions: ``(n_node, )`` """ if "node_y" not in self._ds: _populate_cartesian_xyz_coord(self) @@ -525,23 +591,19 @@ def node_y(self) -> xr.DataArray: @property def node_z(self) -> xr.DataArray: - """Coordinate ``node_z``, which contains the Cartesian y location of - each node in meters. + """Cartesian z location of each node in meters. - Dimensions (``n_node``) + Dimensions: ``(n_node, )`` """ if "node_z" not in self._ds: _populate_cartesian_xyz_coord(self) return self._ds["node_z"] - # ================================================================================================================== - # Spherical Edge Coordinates @property def edge_lon(self) -> xr.DataArray: - """Coordinate ``edge_lon``, which contains the longitude of each edge - in degrees. + """Longitude of the center of each edge in degrees. - Dimensions (``n_edge``) + Dimensions: ``(n_edge, )`` """ if "edge_lon" not in self._ds: _populate_edge_centroids(self) @@ -551,24 +613,20 @@ def edge_lon(self) -> xr.DataArray: @property def edge_lat(self) -> xr.DataArray: - """Coordinate ``edge_lat``, which contains the latitude of each edge in - degrees. + """Latitude of the center of each edge in degrees. - Dimensions (``n_edge``) + Dimensions: ``(n_edge, )`` """ if "edge_lat" not in self._ds: _populate_edge_centroids(self) _set_desired_longitude_range(self._ds) return self._ds["edge_lat"] - # ================================================================================================================== - # Cartesian Edge Coordinates @property def edge_x(self) -> xr.DataArray: - """Coordinate ``edge_x``, which contains the Cartesian x location of - each edge in meters. + """Cartesian x location of the center of each edge in meters. - Dimensions (``n_edge``) + Dimensions: ``(n_edge, )`` """ if "edge_x" not in self._ds: _populate_edge_centroids(self) @@ -577,10 +635,9 @@ def edge_x(self) -> xr.DataArray: @property def edge_y(self) -> xr.DataArray: - """Coordinate ``edge_y``, which contains the Cartesian y location of - each edge in meters. + """Cartesian y location of the center of each edge in meters. - Dimensions (``n_edge``) + Dimensions: ``(n_edge, )`` """ if "edge_y" not in self._ds: _populate_edge_centroids(self) @@ -588,23 +645,19 @@ def edge_y(self) -> xr.DataArray: @property def edge_z(self) -> xr.DataArray: - """Coordinate ``edge_z``, which contains the Cartesian z location of - each edge in meters. + """Cartesian z location of the center of each edge in meters. - Dimensions (``n_edge``) + Dimensions: ``(n_edge, )`` """ if "edge_z" not in self._ds: _populate_edge_centroids(self) return self._ds["edge_z"] - # ================================================================================================================== - # Spherical Face Coordinates @property def face_lon(self) -> xr.DataArray: - """Coordinate ``face_lon``, which contains the longitude of each face - in degrees. + """Longitude of the center of each face in degrees. - Dimensions (``n_face``) + Dimensions: ``(n_face, )`` """ if "face_lon" not in self._ds: _populate_face_centroids(self) @@ -613,10 +666,9 @@ def face_lon(self) -> xr.DataArray: @property def face_lat(self) -> xr.DataArray: - """Coordinate ``face_lat``, which contains the latitude of each face in - degrees. + """Latitude of the center of each face in degrees. - Dimensions (``n_face``) + Dimensions: ``(n_face, )`` """ if "face_lat" not in self._ds: _populate_face_centroids(self) @@ -624,14 +676,11 @@ def face_lat(self) -> xr.DataArray: return self._ds["face_lat"] - # ================================================================================================================== - # Cartesian Face Coordinates @property def face_x(self) -> xr.DataArray: - """Coordinate ``face_x``, which contains the Cartesian x location of - each face in meters. + """Cartesian x location of the center of each face in meters. - Dimensions (``n_face``) + Dimensions: ``(n_face, )`` """ if "face_x" not in self._ds: _populate_face_centroids(self) @@ -640,10 +689,9 @@ def face_x(self) -> xr.DataArray: @property def face_y(self) -> xr.DataArray: - """Coordinate ``face_y``, which contains the Cartesian y location of - each face in meters. + """Cartesian y location of the center of each face in meters. - Dimensions (``n_face``) + Dimensions: ``(n_face, )`` """ if "face_y" not in self._ds: _populate_face_centroids(self) @@ -651,24 +699,19 @@ def face_y(self) -> xr.DataArray: @property def face_z(self) -> xr.DataArray: - """Coordinate ``face_z``, which contains the Cartesian z location of - each face in meters. + """Cartesian z location of the center of each face in meters. - Dimensions (``n_face``) + Dimensions: ``(n_face, )`` """ if "face_z" not in self._ds: _populate_face_centroids(self) return self._ds["face_z"] - # ================================================================================================================== - # (, node) Connectivity @property def face_node_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``face_node_connectivity``, which maps each - face to its corner nodes. + """Indices of the nodes that make up each face. - Dimensions (``n_face``, ``n_max_face_nodes``) and - DataType ``INT_DTYPE``. + Dimensions: ``(n_face, n_max_face_nodes)`` Nodes are in counter-clockwise order. """ @@ -688,11 +731,9 @@ def face_node_connectivity(self) -> xr.DataArray: @property def edge_node_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``edge_node_connectivity``, which maps every - edge to the two nodes that it connects. + """Indices of the two nodes that make up each edge. - Dimensions (``n_edge``, ``two``) and DataType - ``INT_DTYPE``. + Dimensions: ``(n_edge, n_max_edge_nodes)`` Nodes are in arbitrary order. """ @@ -703,18 +744,18 @@ def edge_node_connectivity(self) -> xr.DataArray: @property def node_node_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``node_node_connectivity``.""" - return None + """Indices of the nodes that surround each node.""" + if "node_node_connectivity" not in self._ds: + raise NotImplementedError( + "Construction of `node_node_connectivity` not yet supported." + ) + return self._ds["node_node_connectivity"] - # ================================================================================================================== - # (, edge) Connectivity @property def face_edge_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``face_edge_connectivity``, which maps every - face to its edges. + """Indices of the edges that surround each face. - Dimensions (``n_face``, ``n_max_face_nodes``) and DataType - ``INT_DTYPE``. + Dimensions: ``(n_face, n_max_face_edges)`` """ if "face_edge_connectivity" not in self._ds: _populate_face_edge_connectivity(self) @@ -723,27 +764,45 @@ def face_edge_connectivity(self) -> xr.DataArray: @property def edge_edge_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``edge_edge_connectivity``.""" - return None + """Indices of the edges that surround each edge. + + Dimensions: ``(n_face, n_max_edge_edges)`` + """ + if "edge_edge_connectivity" not in self._ds: + raise NotImplementedError( + "Construction of `edge_edge_connectivity` not yet supported." + ) + + return self._ds["edge_edge_connectivity"] @property def node_edge_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``node_edge_connectivity``.""" - return None + """Indices of the edges that surround each node.""" + if "node_edge_connectivity" not in self._ds: + raise NotImplementedError( + "Construction of `node_edge_connectivity` not yet supported." + ) + + return self._ds["node_edge_connectivity"] - # ================================================================================================================== - # (, face) Connectivity @property def face_face_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``face_face_connectivity``.""" - return None + """Indices of the faces that surround each face. + + Dimensions ``(n_face, n_max_face_faces)`` + """ + if "face_face_connectivity" not in self._ds: + raise NotImplementedError( + "Construction of `face_face_connectivity` not yet supported." + ) + + return self._ds["face_face_connectivity"] @property def edge_face_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``edge_face_connectivity``, which contains the - index of the faces that saddle a given edge. + """Indices of the faces that saddle each edge. - Dimensions (``n_edge``, ``TWO``) and DataType ``INT_DTYPE``. + Dimensions ``(n_edge, two)`` """ if "edge_face_connectivity" not in self._ds: _populate_edge_face_connectivity(self) @@ -752,24 +811,20 @@ def edge_face_connectivity(self) -> xr.DataArray: @property def node_face_connectivity(self) -> xr.DataArray: - """Connectivity Variable ``node_face_connectivity``, which maps every - node to its faces. + """Indices of the faces that surround each node. - Dimensions (``n_node``, ``n_max_faces_per_node``) and DataType - ``INT_DTYPE``. + Dimensions ``(n_node, n_max_node_faces)`` """ if "node_face_connectivity" not in self._ds: _populate_node_face_connectivity(self) return self._ds["node_face_connectivity"] - # ================================================================================================================== - # Distance Quantities @property def edge_node_distances(self): - """Contains the distance between the nodes that saddle a given edge. + """Distances between the two nodes that surround each edge. - Dimensions (``n_edge``) and DataType float. + Dimensions ``(n_edge, )`` """ if "edge_node_distances" not in self._ds: _populate_edge_node_distances(self) @@ -777,16 +832,14 @@ def edge_node_distances(self): @property def edge_face_distances(self): - """Contains the distance between the faces that saddle a given edge. + """Distances between the centers of the faces that saddle each edge. - Dimensions (``n_edge``) and DataType float. + Dimensions ``(n_edge, )`` """ if "edge_face_distances" not in self._ds: _populate_edge_face_distances(self) return self._ds["edge_face_distances"] - # ================================================================================================================== - # Other Grid Descriptor Quantities @property def antimeridian_face_indices(self) -> np.ndarray: """Index of each face that crosses the antimeridian.""" @@ -795,21 +848,22 @@ def antimeridian_face_indices(self) -> np.ndarray: return self._antimeridian_face_indices @property - def face_areas(self) -> np.ndarray: - """Declare face_areas as a property.""" - # if self._face_areas is not None: it allows for using the cached result - if self._face_areas is None: - self._face_areas, self._face_jacobian = self.compute_face_areas() - return self._face_areas - - # ================================================================================================================== + def face_areas(self) -> xr.DataArray: + """The area of each face.""" + from uxarray.conventions.descriptors import FACE_AREAS_DIMS, FACE_AREAS_ATTRS + + if "face_areas" not in self._ds: + face_areas, self._face_jacobian = self.compute_face_areas() + self._ds["face_areas"] = xr.DataArray( + data=face_areas, dims=FACE_AREAS_DIMS, attrs=FACE_AREAS_ATTRS + ) + return self._ds["face_areas"] @property def face_jacobian(self): """Declare face_jacobian as a property.""" - # if self._face_jacobian is not None: it allows for using the cached result if self._face_jacobian is None: - self._face_areas, self._face_jacobian = self.compute_face_areas() + _ = self.face_areas return self._face_jacobian def get_ball_tree( @@ -937,6 +991,10 @@ def encode_as(self, grid_type: str) -> xr.Dataset: If provided grid type or file type is unsupported. """ + warn( + "Grid.encode_as will be deprecated in a future release. Please use Grid.to_xarray instead." + ) + if grid_type == "UGRID": out_ds = _encode_ugrid(self._ds) @@ -1062,6 +1120,43 @@ def compute_face_areas( return self._face_areas, self._face_jacobian + def to_xarray(self, grid_format: Optional[str] = "ugrid"): + """Returns a xarray Dataset representation in a specific grid format + from the Grid object. + + Parameters + ---------- + grid_format: str, optional + The desired grid format for the output dataset. + One of "ugrid", "exodus", or "scrip" + + Returns + ------- + out_ds: xarray.Dataset + Dataset representing the unstructured grid in a given grid format + """ + + if grid_format == "ugrid": + out_ds = _encode_ugrid(self._ds) + + elif grid_format == "exodus": + out_ds = _encode_exodus(self._ds) + + elif grid_format == "scrip": + out_ds = _encode_scrip( + self.face_node_connectivity, + self.node_lon, + self.node_lat, + self.face_areas, + ) + + else: + raise ValueError( + f"Invalid grid_format encountered. Expected one of ['ugrid', 'exodus', 'scrip'] but received: {grid_format}" + ) + + return out_ds + def to_geodataframe( self, override: Optional[bool] = False, diff --git a/uxarray/io/_esmf.py b/uxarray/io/_esmf.py index e21be9e7f..f16abc657 100644 --- a/uxarray/io/_esmf.py +++ b/uxarray/io/_esmf.py @@ -3,6 +3,8 @@ from uxarray.constants import INT_DTYPE, INT_FILL_VALUE +from uxarray.conventions import ugrid + def _read_esmf(in_ds): """Reads in an Xarray dataset containing an ESMF formatted Grid dataset and @@ -37,60 +39,33 @@ def _read_esmf(in_ds): out_ds = xr.Dataset() source_dims_dict = { - "nodeCount": "n_node", - "elementCount": "n_face", - "maxNodePElement": "n_max_face_nodes", + "nodeCount": ugrid.NODE_DIM, + "elementCount": ugrid.FACE_DIM, + "maxNodePElement": ugrid.N_MAX_FACE_NODES_DIM, } if in_ds["nodeCoords"].units == "degrees": # Spherical Coordinates (in degrees) node_lon = in_ds["nodeCoords"].isel(coordDim=0).values - node_lat = in_ds["nodeCoords"].isel(coordDim=1).values - - out_ds["node_lon"] = xr.DataArray( - node_lon, - dims=["n_node"], - attrs={ - "standard_name": "longitude", - "long_name": "longitude of mesh nodes", - "units": "degrees_east", - }, + out_ds[ugrid.NODE_COORDINATES[0]] = xr.DataArray( + node_lon, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LON_ATTRS ) - out_ds["node_lat"] = xr.DataArray( - node_lat, - dims=["n_node"], - attrs={ - "standard_name": "latitude", - "long_name": "latitude of mesh nodes", - "units": "degrees_north", - }, + node_lat = in_ds["nodeCoords"].isel(coordDim=1).values + out_ds[ugrid.NODE_COORDINATES[1]] = xr.DataArray( + node_lat, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LAT_ATTRS ) if "centerCoords" in in_ds: - # parse center coords (face centers) if avaliable - + # parse center coords (face centers) if available face_lon = in_ds["centerCoords"].isel(coordDim=0).values - face_lat = in_ds["centerCoords"].isel(coordDim=1).values - - out_ds["face_lon"] = xr.DataArray( - face_lon, - dims=["n_face"], - attrs={ - "standard_name": "longitude", - "long_name": "longitude of center nodes", - "units": "degrees_east", - }, + out_ds[ugrid.FACE_COORDINATES[0]] = xr.DataArray( + face_lon, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_LON_ATTRS ) - out_ds["face_lat"] = xr.DataArray( - face_lat, - dims=["n_face"], - attrs={ - "standard_name": "latitude", - "long_name": "latitude of center nodes", - "units": "degrees_north", - }, + face_lat = in_ds["centerCoords"].isel(coordDim=1).values + out_ds[ugrid.FACE_COORDINATES[1]] = xr.DataArray( + face_lat, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_LAT_ATTRS ) else: @@ -99,11 +74,10 @@ def _read_esmf(in_ds): ) n_nodes_per_face = in_ds["numElementConn"].values.astype(INT_DTYPE) - out_ds["n_nodes_per_face"] = xr.DataArray( data=n_nodes_per_face, - dims=["n_face"], - attrs={"long_name": "number of non-fill value nodes for each face"}, + dims=ugrid.N_NODES_PER_FACE_DIMS, + attrs=ugrid.N_NODES_PER_FACE_ATTRS, ) if "start_index" in in_ds["elementConn"]: @@ -121,12 +95,8 @@ def _read_esmf(in_ds): out_ds["face_node_connectivity"] = xr.DataArray( data=face_node_connectivity, - dims=["n_face", "n_max_face_nodes"], - attrs={ - "cf_role": "face_node_connectivity", - "_FillValue": INT_FILL_VALUE, - "start_index": INT_DTYPE(0), - }, + dims=ugrid.FACE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_NODE_CONNECTIVITY_ATTRS, ) return out_ds, source_dims_dict diff --git a/uxarray/io/_exodus.py b/uxarray/io/_exodus.py index acbcfa8c7..205b68d36 100644 --- a/uxarray/io/_exodus.py +++ b/uxarray/io/_exodus.py @@ -7,6 +7,8 @@ from uxarray.constants import INT_DTYPE, INT_FILL_VALUE from uxarray.grid.coordinates import _get_lonlat_from_xyz, _get_xyz_from_lonlat +from uxarray.conventions import ugrid + # Exodus Number is one-based. def _read_exodus(ext_ds): @@ -40,63 +42,29 @@ def _read_exodus(ext_ds): pass elif key == "coord": ds["node_x"] = xr.DataArray( - data=ext_ds.coord[0], - dims=["n_node"], - attrs={ - "standard_name": "x", - "long_name": "cartesian x", - "units": "m", - }, + data=ext_ds.coord[0], dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_X_ATTRS ) ds["node_y"] = xr.DataArray( - data=ext_ds.coord[1], - dims=["n_node"], - attrs={ - "standard_name": "y", - "long_name": "cartesian y", - "units": "m", - }, + data=ext_ds.coord[1], dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_Y_ATTRS ) if ext_ds.sizes["num_dim"] > 2: ds["node_z"] = xr.DataArray( data=ext_ds.coord[2], - dims=["n_node"], - attrs={ - "standard_name": "z", - "long_name": "cartesian z", - "units": "m", - }, + dims=[ugrid.NODE_DIM], + attrs=ugrid.NODE_Z_ATTRS, ) elif key == "coordx": ds["node_x"] = xr.DataArray( - data=ext_ds.coordx, - dims=["n_node"], - attrs={ - "standard_name": "x", - "long_name": "cartesian x", - "units": "m", - }, + data=ext_ds.coordx, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_X_ATTRS ) elif key == "coordy": ds["node_y"] = xr.DataArray( - data=ext_ds.coordx, - dims=["n_node"], - attrs={ - "standard_name": "y", - "long_name": "cartesian y", - "units": "m", - }, + data=ext_ds.coordx, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_Y_ATTRS ) elif key == "coordz": if ext_ds.sizes["num_dim"] > 2: ds["node_z"] = xr.DataArray( - data=ext_ds.coordx, - dims=["n_node"], - attrs={ - "standard_name": "z", - "long_name": "cartesian z", - "units": "m", - }, + data=ext_ds.coordx, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_Z_ATTRS ) elif "connect" in key: # check if num face nodes is less than max. @@ -127,14 +95,8 @@ def _read_exodus(ext_ds): ds["face_node_connectivity"] = xr.DataArray( data=face_nodes, - dims=["n_face", "n_max_face_nodes"], - attrs={ - "cf_role": "face_node_connectivity", - "_FillValue": INT_FILL_VALUE, - "start_index": INT_DTYPE( - 0 - ), # NOTE: This might cause an error if numbering has holes - }, + dims=ugrid.FACE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_NODE_CONNECTIVITY_ATTRS, ) # populate lon/lat coordinates @@ -144,22 +106,10 @@ def _read_exodus(ext_ds): # populate dataset ds["node_lon"] = xr.DataArray( - data=lon, - dims=["n_node"], - attrs={ - "standard_name": "longitude", - "long_name": "longitude of mesh nodes", - "units": "degrees_east", - }, + data=lon, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LON_ATTRS ) ds["node_lat"] = xr.DataArray( - data=lat, - dims=["n_node"], - attrs={ - "standard_name": "latitude", - "long_name": "latitude of mesh nodes", - "units": "degrees_north", - }, + data=lat, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LAT_ATTRS ) # set lon/lat coordinates diff --git a/uxarray/io/_mpas.py b/uxarray/io/_mpas.py index dc038c3c8..c175a9a1a 100644 --- a/uxarray/io/_mpas.py +++ b/uxarray/io/_mpas.py @@ -2,6 +2,7 @@ import numpy as np from uxarray.constants import INT_DTYPE, INT_FILL_VALUE +from uxarray.conventions import ugrid, descriptors def _primal_to_ugrid(in_ds, out_ds): @@ -59,9 +60,9 @@ def _primal_to_ugrid(in_ds, out_ds): _parse_global_attrs(in_ds, out_ds) # populate source dims - source_dims_dict["nVertices"] = "n_node" - source_dims_dict[in_ds["verticesOnCell"].dims[0]] = "n_face" - source_dims_dict[in_ds["verticesOnCell"].dims[1]] = "n_max_face_nodes" + source_dims_dict["nVertices"] = ugrid.NODE_DIM + source_dims_dict[in_ds["verticesOnCell"].dims[0]] = ugrid.FACE_DIM + source_dims_dict[in_ds["verticesOnCell"].dims[1]] = ugrid.N_MAX_FACE_NODES_DIM return source_dims_dict @@ -121,9 +122,9 @@ def _dual_to_ugrid(in_ds, out_ds): _parse_global_attrs(in_ds, out_ds) # populate source dims - source_dims_dict[in_ds["latCell"].dims[0]] = "n_node" - source_dims_dict[in_ds["cellsOnVertex"].dims[0]] = "n_face" - source_dims_dict[in_ds["cellsOnVertex"].dims[1]] = "n_max_face_nodes" + source_dims_dict[in_ds["latCell"].dims[0]] = ugrid.NODE_DIM + source_dims_dict[in_ds["cellsOnVertex"].dims[0]] = ugrid.FACE_DIM + source_dims_dict[in_ds["cellsOnVertex"].dims[1]] = ugrid.N_MAX_FACE_NODES_DIM return source_dims_dict @@ -139,23 +140,11 @@ def _parse_node_latlon_coords(in_ds, out_ds, mesh_type): node_lat = np.rad2deg(in_ds["latCell"].values) out_ds["node_lon"] = xr.DataArray( - node_lon, - dims=["n_node"], - attrs={ - "standard_name": "longitude", - "long_name": "longitude of mesh nodes", - "units": "degrees_east", - }, + node_lon, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LON_ATTRS ) out_ds["node_lat"] = xr.DataArray( - node_lat, - dims=["n_node"], - attrs={ - "standard_name": "latitude", - "long_name": "latitude of mesh nodes", - "units": "degrees_north", - }, + node_lat, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LAT_ATTRS ) @@ -173,33 +162,15 @@ def _parse_node_xyz_coords(in_ds, out_ds, mesh_type): node_z = in_ds["zCell"].values out_ds["node_x"] = xr.DataArray( - data=node_x, - dims=["n_node"], - attrs={ - "standard_name": "x", - "long_name": "cartesian node x", - "units": "m", - }, + data=node_x, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_X_ATTRS ) out_ds["node_y"] = xr.DataArray( - data=node_y, - dims=["n_node"], - attrs={ - "standard_name": "y", - "long_name": "cartesian node y", - "units": "m", - }, + data=node_y, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_Y_ATTRS ) out_ds["node_z"] = xr.DataArray( - data=node_z, - dims=["n_node"], - attrs={ - "standard_name": "z", - "long_name": "cartesian node z", - "units": "m", - }, + data=node_z, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_Z_ATTRS ) @@ -214,23 +185,11 @@ def _parse_face_latlon_coords(in_ds, out_ds, mesh_type): face_lat = np.rad2deg(in_ds["latVertex"].values) out_ds["face_lon"] = xr.DataArray( - face_lon, - dims=["n_face"], - attrs={ - "standard_name": "longitude", - "long_name": "longitude of center nodes", - "units": "degrees_east", - }, + face_lon, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_LON_ATTRS ) out_ds["face_lat"] = xr.DataArray( - face_lat, - dims=["n_face"], - attrs={ - "standard_name": "latitude", - "long_name": "latitude of center nodes", - "units": "degrees_north", - }, + face_lat, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_LAT_ATTRS ) @@ -248,33 +207,15 @@ def _parse_face_xyz_coords(in_ds, out_ds, mesh_type): face_z = in_ds["zVertex"].values out_ds["face_x"] = xr.DataArray( - data=face_x, - dims=["n_face"], - attrs={ - "standard_name": "x", - "long_name": "cartesian edge x", - "units": "m", - }, + data=face_x, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_X_ATTRS ) out_ds["face_y"] = xr.DataArray( - data=face_y, - dims=["n_face"], - attrs={ - "standard_name": "y", - "long_name": "cartesian edge y", - "units": "m", - }, + data=face_y, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_Y_ATTRS ) out_ds["face_z"] = xr.DataArray( - data=face_z, - dims=["n_face"], - attrs={ - "standard_name": "z", - "long_name": "cartesian edge z", - "units": "m", - }, + data=face_z, dims=[ugrid.FACE_DIM], attrs=ugrid.FACE_Z_ATTRS ) @@ -286,23 +227,11 @@ def _parse_edge_latlon_coords(in_ds, out_ds, mesh_type): edge_lat = np.rad2deg(in_ds["latEdge"].values) out_ds["edge_lon"] = xr.DataArray( - edge_lon, - dims=["n_edge"], - attrs={ - "standard_name": "longitude", - "long_name": "longitude of edge centers", - "units": "degrees_east", - }, + edge_lon, dims=[ugrid.EDGE_DIM], attrs=ugrid.EDGE_LON_ATTRS ) out_ds["edge_lat"] = xr.DataArray( - edge_lat, - dims=["n_edge"], - attrs={ - "standard_name": "latitude", - "long_name": "latitude of edge centers", - "units": "degrees_north", - }, + edge_lat, dims=[ugrid.EDGE_DIM], attrs=ugrid.EDGE_LAT_ATTRS ) @@ -314,33 +243,15 @@ def _parse_edge_xyz_coords(in_ds, out_ds, mesh_type): edge_z = in_ds["zEdge"].values out_ds["edge_x"] = xr.DataArray( - data=edge_x, - dims=["n_edge"], - attrs={ - "standard_name": "x", - "long_name": "cartesian edge x", - "units": "m", - }, + data=edge_x, dims=[ugrid.EDGE_DIM], attrs=ugrid.EDGE_X_ATTRS ) out_ds["edge_y"] = xr.DataArray( - data=edge_y, - dims=["n_edge"], - attrs={ - "standard_name": "y", - "long_name": "cartesian edge y", - "units": "m", - }, + data=edge_y, dims=[ugrid.EDGE_DIM], attrs=ugrid.EDGE_Y_ATTRS ) out_ds["edge_z"] = xr.DataArray( - data=edge_z, - dims=["n_edge"], - attrs={ - "standard_name": "z", - "long_name": "cartesian edge z", - "units": "m", - }, + data=edge_z, dims=[ugrid.EDGE_DIM], attrs=ugrid.EDGE_Z_ATTRS ) @@ -373,18 +284,12 @@ def _parse_face_nodes(in_ds, out_ds, mesh_type): face_nodes = cellsOnVertex - face_nodes = xr.DataArray( + out_ds["face_node_connectivity"] = xr.DataArray( data=face_nodes, - dims=["n_face", "n_max_face_nodes"], - attrs={ - "cf_role": "face_node_connectivity", - "_FillValue": INT_FILL_VALUE, - "start_index": INT_DTYPE(0), - }, + dims=ugrid.FACE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_NODE_CONNECTIVITY_ATTRS, ) - out_ds["face_node_connectivity"] = face_nodes - def _parse_edge_nodes(in_ds, out_ds, mesh_type): """Parses edge node connectivity for either the Primal or Dual Mesh.""" @@ -413,8 +318,8 @@ def _parse_edge_nodes(in_ds, out_ds, mesh_type): out_ds["edge_node_connectivity"] = xr.DataArray( data=edge_nodes, - dims=["n_edge", "two"], - attrs={"cf_role": "edge_node_connectivity", "start_index": INT_DTYPE(0)}, + dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.EDGE_NODE_CONNECTIVITY_ATTRS, ) @@ -448,8 +353,8 @@ def _parse_node_faces(in_ds, out_ds, mesh_type): out_ds["node_face_connectivity"] = xr.DataArray( data=node_faces, - dims=["n_node", "n_max_faces_per_node"], - attrs={"cf_role": "node_face_connectivity", "start_index": INT_DTYPE(0)}, + dims=ugrid.NODE_FACE_CONNECTIVITY_DIMS, + attrs=ugrid.NODE_FACE_CONNECTIVITY_ATTRS, ) @@ -484,12 +389,8 @@ def _parse_face_edges(in_ds, out_ds, mesh_type): out_ds["face_edge_connectivity"] = xr.DataArray( data=face_edges, - dims=["n_face", "n_max_face_nodes"], - attrs={ - "cf_role": "face_node_connectivity", - "_FillValue": INT_FILL_VALUE, - "start_index": INT_DTYPE(0), - }, + dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, ) @@ -521,8 +422,8 @@ def _parse_edge_faces(in_ds, out_ds, mesh_type): out_ds["edge_face_connectivity"] = xr.DataArray( data=edge_faces, - dims=["n_edge", "two"], - attrs={"cf_role": "edge_face_connectivity", "start_index": INT_DTYPE(0)}, + dims=ugrid.EDGE_FACE_CONNECTIVITY_DIMS, + attrs=ugrid.EDGE_FACE_CONNECTIVITY_ATTRS, ) @@ -531,7 +432,9 @@ def _parse_edge_node_distances(in_ds, out_ds): edge_node_distances = in_ds["dvEdge"].values out_ds["edge_node_distances"] = xr.DataArray( - data=edge_node_distances, dims=["n_edge"], attrs={"start_index": INT_DTYPE(0)} + data=edge_node_distances, + dims=descriptors.EDGE_NODE_DISTANCES_DIMS, + attrs=descriptors.EDGE_NODE_DISTANCES_ATTRS, ) @@ -540,7 +443,9 @@ def _parse_edge_face_distances(in_ds, out_ds): edge_face_distances = in_ds["dcEdge"].values out_ds["edge_face_distances"] = xr.DataArray( - data=edge_face_distances, dims=["n_edge"], attrs={"start_index": INT_DTYPE(0)} + data=edge_face_distances, + dims=descriptors.EDGE_FACE_DISTANCES_DIMS, + attrs=descriptors.EDGE_FACE_DISTANCES_ATTRS, ) diff --git a/uxarray/io/_scrip.py b/uxarray/io/_scrip.py index 5fe312499..c07824147 100644 --- a/uxarray/io/_scrip.py +++ b/uxarray/io/_scrip.py @@ -4,6 +4,8 @@ from uxarray.grid.connectivity import _replace_fill_values from uxarray.constants import INT_DTYPE, INT_FILL_VALUE +from uxarray.conventions import ugrid + def _to_ugrid(in_ds, out_ds): """If input dataset (``in_ds``) file is an unstructured SCRIP file, @@ -46,29 +48,26 @@ def _to_ugrid(in_ds, out_ds): unq_inv = np.reshape(unq_inv, (len(in_ds.grid_size), len(in_ds.grid_corners))) # Create node_lon & node_lat from unsorted, unique grid_corner_lat/lon - out_ds["node_lon"] = xr.DataArray( - unq_lon, - dims=["n_node"], - attrs={ - "standard_name": "longitude", - "long_name": "longitude of mesh nodes", - "units": "degrees_east", - }, + out_ds[ugrid.NODE_COORDINATES[0]] = xr.DataArray( + unq_lon, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LON_ATTRS ) - out_ds["node_lat"] = xr.DataArray( - unq_lat, - dims=["n_node"], - attrs={ - "standard_name": "latitude", - "long_name": "latitude of mesh nodes", - "units": "degrees_north", - }, + 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["face_lon"] = in_ds["grid_center_lon"] - out_ds["face_lat"] = in_ds["grid_center_lat"] + 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( @@ -78,14 +77,8 @@ def _to_ugrid(in_ds, out_ds): # set the face nodes data compiled in "connect" section out_ds["face_node_connectivity"] = xr.DataArray( data=face_nodes, - dims=["n_face", "n_max_face_nodes"], - attrs={ - "cf_role": "face_node_connectivity", - "_FillValue": INT_FILL_VALUE, - "start_index": INT_DTYPE( - 0 - ), # NOTE: This might cause an error if numbering has holes - }, + dims=ugrid.FACE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_NODE_CONNECTIVITY_ATTRS, ) else: diff --git a/uxarray/io/_topology.py b/uxarray/io/_topology.py new file mode 100644 index 000000000..8335796a5 --- /dev/null +++ b/uxarray/io/_topology.py @@ -0,0 +1,66 @@ +import xarray as xr + +import uxarray.conventions.ugrid as ugrid +from uxarray.grid.connectivity import _replace_fill_values +from uxarray.constants import INT_FILL_VALUE, INT_DTYPE + + +def _read_topology( + node_lon, node_lat, face_node_connectivity, fill_value, start_index, **kwargs +): + ds = xr.Dataset() + + for coord in ugrid.SPHERICAL_COORD_NAMES: + # parse spherical coordinates + if coord in ["node_lon", "node_lat"] or coord in kwargs: + if coord == "node_lon": + coord_arr = node_lon + elif coord == "node_lat": + coord_arr = node_lat + else: + coord_arr = kwargs[coord] + + ds[coord] = xr.DataArray( + data=coord_arr, + dims=ugrid.SPHERICAL_COORDS[coord]["dims"], + attrs=ugrid.SPHERICAL_COORDS[coord]["attrs"], + ) + + for coord in ugrid.CARTESIAN_COORD_NAMES: + # parse cartesian coords + if coord in kwargs: + coord_arr = kwargs[coord] + + ds[coord] = xr.DataArray( + data=coord_arr, + dims=ugrid.CARTESIAN_COORDS[coord]["dims"], + attrs=ugrid.CARTESIAN_COORDS[coord]["attrs"], + ) + + for conn in ugrid.CONNECTIVITY_NAMES: + # parse connectivity + if conn == "face_node_connectivity" or conn in kwargs: + if conn == "face_node_connectivity": + conn_arr = face_node_connectivity + else: + conn_arr = kwargs[conn] + + ds[conn] = xr.DataArray( + data=_process_connectivity(conn_arr, fill_value, start_index), + dims=ugrid.CONNECTIVITY[conn]["dims"], + attrs=ugrid.CONNECTIVITY[conn]["attrs"], + ) + + return ds + + +def _process_connectivity(conn, orig_fv, start_index): + """Internal helper for processing connectivity variables, standardizing + fill values and converting to zero-index.""" + + if orig_fv != INT_FILL_VALUE: + conn = _replace_fill_values(conn, orig_fv, INT_FILL_VALUE, INT_DTYPE) + + conn[conn != INT_FILL_VALUE] -= start_index + + return conn diff --git a/uxarray/io/_ugrid.py b/uxarray/io/_ugrid.py index 057800365..7fb87544f 100644 --- a/uxarray/io/_ugrid.py +++ b/uxarray/io/_ugrid.py @@ -1,76 +1,118 @@ import numpy as np - +import xarray as xr from uxarray.grid.connectivity import _replace_fill_values from uxarray.constants import INT_DTYPE, INT_FILL_VALUE -def _read_ugrid(xr_ds): - """UGRID file reader. +import uxarray.conventions.ugrid as ugrid - Parameters: xarray.Dataset, required - Returns: ugrid aware xarray.Dataset - """ - source_dims_dict = {} - # TODO: obtain and change to Mesh2 construct, see Issue #27 +def _read_ugrid(ds): + """Parses an unstructured grid dataset and encodes it in the UGRID + conventions.""" - # get the data variable name that has attribute "cf_role" set to "mesh_topology" - # this is the base xarray.DataArray name - base_xarray_var = list(xr_ds.filter_by_attrs(cf_role="mesh_topology").keys())[0] - # TODO: Allow for parsing datasets with more than just coordinates and face nodes + # Grid Topology + grid_topology_name = list(ds.filter_by_attrs(cf_role="mesh_topology").keys())[0] + ds = ds.rename({grid_topology_name: "grid_topology"}) - xr_ds = xr_ds.rename({base_xarray_var: "Mesh2"}) + # Coordinates - # map and rename coordinates - coord_names = xr_ds["Mesh2"].node_coordinates.split() - if len(coord_names) == 1: - xr_ds = xr_ds.rename({coord_names[0]: "node_lon"}) - elif len(coord_names) == 2: - xr_ds = xr_ds.rename({coord_names[0]: "node_lon", coord_names[1]: "node_lat"}) - # map and rename dimensions - coord_dim_name = xr_ds["node_lon"].dims + # get the names of node_lon and node_lat + node_lon_name, node_lat_name = ds["grid_topology"].node_coordinates.split() + coord_dict = { + node_lon_name: ugrid.NODE_COORDINATES[0], + node_lat_name: ugrid.NODE_COORDINATES[1], + } - xr_ds = xr_ds.rename({coord_dim_name[0]: "n_node"}) + if "edge_coordinates" in ds["grid_topology"]: + # 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] - face_node_names = xr_ds["Mesh2"].face_node_connectivity.split() + if "face_coordinates" in ds["grid_topology"]: + # get the names of face_lon and face_lat, if they exist + face_lon_name, face_lat_name = ds["grid_topology"].edge_coordinates.split() + coord_dict[face_lon_name] = ugrid.FACE_COORDINATES[0] + coord_dict[face_lat_name] = ugrid.FACE_COORDINATES[1] - face_node_name = face_node_names[0] - xr_ds = xr_ds.rename({xr_ds[face_node_name].name: "face_node_connectivity"}) + ds = ds.rename(coord_dict) + # Connectivity - xr_ds = xr_ds.rename( - { - xr_ds["face_node_connectivity"].dims[0]: "n_face", - xr_ds["face_node_connectivity"].dims[1]: "n_max_face_nodes", - } - ) + conn_dict = {} + for conn_name in ugrid.CONNECTIVITY_NAMES: + if conn_name in ds.grid_topology.attrs: + orig_conn_name = ds.grid_topology.attrs[conn_name] + conn_dict[orig_conn_name] = conn_name + elif len(ds.filter_by_attrs(cf_role=conn_name).keys()): + orig_conn_name = list(ds.filter_by_attrs(cf_role=conn_name).keys())[0] + conn_dict[orig_conn_name] = conn_name - xr_ds = xr_ds.set_coords(["node_lon", "node_lat"]) + ds = ds.rename(conn_dict) - # standardize fill values and data type for face nodes - xr_ds = _standardize_fill_values(xr_ds) + for conn_name in conn_dict.values(): + ds = _standardize_connectivity(ds, conn_name) - # populate source dimensions - source_dims_dict[coord_dim_name[0]] = "n_node" - source_dims_dict[xr_ds["face_node_connectivity"].dims[0]] = "n_face" - source_dims_dict[xr_ds["face_node_connectivity"].dims[1]] = "n_max_face_nodes" + dim_dict = {} - return xr_ds, source_dims_dict + # Rename Core Dims (node, edge, face) + if "node_dimension" in ds["grid_topology"]: + dim_dict[ds["grid_topology"].node_dimension] = ugrid.NODE_DIM + else: + dim_dict[ds["node_lon"].dims[0]] = ugrid.NODE_DIM + + if "face_dimension" in ds["grid_topology"]: + dim_dict[ds["grid_topology"].face_dimension] = ugrid.FACE_DIM + else: + dim_dict[ds["face_node_connectivity"].dims[0]] = ugrid.FACE_DIM + + if "edge_dimension" in ds["grid_topology"]: + # edge dimension is not always provided + dim_dict[ds["grid_topology"].edge_dimension] = ugrid.EDGE_DIM + else: + if "edge_lon" in ds: + dim_dict[ds["edge_lon"].dims[0]] = ugrid.EDGE_DIM + + dim_dict[ds["face_node_connectivity"].dims[1]] = ugrid.N_MAX_FACE_NODES_DIM + + ds = ds.swap_dims(dim_dict) + + return ds, dim_dict def _encode_ugrid(ds): - """Encodes UGRID file . - Parameters - ---------- - ds : xarray.Dataset - Dataset to be encoded to file + """Encodes an unstructured grid represented under a ``Grid`` object as a + ``xr.Dataset`` with an updated grid topology variable.""" + + if "grid_topology" in ds: + ds = ds.drop_vars(["grid_topology"]) + + grid_topology = ugrid.BASE_GRID_TOPOLOGY_ATTRS + + if "n_edge" in ds: + grid_topology["edge_dimension"] = "n_edge" + + if "face_lon" in ds: + grid_topology["face_coordinates"] = "face_lon face_lat" + if "edge_lon" in ds: + grid_topology["edge_coordinates"] = "edge_lon edge_lat" + + # TODO: Encode spherical (i.e. node_x) coordinates eventually (need to extend ugrid conventions) + + for conn_name in ugrid.CONNECTIVITY_NAMES: + if conn_name in ds: + grid_topology[conn_name] = conn_name + + grid_topology_da = xr.DataArray(data=-1, attrs=grid_topology) + + ds["grid_topology"] = grid_topology_da - Uses to_netcdf from xarray object. - """ return ds -def _standardize_fill_values(ds): - """Standardizes the fill values and data type of index variables. +def _standardize_connectivity(ds, conn_name): + """Standardizes the fill values and data type for a given connectivity + variable. Parameters ---------- @@ -83,31 +125,37 @@ def _standardize_fill_values(ds): Input Dataset with correct index variables """ - # original face nodes - face_nodes = ds["face_node_connectivity"].values + # original connectivity + conn = ds[conn_name].values # original fill value, if one exists - if "_FillValue" in ds["face_node_connectivity"].attrs: - original_fv = ds["face_node_connectivity"]._FillValue - elif np.isnan(ds["face_node_connectivity"].values).any(): + if "_FillValue" in ds[conn_name].attrs: + original_fv = ds[conn_name]._FillValue + elif np.isnan(ds[conn_name].values).any(): original_fv = np.nan else: original_fv = None # if current dtype and fill value are not standardized - if face_nodes.dtype != INT_DTYPE or original_fv != INT_FILL_VALUE: + if conn.dtype != INT_DTYPE or original_fv != INT_FILL_VALUE: # replace fill values and set correct dtype - new_face_nodes = _replace_fill_values( - grid_var=face_nodes, + new_conn = _replace_fill_values( + grid_var=conn, original_fill=original_fv, new_fill=INT_FILL_VALUE, new_dtype=INT_DTYPE, ) - # reassign data to use updated face nodes - ds["face_node_connectivity"].data = new_face_nodes + + if "start_index" in ds[conn_name].attrs: + new_conn -= INT_DTYPE(ds[conn_name].start_index) + else: + new_conn -= INT_DTYPE(new_conn.min()) + + # reassign data to use updated connectivity + ds[conn_name].data = new_conn # use new fill value - ds["face_node_connectivity"].attrs["_FillValue"] = INT_FILL_VALUE + ds[conn_name].attrs["_FillValue"] = INT_FILL_VALUE return ds