From d6e26e9d0ea8e9c43997eb093829f0d518ca8eda Mon Sep 17 00:00:00 2001 From: Aaron Zedwick <95507181+aaronzedwick@users.noreply.github.com> Date: Wed, 9 Oct 2024 15:59:32 -0500 Subject: [PATCH] Dual Mesh Construction (#859) * Added dual mesh initial support * Updated Mesh Construction Method * Removed Old Files * Updated API added comprehensive test * Fixed pre-commit * optimize code * update dual mesh construction and duplicate node validation * progress on migrating implementation to Grid * Testing Notebook * Update Untitled.ipynb * Updated Notebook * Update dual-mesh.ipynb * Update Notebook * Update docs * Update dual-mesh.ipynb * Updated Merge Duplicate Nodes and cleaned up code * Added Duplicate Nodes Test * Duplicate nodes error, updated index, fix codecov, * Fixed pre-commit * fixed duplicate check * Moved to dual.py remove duplication merge * Updated API and Benchmarks * Updated face construction and ordering * Optimization and clean up * Update dual-mesh.ipynb * Updated userguide.rst * Fixed pre-commit * Update docs/userguide.rst Co-authored-by: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> * Added docstrings * Update dual.py * Added dual mesh support of DataArrays * Fixed asv config file * Added name to data array * Fixed constructed node face connectivity error * Updated user guide * Fixed docs * Removed method parameter * Added support for UxDatasets * Update dual user guide * Update index.rst * Updated Test Cases * Update test_grid.py * Updated colormaps * pre-commit fix * Notebook Update * Update dual-mesh.ipynb * Update dual-mesh.ipynb * Fixed Dimension Bug * Fixed pre-commit * Update mpas_ocean.py * Added warnings * Fixed test failures * Updated notebook and jit * Merge pre-commit fix * Updated notebook * Update Dual Notebook * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Updated docstrings, benchmarks, and errors * update notebook to use updated plotting api --------- Co-authored-by: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Co-authored-by: Philip Chmielowiec Co-authored-by: Rajeev Jain Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- benchmarks/mpas_ocean.py | 7 + docs/api.rst | 19 ++ docs/internal_api/index.rst | 3 + docs/user-guide/dual-mesh.ipynb | 358 ++++++++++++++++++++++++++++++++ docs/userguide.rst | 4 + test/test_dataarray.py | 10 +- test/test_dataset.py | 9 + test/test_grid.py | 35 +++- uxarray/core/dataarray.py | 47 +++++ uxarray/core/dataset.py | 53 +++++ uxarray/grid/connectivity.py | 60 ++---- uxarray/grid/dual.py | 239 +++++++++++++++++++++ uxarray/grid/grid.py | 34 ++- uxarray/grid/validation.py | 47 ++++- 14 files changed, 876 insertions(+), 49 deletions(-) create mode 100644 docs/user-guide/dual-mesh.ipynb create mode 100644 uxarray/grid/dual.py diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 0c761434c..e8f783aa3 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -95,6 +95,7 @@ def time_dataarray_to_polycollection(self, resolution, periodic_elements): class ConstructTreeStructures(DatasetBenchmark): + def time_kd_tree(self, resolution): self.uxds.uxgrid.get_kd_tree() @@ -137,6 +138,11 @@ class HoleEdgeIndices(DatasetBenchmark): def time_construct_hole_edge_indices(self, resolution): ux.grid.geometry._construct_hole_edge_indices(self.uxds.uxgrid.edge_face_connectivity) + +class DualMesh(DatasetBenchmark): + def time_dual_mesh_construction(self, resolution): + self.uxds.uxgrid.get_dual() + class ConstructFaceLatLon(GridBenchmark): def time_welzl(self, resolution): self.uxgrid.construct_face_centers(method='welzl') @@ -144,6 +150,7 @@ def time_welzl(self, resolution): def time_cartesian_averaging(self, resolution): self.uxgrid.construct_face_centers(method='cartesian average') + class CheckNorm: param_names = ['resolution'] params = ['480km', '120km'] diff --git a/docs/api.rst b/docs/api.rst index 1c5b62ef6..4c6469279 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -331,6 +331,16 @@ Mathematical Operators UxDataArray.gradient UxDataArray.difference + +Dual Mesh Construction +---------------------- +.. autosummary:: + :toctree: generated/ + + Grid.get_dual + UxDataArray.get_dual + UxDataset.get_dual + Aggregations ------------ @@ -362,6 +372,15 @@ on each face. +Intersections +~~~~~~~~~~~~~ + +.. autosummary:: + :toctree: generated/ + + grid.intersections.gca_gca_intersection + grid.intersections.gca_const_lat_intersection + Spherical Geometry ------------------ diff --git a/docs/internal_api/index.rst b/docs/internal_api/index.rst index 34586a506..c77efca7e 100644 --- a/docs/internal_api/index.rst +++ b/docs/internal_api/index.rst @@ -200,6 +200,9 @@ Validation grid.validation._check_connectivity grid.validation._check_duplicate_nodes grid.validation._check_area + grid.validation._check_duplicate_nodes_indices + grid.validation._find_duplicate_nodes + Accurate Computing Utils ------------------------ diff --git a/docs/user-guide/dual-mesh.ipynb b/docs/user-guide/dual-mesh.ipynb new file mode 100644 index 000000000..f1147aeb7 --- /dev/null +++ b/docs/user-guide/dual-mesh.ipynb @@ -0,0 +1,358 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0660f9ec-8fa3-44a7-b362-3aa6526d8550", + "metadata": {}, + "source": [ + "# Dual Grid Construction" + ] + }, + { + "cell_type": "markdown", + "id": "974ebf95-c0b1-4397-be4d-bd5cf42cc8bc", + "metadata": {}, + "source": [ + "A dual grid is a secondary grid built from a grid, by constructing new faces over the original grids nodes. One property of this means that if you take the dual grid of a dual grid, the orignal grid will be constructed. However, this will only work properly if the grid is not partial. UXarray allows the constructing of this dual grid on all three of our data structures using the following methods:\n", + "\n", + "* `Grid.get_dual()`\n", + "* `UxDataArray.get_dual()`\n", + "* `UxDataset.get_dual()`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4fd7d3b3-5186-4157-841e-1b4c9f38bb62", + "metadata": { + "ExecuteTime": { + "start_time": "2024-10-09T20:06:48.348972Z" + }, + "jupyter": { + "is_executing": true + } + }, + "outputs": [], + "source": [ + "import uxarray as ux\n", + "import warnings\n", + "import cartopy.crs as ccrs\n", + "\n", + "warnings.filterwarnings(\"ignore\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1eebd05d-506e-4227-aa8e-f1212aae30f2", + "metadata": {}, + "outputs": [], + "source": [ + "file = \"../../test/meshfiles/mpas/QU/mesh.QU.1920km.151026.nc\"\n", + "\n", + "uxds = ux.open_dataset(file, file)" + ] + }, + { + "cell_type": "markdown", + "id": "1e6ddbd8-7034-4ddd-bdd0-4454d76ea390", + "metadata": {}, + "source": [ + "## Computing the Dual Grid" + ] + }, + { + "cell_type": "markdown", + "id": "dc3b2464-b824-413f-b2e6-40b1e7049f3b", + "metadata": {}, + "source": [ + "When computing the dual of a grid, the original grid is typically referred to as the \"Primal\" grid. The corner nodes of the Primal grid become the face centers of the Dual grid, and the face centers of the Primal grid become the corner nodes of the Dual grid. This means that variables that were originally face-centered on the Primal grid will be node-centered on the Dual grid, and vice versa. Using UXarray we can construct the dual mesh using `grid.compute_dual()`, which returns a new `Grid` object." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6486e0e4-1908-4363-8196-2efa9e521e56", + "metadata": {}, + "outputs": [], + "source": [ + "grid = uxds.uxgrid" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "925a82d6-2947-4b15-a3ff-0b0612bc5ba0", + "metadata": {}, + "outputs": [], + "source": [ + "dual = grid.get_dual()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12e9a8a5-526f-4b30-9a0b-de2d6d809f1d", + "metadata": {}, + "outputs": [], + "source": [ + "(\n", + " grid.plot(title=\"Primal Grid\", backend=\"bokeh\", projection=ccrs.Orthographic())\n", + " + dual.plot(title=\"Dual Grid\", backend=\"bokeh\", projection=ccrs.Orthographic())\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "550be213-245f-4c2e-8a59-c18b36d58e0f", + "metadata": {}, + "outputs": [], + "source": [ + "(\n", + " grid.plot(backend=\"bokeh\", projection=ccrs.Orthographic(), color=\"red\")\n", + " * dual.plot(backend=\"bokeh\", projection=ccrs.Orthographic(), color=\"blue\")\n", + ").opts(title=\"Primal (Red) & Dual (Blue) Grids\")" + ] + }, + { + "cell_type": "markdown", + "id": "721498b4-d587-4e4e-b1dc-6d8948a7f3f6", + "metadata": {}, + "source": [ + "## Computing the Dual Grid with Data" + ] + }, + { + "cell_type": "markdown", + "id": "9640cb00-6896-49de-b607-1721a8462d53", + "metadata": {}, + "source": [ + "We can also take a dual mesh of a `UxDataArray`. The concept for constructing the `Grid` remains the same, and what is constructed will be identical. The difference is the data stored inside the `UxDataArray` will be transfered with the dual mesh. The key differences is the location that the data is stored. The data transfer process works as follows:\n", + "\n", + "* Face centered data becomes node centered, as each face becomes a node in the dual mesh.\n", + "* Node centered data becomes face centered, as each node becomes a face in the dual mesh.\n", + "* Edge centered data remains unchanged, as the edge centers will remain in the same place, despite the edges themselves being different." + ] + }, + { + "cell_type": "markdown", + "id": "c55b2285-046d-460c-8e64-1bb204a36177", + "metadata": {}, + "source": [ + "### Face Centered Data" + ] + }, + { + "cell_type": "markdown", + "id": "ebdbc2da-938b-455a-9c84-090faf0e25f9", + "metadata": {}, + "source": [ + "When constructing the dual mesh paired with a face centered variable, the data becomes node centered. We can then plot this using `topological_mean` to get the dual data to the faces for proper visualization comparisions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "75e86af7-de9c-4013-91f6-9fd9e9c0c3ad", + "metadata": {}, + "outputs": [], + "source": [ + "uxds_dual_face = uxds[\"latCell\"].get_dual()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ff1407d4-910e-4b89-a6eb-6c4fbed9488a", + "metadata": {}, + "outputs": [], + "source": [ + "(\n", + " uxds[\"latCell\"].plot.polygons(\n", + " rasterize=True,\n", + " backend=\"matplotlib\",\n", + " title=\"Face centered data on Primal Mesh\",\n", + " cmap=ux.cmaps.sequential_green_blue,\n", + " projection=ccrs.Orthographic(),\n", + " )\n", + " + uxds_dual_face.topological_mean(destination=\"face\").plot.polygons(\n", + " rasterize=True,\n", + " backend=\"matplotlib\",\n", + " title=\"Node Centered Data on Dual Mesh\",\n", + " cmap=ux.cmaps.sequential_green_blue,\n", + " projection=ccrs.Orthographic(),\n", + " )\n", + ").cols(1).opts(fig_size=125)" + ] + }, + { + "cell_type": "markdown", + "id": "62d27519-aced-47b0-841e-447bf3a200c2", + "metadata": {}, + "source": [ + "### Node Centered Data" + ] + }, + { + "cell_type": "markdown", + "id": "6819dd2b-91a1-465f-8e2e-11cf14b81a7e", + "metadata": {}, + "source": [ + "When constructing the dual mesh paired with a node centered variable, the data becomes face centered. A benefit of computing the dual of a node centered variable is that it allows us to visualize the data as polygons." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c68473ae-dac6-422e-b322-71cfbfe3f876", + "metadata": {}, + "outputs": [], + "source": [ + "uxds_dual_node = uxds[\"latVertex\"].get_dual()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1cec4e98-d35f-41cc-b3a2-b2bcf1db2094", + "metadata": {}, + "outputs": [], + "source": [ + "(\n", + " uxds[\"latVertex\"]\n", + " .topological_mean(destination=\"face\")\n", + " .plot.polygons(\n", + " rasterize=True,\n", + " backend=\"matplotlib\",\n", + " title=\"Face centered data on Primal Mesh\",\n", + " cmap=ux.cmaps.sequential_green_blue,\n", + " projection=ccrs.Orthographic(),\n", + " )\n", + " + uxds_dual_node.plot.polygons(\n", + " rasterize=True,\n", + " backend=\"matplotlib\",\n", + " title=\"Node Centered Data on Dual Mesh\",\n", + " cmap=ux.cmaps.sequential_green_blue,\n", + " projection=ccrs.Orthographic(),\n", + " )\n", + ").cols(1).opts(fig_size=125)" + ] + }, + { + "cell_type": "markdown", + "id": "99ab5fd9-76ea-4b11-91ee-95bf535e3e7f", + "metadata": {}, + "source": [ + "### Edge Centered Data" + ] + }, + { + "cell_type": "markdown", + "id": "4a4b8f5a-4158-4f45-b7a7-9e82fcdac353", + "metadata": {}, + "source": [ + "When constructing the dual mesh paired with an edge centered variable, the data will stay edge centered. However, a plotting example cannot be shown, as the `topological_mean` needed to visualize edge centered data is not currently implemented for edge centered data." + ] + }, + { + "cell_type": "markdown", + "id": "1dafeee0-3d1d-49dd-aa41-8222463ee3b3", + "metadata": {}, + "source": [ + "### UxDataset" + ] + }, + { + "cell_type": "markdown", + "id": "36df7eac-d50d-4bae-883c-310529a1a5f0", + "metadata": {}, + "source": [ + "We can also construct a dual mesh from an entire dataset, which will convert the whole `UxDataset` to its dual mesh form. Below we can see the dataset before the dual mesh is constructed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "97a884d0-5966-4bfc-b505-ceb3c37c4b47", + "metadata": {}, + "outputs": [], + "source": [ + "uxds" + ] + }, + { + "cell_type": "markdown", + "id": "c65085a8-4111-4019-860c-dd28a7f3aefc", + "metadata": {}, + "source": [ + "Now we can construct the dual and see the new dataset that is returned." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9c14a1bd-d06a-4ccb-92a5-9054e063e740", + "metadata": {}, + "outputs": [], + "source": [ + "uxds_dual = uxds.get_dual()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7c039b32-03cd-4b83-b41e-dc5e06d118a0", + "metadata": {}, + "outputs": [], + "source": [ + "uxds_dual" + ] + }, + { + "cell_type": "markdown", + "id": "fcac470d-4025-4fe8-afde-4a7bdf05c3ec", + "metadata": {}, + "source": [ + "As you can see, transforms the whole dataset. We can now take any variable and plot it, as shown below." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "85dedc57-f179-43b9-a55c-c0b469330988", + "metadata": {}, + "outputs": [], + "source": [ + "uxds_dual[\"latVertex\"].plot.polygons(\n", + " rasterize=True,\n", + " backend=\"matplotlib\",\n", + " title=\"latVertex from UxDataset dual mesh\",\n", + " cmap=ux.cmaps.sequential_green_blue,\n", + " projection=ccrs.Orthographic(),\n", + ").opts(fig_size=120)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/userguide.rst b/docs/userguide.rst index 3342d99f7..007079780 100644 --- a/docs/userguide.rst +++ b/docs/userguide.rst @@ -58,6 +58,9 @@ These user guides provide detailed explanations of the core functionality in UXa `Face Area Calculations `_ Methods for computing the area of each face +`Dual Mesh Construction `_ + Construct the Dual Mesh of an unstructured grid + Supplementary Guides -------------------- @@ -84,4 +87,5 @@ These user guides provide additional detail about specific features in UXarray. user-guide/calculus.ipynb user-guide/tree_structures.ipynb user-guide/area_calc.ipynb + user-guide/dual-mesh.ipynb user-guide/holoviz.ipynb diff --git a/test/test_dataarray.py b/test/test_dataarray.py index 9e5a4a027..092c658af 100644 --- a/test/test_dataarray.py +++ b/test/test_dataarray.py @@ -9,7 +9,7 @@ from uxarray.grid.geometry import _build_polygon_shells, _build_corrected_polygon_shells -from uxarray.core.dataset import UxDataset +from uxarray.core.dataset import UxDataset, UxDataArray current_path = Path(os.path.dirname(os.path.realpath(__file__))) @@ -33,6 +33,14 @@ def test_to_dataset(self): assert isinstance(uxds_converted, UxDataset) assert uxds_converted.uxgrid == uxds.uxgrid + def test_get_dual(self): + """Tests the creation of the dual mesh on a data array.""" + uxds = ux.open_dataset(gridfile_ne30, dsfile_var2_ne30) + dual = uxds['psi'].get_dual() + + assert isinstance(dual, UxDataArray) + self.assertTrue(dual._node_centered()) + class TestGeometryConversions(TestCase): diff --git a/test/test_dataset.py b/test/test_dataset.py index a4efb6573..33833bef0 100644 --- a/test/test_dataset.py +++ b/test/test_dataset.py @@ -5,6 +5,7 @@ import xarray as xr import uxarray as ux +from uxarray import UxDataset try: import constants @@ -67,6 +68,14 @@ def test_ugrid_dim_names(self): for dim in ugrid_dims: assert dim in uxds_remap.dims + def test_get_dual(self): + """Tests the creation of the dual mesh on a data set.""" + uxds = ux.open_dataset(gridfile_ne30, dsfile_var2_ne30) + dual = uxds.get_dual() + + assert isinstance(dual, UxDataset) + self.assertTrue(len(uxds.data_vars) == len(dual.data_vars)) + # commented out due to often failures on the dataset hosting # def test_read_from_https(self): # """Tests reading a dataset from a HTTPS link.""" diff --git a/test/test_grid.py b/test/test_grid.py index 2e5547311..81a6d93d0 100644 --- a/test/test_grid.py +++ b/test/test_grid.py @@ -6,6 +6,7 @@ from unittest import TestCase from pathlib import Path +import uxarray import uxarray as ux from uxarray.grid.connectivity import _populate_face_edge_connectivity, _build_edge_face_connectivity, \ @@ -17,6 +18,8 @@ from uxarray.grid.arcs import extreme_gca_latitude +from uxarray.grid.validation import _find_duplicate_nodes + try: import constants except ImportError: @@ -31,6 +34,8 @@ gridfile_fesom = current_path / "meshfiles" / "ugrid" / "fesom" / "fesom.mesh.diag.nc" gridfile_geoflow = current_path / "meshfiles" / "ugrid" / "geoflow-small" / "grid.nc" gridfile_mpas = current_path / 'meshfiles' / "mpas" / "QU" / 'mesh.QU.1920km.151026.nc' +gridfile_mpas_two = current_path / 'meshfiles' / "mpas" / "QU" / 'oQU480.231010.nc' +gridfile_geos = current_path / 'meshfiles' / "geos-cs" / "c12" / 'test-c12.native.nc4' gridfile_mpas_holes = current_path / 'meshfiles' / "mpas" / "QU" / 'oQU480.231010.nc' dsfile_vortex_CSne30 = current_path / "meshfiles" / "ugrid" / "outCSne30" / "outCSne30_vortex.nc" @@ -954,8 +959,34 @@ def test_populate_bounds_GCA_mix(self): nt.assert_allclose(grid.bounds.values, expected_bounds, atol=ERROR_TOLERANCE) def test_populate_bounds_MPAS(self): - uxgrid = ux.open_grid(self.gridfile_mpas) - bounds_xarray = uxgrid.bounds + uxgrid = ux.open_grid(self.gridfile_mpas) + bounds_xarray = uxgrid.bounds + + +class TestDualMesh(TestCase): + """Test Dual Mesh Construction.""" + + def test_dual_mesh_mpas(self): + # Open a grid with and without dual + grid = ux.open_grid(gridfile_mpas, use_dual=False) + mpas_dual = ux.open_grid(gridfile_mpas, use_dual=True) + + # Construct Dual + dual = grid.get_dual() + + # Assert the dimensions are the same + assert dual.n_face == mpas_dual.n_face + assert dual.n_node == mpas_dual.n_node + assert dual.n_max_face_nodes == mpas_dual.n_max_face_nodes + + # Assert the faces are the same + nt.assert_equal(dual.face_node_connectivity.values, mpas_dual.face_node_connectivity.values) + + def test_dual_duplicate(self): + # Test that dual mesh throws an exception if duplicate nodes exist + dataset = ux.open_dataset(gridfile_geoflow, gridfile_geoflow) + + nt.assert_raises(RuntimeError, dataset.get_dual) class TestNormalizeExistingCoordinates(TestCase): diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 3b64814da..ece6387cb 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Optional, Hashable, Literal +import uxarray from uxarray.formatting_html import array_repr from html import escape @@ -14,6 +15,8 @@ from uxarray.grid import Grid import uxarray.core.dataset +from uxarray.grid.dual import construct_dual +from uxarray.grid.validation import _check_duplicate_nodes_indices if TYPE_CHECKING: from uxarray.core.dataset import UxDataset @@ -33,6 +36,7 @@ from uxarray.core.aggregation import _uxda_grid_aggregate import warnings +from warnings import warn import cartopy.crs as ccrs @@ -1108,3 +1112,46 @@ def _slice_from_grid(self, sliced_grid): dims=self.dims, attrs=self.attrs, ) + + def get_dual(self): + """Compute the dual mesh for a data array, returns a new data array + object. + + Returns + -------- + dual : uxda + Dual Mesh `uxda` constructed + """ + + if _check_duplicate_nodes_indices(self.uxgrid): + raise RuntimeError("Duplicate nodes found, cannot construct dual") + + if self.uxgrid.hole_edge_indices.size != 0: + warn( + "This mesh is partial, which could cause inconsistent results and data will be lost", + Warning, + ) + + # Get dual mesh node face connectivity + dual_node_face_conn = construct_dual(grid=self.uxgrid) + + # Construct dual mesh + dual = self.uxgrid.from_topology( + self.uxgrid.face_lon.values, + self.uxgrid.face_lat.values, + dual_node_face_conn, + ) + + # Dictionary to swap dimensions + dim_map = {"n_face": "n_node", "n_node": "n_face"} + + # Get correct dimensions for the dual + dims = [dim_map.get(dim, dim) for dim in self.dims] + + # Get the values from the data array + data = np.array(self.values) + + # Construct the new data array + uxda = uxarray.UxDataArray(uxgrid=dual, data=data, dims=dims, name=self.name) + + return uxda diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index fe61bbf34..a9ad3fbdf 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -7,8 +7,11 @@ from typing import Optional, IO +import uxarray from uxarray.grid import Grid from uxarray.core.dataarray import UxDataArray +from uxarray.grid.dual import construct_dual +from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.plot.accessor import UxDatasetPlotAccessor @@ -330,3 +333,53 @@ def to_array(self) -> UxDataArray: xarr = super().to_array() return UxDataArray(xarr, uxgrid=self.uxgrid) + + def get_dual(self): + """Compute the dual mesh for a dataset, returns a new dataset object. + + Returns + -------- + dual : uxds + Dual Mesh `uxds` constructed + """ + + if _check_duplicate_nodes_indices(self.uxgrid): + raise RuntimeError("Duplicate nodes found, cannot construct dual") + + if self.uxgrid.hole_edge_indices.size != 0: + warn( + "This mesh is partial, which could cause inconsistent results and data will be lost", + Warning, + ) + + # Get dual mesh node face connectivity + dual_node_face_conn = construct_dual(grid=self.uxgrid) + + # Construct dual mesh + dual = self.uxgrid.from_topology( + self.uxgrid.face_lon.values, + self.uxgrid.face_lat.values, + dual_node_face_conn, + ) + + # Initialize new dataset + dataset = uxarray.UxDataset(uxgrid=dual) + + # Dictionary to swap dimensions + dim_map = {"n_face": "n_node", "n_node": "n_face"} + + # For each data array in the dataset, reconstruct the data array with the dual mesh + for var in self.data_vars: + # Get correct dimensions for the dual + dims = [dim_map.get(dim, dim) for dim in self[var].dims] + + # Get the values from the data array + data = np.array(self[var].values) + + # Construct the new data array + uxda = uxarray.UxDataArray(uxgrid=dual, data=data, dims=dims, name=var) + + # Add data array to dataset + dataset[var] = uxda + + return dataset diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 4fdddd722..d94d9eba6 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -1,7 +1,6 @@ import numpy as np import xarray as xr -from scipy import sparse from uxarray.constants import INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid @@ -330,48 +329,27 @@ def _build_node_faces_connectivity(face_nodes, n_node): RuntimeError If the Mesh object does not contain a 'face_node_connectivity' variable. """ - # First we need to build a matrix such that: the row indices are face indexes and the column indices are node - # indexes (similar to an adjacency matrix) - face_indices, node_indices, non_filled_element_flags = _face_nodes_to_sparse_matrix( - face_nodes - ) - coo_matrix = sparse.coo_matrix( - (non_filled_element_flags, (node_indices, face_indices)) + + node_face_conn = {node_i: [] for node_i in range(n_node)} + for face_i, face_nodes in enumerate(face_nodes): + for node_i in face_nodes: + if node_i != INT_FILL_VALUE: + node_face_conn[node_i].append(face_i) + + n_max_node_faces = -1 + for face_indicies in node_face_conn.values(): + if len(face_indicies) > n_max_node_faces: + n_max_node_faces = len(face_indicies) + + node_face_connectivity = np.full( + (n_node, n_max_node_faces), INT_FILL_VALUE, dtype=INT_DTYPE ) - csr_matrix = coo_matrix.tocsr() - # get the row and column indices of the non-zero elements - rows, cols = csr_matrix.nonzero() - # Find the frequency of each face to determine the maximum number of faces per node - freq = np.bincount(rows) - nMaxNumFacesPerNode = freq.max() - - node_face_connectivity = [[]] * n_node - - # find the indices where the array changes value - change_indices = np.where(np.diff(rows) != 0)[0] + 1 - - # split the array at the change indices to get subarrays of consecutive same elements - subarrays = np.split(rows, change_indices) - - # get the start and end indices for each subarray - start_indices = np.cumsum([0] + [len(subarray) for subarray in subarrays[:-1]]) - end_indices = np.cumsum([len(subarray) for subarray in subarrays]) - 1 - - for node_index in range(n_node): - node_face_connectivity[node_index] = cols[ - start_indices[node_index] : end_indices[node_index] + 1 - ] - if len(node_face_connectivity[node_index]) < nMaxNumFacesPerNode: - node_face_connectivity[node_index] = np.append( - node_face_connectivity[node_index], - np.full( - nMaxNumFacesPerNode - len(node_face_connectivity[node_index]), - INT_FILL_VALUE, - dtype=INT_DTYPE, - ), - ) - return node_face_connectivity, nMaxNumFacesPerNode + for node_idx, face_indices in enumerate(node_face_conn.values()): + n_faces = len(face_indices) + node_face_connectivity[node_idx, 0:n_faces] = face_indices + + return node_face_connectivity, n_max_node_faces def _face_nodes_to_sparse_matrix(dense_matrix: np.ndarray) -> tuple: diff --git a/uxarray/grid/dual.py b/uxarray/grid/dual.py new file mode 100644 index 000000000..5dc5cb4f5 --- /dev/null +++ b/uxarray/grid/dual.py @@ -0,0 +1,239 @@ +import numpy as np +from uxarray.constants import INT_FILL_VALUE, INT_DTYPE + +from numba import njit +from uxarray.constants import ENABLE_JIT_CACHE + + +def construct_dual(grid): + """Constructs a dual mesh from a given grid, by connecting the face centers + to make a grid centered over the primal mesh face centers, with the nodes + of the primal mesh being the face centers of the dual mesh. + + Parameters + ---------- + grid: Grid + Primal mesh to construct dual from + + Returns + -------- + new_node_face_connectivity : ndarray + Constructed node_face_connectivity for the dual mesh + """ + + # Get the dual node xyz, which is the face centers + dual_node_x = grid.face_x.values + dual_node_y = grid.face_y.values + dual_node_z = grid.face_z.values + + # Get other information from the grid needed + n_node = grid.n_node + node_x = grid.node_x.values + node_y = grid.node_y.values + node_z = grid.node_z.values + node_face_connectivity = grid.node_face_connectivity.values + + # Get an array with the number of edges for each face + n_edges_mask = node_face_connectivity != INT_FILL_VALUE + n_edges = np.sum(n_edges_mask, axis=1) + + # Construct and return the faces + new_node_face_connectivity = construct_faces( + n_node, + n_edges, + dual_node_x, + dual_node_y, + dual_node_z, + node_face_connectivity, + node_x, + node_y, + node_z, + ) + + return new_node_face_connectivity + + +@njit(cache=ENABLE_JIT_CACHE) +def construct_faces( + n_node, + n_edges, + dual_node_x, + dual_node_y, + dual_node_z, + node_face_connectivity, + node_x, + node_y, + node_z, +): + """Construct the faces of the dual mesh based on a given + node_face_connectivity. + + Parameters + ---------- + n_node: np.ndarray + number of nodes in the primal mesh + n_edges: np.ndarray + array of the number of edges for each dual face + dual_node_x: np.ndarray + x node coordinates for the dual mesh + dual_node_y: np.ndarray + y node coordinates for the dual mesh + dual_node_z: np.ndarray + z node coordinates for the dual mesh + node_face_connectivity: np.ndarray + `node_face_connectivity` of the primal mesh + node_x: np.ndarray + x node coordinates from the primal mesh + node_y: np.ndarray + y node coordinates from the primal mesh + node_z: np.ndarray + z node coordinates from the primal mesh + + + Returns + -------- + node_face_connectivity : ndarray + Constructed node_face_connectivity for the dual mesh + """ + correction = 0 + max_edges = len(node_face_connectivity[0]) + construct_node_face_connectivity = np.full( + (np.sum(n_edges > 2), max_edges), INT_FILL_VALUE, dtype=INT_DTYPE + ) + for i in range(n_node): + # If we have less than 3 edges, we can't construct anything but a line + if n_edges[i] < 3: + correction += 1 + continue + + # Construct temporary face to hold unordered face nodes + temp_face = np.array( + [INT_FILL_VALUE for _ in range(n_edges[i])], dtype=INT_DTYPE + ) + + # Get a list of the valid non fill value nodes + valid_node_indices = node_face_connectivity[i][0 : n_edges[i]] + index = 0 + + # Connect the face centers around the node to make dual face + for node_idx in valid_node_indices: + temp_face[index] = node_idx + index += 1 + + # Order the nodes using the angles so the faces have nodes in counter-clockwise sequence + node_central = np.array([node_x[i], node_y[i], node_z[i]]) + node_0 = np.array( + [ + dual_node_x[temp_face[0]], + dual_node_y[temp_face[0]], + dual_node_z[temp_face[0]], + ] + ) + + # Order the face nodes properly in a counter-clockwise fashion + if temp_face[0] is not INT_FILL_VALUE: + _face = _order_nodes( + temp_face, + node_0, + node_central, + n_edges[i], + dual_node_x, + dual_node_y, + dual_node_z, + max_edges, + ) + construct_node_face_connectivity[i - correction] = _face + return construct_node_face_connectivity + + +@njit(cache=ENABLE_JIT_CACHE) +def _order_nodes( + temp_face, + node_0, + node_central, + n_edges, + dual_node_x, + dual_node_y, + dual_node_z, + max_edges, +): + """Correctly order the nodes of each face in a counter-clockwise fashion. + + Parameters + ---------- + temp_face: np.ndarray + Face holding unordered nodes + node_0: np.ndarray + Starting node + node_central: np.ndarray + Face center + n_edges: int + Number of edges in the face + dual_node_x: np.ndarray + x node coordinates for the dual mesh + dual_node_y: np.ndarray + y node coordinates for the dual mesh + dual_node_z: np.ndarray + z node coordinates for the dual mesh + max_edges: int + Max number of edges a face could have + + + Returns + -------- + final_face : np.ndarray + The face in proper counter-clockwise order + """ + node_zero = node_0 - node_central + + node_cross = np.cross(node_0, node_central) + node_zero_mag = np.linalg.norm(node_zero) + + d_angles = np.zeros(n_edges, dtype=np.float64) + d_angles[0] = 0.0 + final_face = np.array([INT_FILL_VALUE for _ in range(max_edges)], dtype=INT_DTYPE) + for j in range(1, n_edges): + _cur_face_temp_idx = temp_face[j] + + if _cur_face_temp_idx is not INT_FILL_VALUE: + sub = np.array( + [ + dual_node_x[_cur_face_temp_idx], + dual_node_y[_cur_face_temp_idx], + dual_node_z[_cur_face_temp_idx], + ] + ) + node_diff = sub - node_central + node_diff_mag = np.linalg.norm(node_diff) + + d_side = np.dot(node_cross, node_diff) + d_dot_norm = np.dot(node_zero, node_diff) / (node_zero_mag * node_diff_mag) + + if d_dot_norm > 1.0: + d_dot_norm = 1.0 + + d_angles[j] = np.arccos(d_dot_norm) + + if d_side > 0.0: + d_angles[j] = -d_angles[j] + 2.0 * np.pi + + d_current_angle = 0.0 + + final_face[0] = temp_face[0] + for j in range(1, n_edges): + ix_next_node = -1 + d_next_angle = 2.0 * np.pi + + for k in range(1, n_edges): + if d_current_angle < d_angles[k] < d_next_angle: + ix_next_node = k + d_next_angle = d_angles[k] + + if ix_next_node == -1: + continue + + final_face[j] = temp_face[ix_next_node] + + d_current_angle = d_next_angle + + return final_face diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 3c19609f1..1d9e52811 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -75,10 +75,15 @@ from uxarray.grid.validation import ( _check_connectivity, _check_duplicate_nodes, + _check_duplicate_nodes_indices, _check_area, _check_normalization, ) + +from uxarray.conventions import ugrid + + from xarray.core.utils import UncachedAccessor from warnings import warn @@ -87,6 +92,8 @@ import copy +from uxarray.grid.dual import construct_dual + class Grid: """Represents a two-dimensional unstructured grid encoded following the @@ -458,7 +465,7 @@ def construct_face_centers(self, method="cartesian average"): def __repr__(self): """Constructs a string representation of the contents of a ``Grid``.""" - from uxarray.conventions import ugrid, descriptors + from uxarray.conventions import descriptors prefix = "\n" original_grid_str = f"Original Grid Type: {self.source_grid_spec}\n" @@ -1910,3 +1917,28 @@ def isel(self, **dim_kwargs): raise ValueError( "Indexing must be along a grid dimension: ('n_node', 'n_edge', 'n_face')" ) + + def get_dual(self): + """Compute the dual for a grid, which constructs a new grid centered + around the nodes, where the nodes of the primal become the face centers + of the dual, and the face centers of the primal become the nodes of the + dual. Returns a new `Grid` object. + + Returns + -------- + dual : Grid + Dual Mesh Grid constructed + """ + + if _check_duplicate_nodes_indices(self): + raise RuntimeError("Duplicate nodes found, cannot construct dual") + + # Get dual mesh node face connectivity + dual_node_face_conn = construct_dual(grid=self) + + # Construct dual mesh + dual = self.from_topology( + self.face_lon.values, self.face_lat.values, dual_node_face_conn + ) + + return dual diff --git a/uxarray/grid/validation.py b/uxarray/grid/validation.py index b42b8cbb3..e4a11aca3 100644 --- a/uxarray/grid/validation.py +++ b/uxarray/grid/validation.py @@ -1,7 +1,8 @@ import numpy as np from warnings import warn -from uxarray.constants import ERROR_TOLERANCE + +from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE # validation helper functions @@ -20,7 +21,6 @@ def _check_connectivity(grid): # check if the size of unique nodes in connectivity is equal to the number of nodes if nodes_in_conn.size == grid.n_node: - print("-All nodes are referenced by at least one element.") return True else: warn( @@ -48,10 +48,23 @@ def _check_duplicate_nodes(grid): ) return False else: - print("-No duplicate nodes found in the mesh.") return True +def _check_duplicate_nodes_indices(grid): + """Check if there are duplicate node indices, returns True if there are.""" + + # Create a duplication dictionary + duplicate_node_dict = _find_duplicate_nodes(grid) + + for face_nodes in grid.face_node_connectivity.values: + for node in face_nodes: + if node in duplicate_node_dict.keys(): + return True + + return False + + def _check_area(grid): """Check if each face area is greater than our constant ERROR_TOLERANCE.""" areas = grid.face_areas @@ -63,10 +76,36 @@ def _check_area(grid): ) return False else: - print("-No face area is close to zero.") return True +def _find_duplicate_nodes(grid): + # list of tuple indices + lonlat_t = [ + (lon, lat) for lon, lat in zip(grid.node_lon.values, grid.node_lat.values) + ] + + # # Dictionary to track first occurrence and subsequent indices + occurrences = {} + + # Iterate through the list and track occurrences + for index, tpl in enumerate(lonlat_t): + if tpl in occurrences: + occurrences[tpl].append((INT_DTYPE(index))) + else: + occurrences[tpl] = [INT_DTYPE(index)] + + duplicate_dict = {} + + for tpl, indices in occurrences.items(): + if len(indices) > 1: + source_idx = indices[0] + for duplicate_idx in indices[1:]: + duplicate_dict[duplicate_idx] = source_idx + + return duplicate_dict + + def _check_normalization(grid): """Checks whether all the cartesiain coordinates are normalized."""