From 3b27d3f3ccf1d4a84a28f6538b2a91f79e051160 Mon Sep 17 00:00:00 2001 From: Hongyu Chen Date: Sat, 29 Jun 2024 21:39:03 -0700 Subject: [PATCH] Optimize the `get_cartesian_face_edge_nodes` and `get_lonlat_rad_face_edge_nodes` (#799) * Initial Draft Algorithm * Revert "Initial Draft Algorithm" This reverts commit 99573b8edfc2dc2fc329e91d395da4d907b42e4a. * Initial commit * Revert "Initial commit" This reverts commit 1ba68804f2f8f6fd062c9eeb63ef0b600549c9f0. * initial commit * Finished implementation * Fixed faces with holes * Fix precommit * test * add examples to doc strings * initial commit for vectorizing _get_cartesian_face_edge_nodes and _get_lonlat_rad_face_edge_nodes * add examples, modify test_helpers to test the new implementation(cur fail cartesian with fill) * Fix precommit * continue testcase * fixed input with INT_FILL_VALUE in face_node_conn * fix test case * more tests * rename helper * rename * update usage of _get_cartesian_face_edge_nodes and _get_lonlat_rad_face_edge_nodes * update helper index --------- Co-authored-by: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Co-authored-by: amberchen122 Co-authored-by: amberchen122 --- docs/internal_api/index.rst | 1 + test/test_geometry.py | 423 +++++++++++++++++++++++++++++++++--- test/test_grid.py | 7 + test/test_helpers.py | 304 +++++++++++++++++++++----- uxarray/grid/geometry.py | 42 ++-- uxarray/grid/utils.py | 278 +++++++++++++++++------- 6 files changed, 872 insertions(+), 183 deletions(-) diff --git a/docs/internal_api/index.rst b/docs/internal_api/index.rst index b233e2413..b263f4ec8 100644 --- a/docs/internal_api/index.rst +++ b/docs/internal_api/index.rst @@ -176,6 +176,7 @@ Utils grid.utils._newton_raphson_solver_for_gca_constLat grid.utils._inv_jacobian + grid.utils._swap_first_fill_value_with_last grid.utils._get_cartesiain_face_edge_nodes grid.utils._get_lonlat_rad_face_edge_nodes diff --git a/test/test_geometry.py b/test/test_geometry.py index b73e2d8b9..953942293 100644 --- a/test/test_geometry.py +++ b/test/test_geometry.py @@ -439,6 +439,128 @@ def test_insert_pt_in_empty_state(self): class TestLatlonBoundsGCA(TestCase): + + def _get_cartesian_face_edge_nodes_testcase_helper( + self,face_nodes_ind, face_edges_ind, edge_nodes_grid, node_x, node_y, node_z + ): + """This function is only used to help generating the testcase and + should not be used in the actual implementation. Construct an array to + hold the edge Cartesian coordinates connectivity for a face in a grid. + + Parameters + ---------- + face_nodes_ind : np.ndarray, shape (n_nodes,) + The ith entry of Grid.face_node_connectivity, where n_nodes is the number of nodes in the face. + face_edges_ind : np.ndarray, shape (n_edges,) + The ith entry of Grid.face_edge_connectivity, where n_edges is the number of edges in the face. + edge_nodes_grid : np.ndarray, shape (n_edges, 2) + The entire Grid.edge_node_connectivity, where n_edges is the total number of edges in the grid. + node_x : np.ndarray, shape (n_nodes,) + The values of Grid.node_x, where n_nodes_total is the total number of nodes in the grid. + node_y : np.ndarray, shape (n_nodes,) + The values of Grid.node_y. + node_z : np.ndarray, shape (n_nodes,) + The values of Grid.node_z. + + Returns + ------- + face_edges : np.ndarray, shape (n_edges, 2, 3) + Face edge connectivity in Cartesian coordinates, where n_edges is the number of edges for the specific face. + + Notes + ----- + - The function assumes that the inputs are well-formed and correspond to the same face. + - The output array contains the Cartesian coordinates for each edge of the face. + """ + + # Create a mask that is True for all values not equal to INT_FILL_VALUE + mask = face_edges_ind != INT_FILL_VALUE + + # Use the mask to select only the elements not equal to INT_FILL_VALUE + valid_edges = face_edges_ind[mask] + face_edges = edge_nodes_grid[valid_edges] + + # Ensure counter-clockwise order of edge nodes + # Start with the first two nodes + face_edges[0] = [face_nodes_ind[0], face_nodes_ind[1]] + + for idx in range(1, len(face_edges)): + if face_edges[idx][0] != face_edges[idx - 1][1]: + # Swap the node index in this edge if not in counter-clockwise order + face_edges[idx] = face_edges[idx][::-1] + + # Fetch coordinates for each node in the face edges + cartesian_coordinates = np.array( + [ + [[node_x[node], node_y[node], node_z[node]] for node in edge] + for edge in face_edges + ] + ) + + return cartesian_coordinates + + def _get_lonlat_rad_face_edge_nodes_testcase_helper( + self,face_nodes_ind, face_edges_ind, edge_nodes_grid, node_lon, node_lat + ): + """This function is only used to help generating the testcase and + should not be used in the actual implementation. Construct an array to + hold the edge lat lon in radian connectivity for a face in a grid. + + Parameters + ---------- + face_nodes_ind : np.ndarray, shape (n_nodes,) + The ith entry of Grid.face_node_connectivity, where n_nodes is the number of nodes in the face. + face_edges_ind : np.ndarray, shape (n_edges,) + The ith entry of Grid.face_edge_connectivity, where n_edges is the number of edges in the face. + edge_nodes_grid : np.ndarray, shape (n_edges, 2) + The entire Grid.edge_node_connectivity, where n_edges is the total number of edges in the grid. + node_lon : np.ndarray, shape (n_nodes,) + The values of Grid.node_lon. + node_lat : np.ndarray, shape (n_nodes,) + The values of Grid.node_lat, where n_nodes_total is the total number of nodes in the grid. + Returns + ------- + face_edges : np.ndarray, shape (n_edges, 2, 3) + Face edge connectivity in Cartesian coordinates, where n_edges is the number of edges for the specific face. + + Notes + ----- + - The function assumes that the inputs are well-formed and correspond to the same face. + - The output array contains the Latitude and longitude coordinates in radian for each edge of the face. + """ + + # Create a mask that is True for all values not equal to INT_FILL_VALUE + mask = face_edges_ind != INT_FILL_VALUE + + # Use the mask to select only the elements not equal to INT_FILL_VALUE + valid_edges = face_edges_ind[mask] + face_edges = edge_nodes_grid[valid_edges] + + # Ensure counter-clockwise order of edge nodes + # Start with the first two nodes + face_edges[0] = [face_nodes_ind[0], face_nodes_ind[1]] + + for idx in range(1, len(face_edges)): + if face_edges[idx][0] != face_edges[idx - 1][1]: + # Swap the node index in this edge if not in counter-clockwise order + face_edges[idx] = face_edges[idx][::-1] + + # Fetch coordinates for each node in the face edges + lonlat_coordinates = np.array( + [ + [ + [ + np.mod(np.deg2rad(node_lon[node]), 2 * np.pi), + np.deg2rad(node_lat[node]), + ] + for node in edge + ] + for edge in face_edges + ] + ) + + return lonlat_coordinates + def test_populate_bounds_normal(self): # Generate a normal face that is not crossing the antimeridian or the poles vertices_lonlat = [[10.0, 60.0], [10.0, 10.0], [50.0, 10.0], [50.0, 60.0]] @@ -456,12 +578,12 @@ def test_populate_bounds_normal(self): lon_min = np.deg2rad(10.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -485,12 +607,12 @@ def test_populate_bounds_antimeridian(self): lon_min = np.deg2rad(350.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -513,12 +635,12 @@ def test_populate_bounds_node_on_pole(self): lon_min = np.deg2rad(10.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -541,12 +663,12 @@ def test_populate_bounds_edge_over_pole(self): lon_min = np.deg2rad(210.0) lon_max = np.deg2rad(30.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -569,12 +691,12 @@ def test_populate_bounds_pole_inside(self): lon_min = 0 lon_max = 2 * np.pi grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -586,6 +708,127 @@ def test_populate_bounds_pole_inside(self): class TestLatlonBoundsLatLonFace(TestCase): + def _get_cartesian_face_edge_nodes_testcase_helper( + self,face_nodes_ind, face_edges_ind, edge_nodes_grid, node_x, node_y, node_z + ): + """This function is only used to help generating the testcase and + should not be used in the actual implementation. Construct an array to + hold the edge Cartesian coordinates connectivity for a face in a grid. + + Parameters + ---------- + face_nodes_ind : np.ndarray, shape (n_nodes,) + The ith entry of Grid.face_node_connectivity, where n_nodes is the number of nodes in the face. + face_edges_ind : np.ndarray, shape (n_edges,) + The ith entry of Grid.face_edge_connectivity, where n_edges is the number of edges in the face. + edge_nodes_grid : np.ndarray, shape (n_edges, 2) + The entire Grid.edge_node_connectivity, where n_edges is the total number of edges in the grid. + node_x : np.ndarray, shape (n_nodes,) + The values of Grid.node_x, where n_nodes_total is the total number of nodes in the grid. + node_y : np.ndarray, shape (n_nodes,) + The values of Grid.node_y. + node_z : np.ndarray, shape (n_nodes,) + The values of Grid.node_z. + + Returns + ------- + face_edges : np.ndarray, shape (n_edges, 2, 3) + Face edge connectivity in Cartesian coordinates, where n_edges is the number of edges for the specific face. + + Notes + ----- + - The function assumes that the inputs are well-formed and correspond to the same face. + - The output array contains the Cartesian coordinates for each edge of the face. + """ + + # Create a mask that is True for all values not equal to INT_FILL_VALUE + mask = face_edges_ind != INT_FILL_VALUE + + # Use the mask to select only the elements not equal to INT_FILL_VALUE + valid_edges = face_edges_ind[mask] + face_edges = edge_nodes_grid[valid_edges] + + # Ensure counter-clockwise order of edge nodes + # Start with the first two nodes + face_edges[0] = [face_nodes_ind[0], face_nodes_ind[1]] + + for idx in range(1, len(face_edges)): + if face_edges[idx][0] != face_edges[idx - 1][1]: + # Swap the node index in this edge if not in counter-clockwise order + face_edges[idx] = face_edges[idx][::-1] + + # Fetch coordinates for each node in the face edges + cartesian_coordinates = np.array( + [ + [[node_x[node], node_y[node], node_z[node]] for node in edge] + for edge in face_edges + ] + ) + + return cartesian_coordinates + + def _get_lonlat_rad_face_edge_nodes_testcase_helper( + self,face_nodes_ind, face_edges_ind, edge_nodes_grid, node_lon, node_lat + ): + """This function is only used to help generating the testcase and + should not be used in the actual implementation. Construct an array to + hold the edge lat lon in radian connectivity for a face in a grid. + + Parameters + ---------- + face_nodes_ind : np.ndarray, shape (n_nodes,) + The ith entry of Grid.face_node_connectivity, where n_nodes is the number of nodes in the face. + face_edges_ind : np.ndarray, shape (n_edges,) + The ith entry of Grid.face_edge_connectivity, where n_edges is the number of edges in the face. + edge_nodes_grid : np.ndarray, shape (n_edges, 2) + The entire Grid.edge_node_connectivity, where n_edges is the total number of edges in the grid. + node_lon : np.ndarray, shape (n_nodes,) + The values of Grid.node_lon. + node_lat : np.ndarray, shape (n_nodes,) + The values of Grid.node_lat, where n_nodes_total is the total number of nodes in the grid. + Returns + ------- + face_edges : np.ndarray, shape (n_edges, 2, 3) + Face edge connectivity in Cartesian coordinates, where n_edges is the number of edges for the specific face. + + Notes + ----- + - The function assumes that the inputs are well-formed and correspond to the same face. + - The output array contains the Latitude and longitude coordinates in radian for each edge of the face. + """ + + # Create a mask that is True for all values not equal to INT_FILL_VALUE + mask = face_edges_ind != INT_FILL_VALUE + + # Use the mask to select only the elements not equal to INT_FILL_VALUE + valid_edges = face_edges_ind[mask] + face_edges = edge_nodes_grid[valid_edges] + + # Ensure counter-clockwise order of edge nodes + # Start with the first two nodes + face_edges[0] = [face_nodes_ind[0], face_nodes_ind[1]] + + for idx in range(1, len(face_edges)): + if face_edges[idx][0] != face_edges[idx - 1][1]: + # Swap the node index in this edge if not in counter-clockwise order + face_edges[idx] = face_edges[idx][::-1] + + # Fetch coordinates for each node in the face edges + lonlat_coordinates = np.array( + [ + [ + [ + np.mod(np.deg2rad(node_lon[node]), 2 * np.pi), + np.deg2rad(node_lat[node]), + ] + for node in edge + ] + for edge in face_edges + ] + ) + + return lonlat_coordinates + def test_populate_bounds_normal(self): # Generate a normal face that is not crossing the antimeridian or the poles vertices_lonlat = [[10.0, 60.0], [10.0, 10.0], [50.0, 10.0], [50.0, 60.0]] @@ -598,12 +841,12 @@ def test_populate_bounds_normal(self): lon_min = np.deg2rad(10.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -625,12 +868,12 @@ def test_populate_bounds_antimeridian(self): lon_min = np.deg2rad(350.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -652,12 +895,12 @@ def test_populate_bounds_node_on_pole(self): lon_min = np.deg2rad(10.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -680,12 +923,12 @@ def test_populate_bounds_edge_over_pole(self): lon_min = np.deg2rad(210.0) lon_max = np.deg2rad(30.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -708,12 +951,12 @@ def test_populate_bounds_pole_inside(self): lon_min = 0 lon_max = 2 * np.pi grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -725,6 +968,126 @@ def test_populate_bounds_pole_inside(self): class TestLatlonBoundsGCAList(TestCase): + def _get_cartesian_face_edge_nodes_testcase_helper( + self,face_nodes_ind, face_edges_ind, edge_nodes_grid, node_x, node_y, node_z + ): + """This function is only used to help generating the testcase and + should not be used in the actual implementation. Construct an array to + hold the edge Cartesian coordinates connectivity for a face in a grid. + + Parameters + ---------- + face_nodes_ind : np.ndarray, shape (n_nodes,) + The ith entry of Grid.face_node_connectivity, where n_nodes is the number of nodes in the face. + face_edges_ind : np.ndarray, shape (n_edges,) + The ith entry of Grid.face_edge_connectivity, where n_edges is the number of edges in the face. + edge_nodes_grid : np.ndarray, shape (n_edges, 2) + The entire Grid.edge_node_connectivity, where n_edges is the total number of edges in the grid. + node_x : np.ndarray, shape (n_nodes,) + The values of Grid.node_x, where n_nodes_total is the total number of nodes in the grid. + node_y : np.ndarray, shape (n_nodes,) + The values of Grid.node_y. + node_z : np.ndarray, shape (n_nodes,) + The values of Grid.node_z. + + Returns + ------- + face_edges : np.ndarray, shape (n_edges, 2, 3) + Face edge connectivity in Cartesian coordinates, where n_edges is the number of edges for the specific face. + + Notes + ----- + - The function assumes that the inputs are well-formed and correspond to the same face. + - The output array contains the Cartesian coordinates for each edge of the face. + """ + + # Create a mask that is True for all values not equal to INT_FILL_VALUE + mask = face_edges_ind != INT_FILL_VALUE + + # Use the mask to select only the elements not equal to INT_FILL_VALUE + valid_edges = face_edges_ind[mask] + face_edges = edge_nodes_grid[valid_edges] + + # Ensure counter-clockwise order of edge nodes + # Start with the first two nodes + face_edges[0] = [face_nodes_ind[0], face_nodes_ind[1]] + + for idx in range(1, len(face_edges)): + if face_edges[idx][0] != face_edges[idx - 1][1]: + # Swap the node index in this edge if not in counter-clockwise order + face_edges[idx] = face_edges[idx][::-1] + + # Fetch coordinates for each node in the face edges + cartesian_coordinates = np.array( + [ + [[node_x[node], node_y[node], node_z[node]] for node in edge] + for edge in face_edges + ] + ) + + return cartesian_coordinates + + def _get_lonlat_rad_face_edge_nodes_testcase_helper( + self,face_nodes_ind, face_edges_ind, edge_nodes_grid, node_lon, node_lat + ): + """This function is only used to help generating the testcase and + should not be used in the actual implementation. Construct an array to + hold the edge lat lon in radian connectivity for a face in a grid. + + Parameters + ---------- + face_nodes_ind : np.ndarray, shape (n_nodes,) + The ith entry of Grid.face_node_connectivity, where n_nodes is the number of nodes in the face. + face_edges_ind : np.ndarray, shape (n_edges,) + The ith entry of Grid.face_edge_connectivity, where n_edges is the number of edges in the face. + edge_nodes_grid : np.ndarray, shape (n_edges, 2) + The entire Grid.edge_node_connectivity, where n_edges is the total number of edges in the grid. + node_lon : np.ndarray, shape (n_nodes,) + The values of Grid.node_lon. + node_lat : np.ndarray, shape (n_nodes,) + The values of Grid.node_lat, where n_nodes_total is the total number of nodes in the grid. + Returns + ------- + face_edges : np.ndarray, shape (n_edges, 2, 3) + Face edge connectivity in Cartesian coordinates, where n_edges is the number of edges for the specific face. + + Notes + ----- + - The function assumes that the inputs are well-formed and correspond to the same face. + - The output array contains the Latitude and longitude coordinates in radian for each edge of the face. + """ + + # Create a mask that is True for all values not equal to INT_FILL_VALUE + mask = face_edges_ind != INT_FILL_VALUE + + # Use the mask to select only the elements not equal to INT_FILL_VALUE + valid_edges = face_edges_ind[mask] + face_edges = edge_nodes_grid[valid_edges] + + # Ensure counter-clockwise order of edge nodes + # Start with the first two nodes + face_edges[0] = [face_nodes_ind[0], face_nodes_ind[1]] + + for idx in range(1, len(face_edges)): + if face_edges[idx][0] != face_edges[idx - 1][1]: + # Swap the node index in this edge if not in counter-clockwise order + face_edges[idx] = face_edges[idx][::-1] + + # Fetch coordinates for each node in the face edges + lonlat_coordinates = np.array( + [ + [ + [ + np.mod(np.deg2rad(node_lon[node]), 2 * np.pi), + np.deg2rad(node_lat[node]), + ] + for node in edge + ] + for edge in face_edges + ] + ) + + return lonlat_coordinates def test_populate_bounds_normal(self): # Generate a normal face that is not crossing the antimeridian or the poles @@ -738,12 +1101,12 @@ def test_populate_bounds_normal(self): lon_min = np.deg2rad(10.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -766,12 +1129,12 @@ def test_populate_bounds_antimeridian(self): lon_min = np.deg2rad(350.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -794,12 +1157,12 @@ def test_populate_bounds_node_on_pole(self): lon_min = np.deg2rad(10.0) lon_max = np.deg2rad(50.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -822,12 +1185,12 @@ def test_populate_bounds_edge_over_pole(self): lon_min = np.deg2rad(210.0) lon_max = np.deg2rad(30.0) grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, @@ -850,12 +1213,12 @@ def test_populate_bounds_pole_inside(self): lon_min = 0 lon_max = 2 * np.pi grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_edges_connectivity_cartesian = self._get_cartesian_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_x.values, grid.node_y.values, grid.node_z.values) - face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_edges_connectivity_lonlat = self._get_lonlat_rad_face_edge_nodes_testcase_helper( grid.face_node_connectivity.values[0], grid.face_edge_connectivity.values[0], grid.edge_node_connectivity.values, grid.node_lon.values, diff --git a/test/test_grid.py b/test/test_grid.py index b9be85816..01d37a0ae 100644 --- a/test/test_grid.py +++ b/test/test_grid.py @@ -913,6 +913,7 @@ def test_from_face_vertices(self): class TestLatlonBounds(TestCase): + gridfile_mpas = current_path / "meshfiles" / "mpas" / "QU" / "oQU480.231010.nc" def test_populate_bounds_GCA_mix(self): face_1 = [[10.0, 60.0], [10.0, 10.0], [50.0, 10.0], [50.0, 60.0]] face_2 = [[350, 60.0], [350, 10.0], [50.0, 10.0], [50.0, 60.0]] @@ -932,3 +933,9 @@ def test_populate_bounds_GCA_mix(self): bounds_xarray = grid.bounds face_bounds = bounds_xarray.values nt.assert_allclose(grid.bounds.values, expected_bounds, atol=ERROR_TOLERANCE) + + def test_populate_bounds_MPAS(self): + xrds = xr.open_dataset(self.gridfile_mpas) + uxgrid = ux.Grid.from_dataset(xrds, use_dual=True) + bounds_xarray = uxgrid.bounds + pass diff --git a/test/test_helpers.py b/test/test_helpers.py index 42cd72def..202c51bfe 100644 --- a/test/test_helpers.py +++ b/test/test_helpers.py @@ -241,111 +241,303 @@ def test_angle_of_2_vectors(self): class TestFaceEdgeConnectivityHelper(TestCase): - def test_get_cartesian_face_edge_nodes(self): - # Load the dataset - - uxds = ux.open_dataset(gridfile_geoflowsmall_grid, - gridfile_geoflowsmall_var) + def test_get_cartesian_face_edge_nodes_pipeline(self): + # Create the vertices for the grid, based around the North Pole + vertices = [[0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], [0.5, -0.5, 0.5]] - # Initialize array + # Normalize the vertices + vertices = [x / np.linalg.norm(x) for x in vertices] - face_edges_connectivity_cartesian = [] + # Construct the grid from the vertices + grid = ux.Grid.from_face_vertices(vertices, latlon=False) - # Get the connectivity + # Extract the necessary grid data + face_node_conn = grid.face_node_connectivity.values + n_nodes_per_face = np.array([len(face) for face in face_node_conn]) + n_face = len(face_node_conn) - for i in range(len(uxds.uxgrid.face_node_connectivity)): - face_edges_connectivity_cartesian.append( - _get_cartesian_face_edge_nodes( - uxds.uxgrid.face_node_connectivity.values[i], - uxds.uxgrid.face_edge_connectivity.values[i], - uxds.uxgrid.edge_node_connectivity.values, - uxds.uxgrid.node_x.values, uxds.uxgrid.node_y.values, - uxds.uxgrid.node_z.values)) + n_max_face_edges = max(n_nodes_per_face) + node_x = grid.node_x.values + node_y = grid.node_y.values + node_z = grid.node_z.values - # Stack the arrays to get the desired (3,3) array + # Call the function to test + face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z + ) - face_edges_connectivity_cartesian = np.vstack( - face_edges_connectivity_cartesian) + # Check that the face_edges_connectivity_cartesian works as an input to _pole_point_inside_polygon + result = ux.grid.geometry._pole_point_inside_polygon( + 'North', face_edges_connectivity_cartesian[0] + ) - assert (face_edges_connectivity_cartesian.ndim == 3) + # Assert that the result is True + self.assertTrue(result) - def test_get_cartesian_face_edge_nodes_pipeline(self): + def test_get_cartesian_face_edge_nodes_filled_value(self): # Create the vertices for the grid, based around the North Pole + vertices = [[0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], [0.5, -0.5, 0.5]] - vertices = [[0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], - [0.5, -0.5, 0.5]] - - #Normalize the vertices + # Normalize the vertices vertices = [x / np.linalg.norm(x) for x in vertices] + vertices.append([INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]) # Construct the grid from the vertices grid = ux.Grid.from_face_vertices(vertices, latlon=False) + + # Extract the necessary grid data + face_node_conn = grid.face_node_connectivity.values + n_nodes_per_face = np.array([len(face) for face in face_node_conn]) + n_face = len(face_node_conn) + n_max_face_edges = max(n_nodes_per_face) + node_x = grid.node_x.values + node_y = grid.node_y.values + node_z = grid.node_z.values + + # Call the function to test face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( - grid.face_node_connectivity.values[0], - grid.face_edge_connectivity.values[0], - grid.edge_node_connectivity.values, grid.node_x.values, - grid.node_y.values, grid.node_z.values) + face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z + ) # Check that the face_edges_connectivity_cartesian works as an input to _pole_point_inside_polygon result = ux.grid.geometry._pole_point_inside_polygon( - 'North', face_edges_connectivity_cartesian) + 'North', face_edges_connectivity_cartesian[0] + ) # Assert that the result is True - self.assertTrue(result) - def test_get_cartesian_face_edge_nodes_filled_value(self): - # Create the vertices for the grid, based around the North Pole + def test_get_cartesian_face_edge_nodes_filled_value2(self): + # The face vertices order in counter-clockwise + # face_conn = [[0,1,2],[1,3,4,2]] + + #Each vertex is a 2D vector represent the longitude and latitude in degree. Call the node_lonlat_to_xyz to convert it to 3D vector + v0_deg = [10,10] + v1_deg = [15,15] + v2_deg = [5,15] + v3_deg = [15,45] + v4_deg = [5,45] + + # First convert them into radians + v0_rad = np.deg2rad(v0_deg) + v1_rad = np.deg2rad(v1_deg) + v2_rad = np.deg2rad(v2_deg) + v3_rad = np.deg2rad(v3_deg) + v4_rad = np.deg2rad(v4_deg) + + # It should look like following when passing in the _get_cartesian_face_edge_nodes + # [[v0_cart,v1_cart,v2_cart, [Fill_Value,Fill_Value,Fill_Value]],[v1_cart,v3_cart,v4_cart,v2_cart]] + v0_cart = _lonlat_rad_to_xyz(v0_rad[0],v0_rad[1]) + v1_cart = _lonlat_rad_to_xyz(v1_rad[0],v1_rad[1]) + v2_cart = _lonlat_rad_to_xyz(v2_rad[0],v2_rad[1]) + v3_cart = _lonlat_rad_to_xyz(v3_rad[0],v3_rad[1]) + v4_cart = _lonlat_rad_to_xyz(v4_rad[0],v4_rad[1]) + + face_node_conn = np.array([[0, 1, 2, INT_FILL_VALUE],[1, 3, 4, 2]]) + n_face = 2 + n_max_face_edges = 4 + n_nodes_per_face = np.array([len(face) for face in face_node_conn]) + node_x = np.array([v0_cart[0],v1_cart[0],v2_cart[0],v3_cart[0],v4_cart[0]]) + node_y = np.array([v0_cart[1],v1_cart[1],v2_cart[1],v3_cart[1],v4_cart[1]]) + node_z = np.array([v0_cart[2],v1_cart[2],v2_cart[2],v3_cart[2],v4_cart[2]]) + + # call the function to test + face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( + face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z + ) + + # Define correct result + correct_result = np.array([ + [ + [ + [v0_cart[0], v0_cart[1], v0_cart[2]], + [v1_cart[0], v1_cart[1], v1_cart[2]] + ], + [ + [v1_cart[0], v1_cart[1], v1_cart[2]], + [v2_cart[0], v2_cart[1], v2_cart[2]] + ], + [ + [v2_cart[0], v2_cart[1], v2_cart[2]], + [v0_cart[0], v0_cart[1], v0_cart[2]] + ], + [ + [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], + [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE] + ] + ], + [ + [ + [v1_cart[0], v1_cart[1], v1_cart[2]], + [v3_cart[0], v3_cart[1], v3_cart[2]] + ], + [ + [v3_cart[0], v3_cart[1], v3_cart[2]], + [v4_cart[0], v4_cart[1], v4_cart[2]] + ], + [ + [v4_cart[0], v4_cart[1], v4_cart[2]], + [v2_cart[0], v2_cart[1], v2_cart[2]] + ], + [ + [v2_cart[0], v2_cart[1], v2_cart[2]], + [v1_cart[0], v1_cart[1], v1_cart[2]] + ] + ] + ]) + - vertices = [[0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], - [0.5, -0.5, 0.5]] + # Assert that the result is correct + self.assertEqual(face_edges_connectivity_cartesian.shape, correct_result.shape) - #Normalize the vertices + + def test_get_lonlat_face_edge_nodes_pipeline(self): + # Create the vertices for the grid, based around the North Pole + vertices = [[0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], [0.5, -0.5, 0.5]] + + # Normalize the vertices vertices = [x / np.linalg.norm(x) for x in vertices] - vertices.append([INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]) # Construct the grid from the vertices grid = ux.Grid.from_face_vertices(vertices, latlon=False) - face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes( - grid.face_node_connectivity.values[0], - grid.face_edge_connectivity.values[0], - grid.edge_node_connectivity.values, grid.node_x.values, - grid.node_y.values, grid.node_z.values) + + # Extract the necessary grid data + face_node_conn = grid.face_node_connectivity.values + n_nodes_per_face = np.array([len(face) for face in face_node_conn]) + n_face = len(face_node_conn) + n_max_face_edges = max(n_nodes_per_face) + node_lon = grid.node_lon.values + node_lat = grid.node_lat.values + + # Call the function to test + face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_node_conn, n_face, n_max_face_edges, node_lon, node_lat + ) + + # Convert the first face's edges to Cartesian coordinates + face_edges_connectivity_lonlat = face_edges_connectivity_lonlat[0] + face_edges_connectivity_cartesian = [] + for edge in face_edges_connectivity_lonlat: + edge_cart = [_lonlat_rad_to_xyz(*node) for node in edge] + face_edges_connectivity_cartesian.append(edge_cart) # Check that the face_edges_connectivity_cartesian works as an input to _pole_point_inside_polygon result = ux.grid.geometry._pole_point_inside_polygon( - 'North', face_edges_connectivity_cartesian) + 'North', np.array(face_edges_connectivity_cartesian) + ) # Assert that the result is True self.assertTrue(result) - def test_get_lonlat_face_edge_nodes_pipeline(self): + def test_get_lonlat_face_edge_nodes_filled_value(self): # Create the vertices for the grid, based around the North Pole + vertices = [[0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], [0.5, -0.5, 0.5]] - vertices = [[0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], - [0.5, -0.5, 0.5]] - - #Normalize the vertices + # Normalize the vertices vertices = [x / np.linalg.norm(x) for x in vertices] + vertices.append([INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]) # Construct the grid from the vertices grid = ux.Grid.from_face_vertices(vertices, latlon=False) + + # Extract the necessary grid data + face_node_conn = grid.face_node_connectivity.values + n_nodes_per_face = np.array([len(face) for face in face_node_conn]) + n_face = len(face_node_conn) + n_max_face_edges = max(n_nodes_per_face) + node_lon = grid.node_lon.values + node_lat = grid.node_lat.values + + # Call the function to test face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( - grid.face_node_connectivity.values[0], - grid.face_edge_connectivity.values[0], - grid.edge_node_connectivity.values, grid.node_lon.values, - grid.node_lat.values) + face_node_conn, n_face, n_max_face_edges, node_lon, node_lat + ) - # Convert all the values into cartesian coordinates + # Convert the first face's edges to Cartesian coordinates + face_edges_connectivity_lonlat = face_edges_connectivity_lonlat[0] face_edges_connectivity_cartesian = [] - for i in range(len(face_edges_connectivity_lonlat)): - edge = face_edges_connectivity_lonlat[i] + for edge in face_edges_connectivity_lonlat: edge_cart = [_lonlat_rad_to_xyz(*node) for node in edge] face_edges_connectivity_cartesian.append(edge_cart) # Check that the face_edges_connectivity_cartesian works as an input to _pole_point_inside_polygon result = ux.grid.geometry._pole_point_inside_polygon( - 'North', np.array(face_edges_connectivity_cartesian)) + 'North', np.array(face_edges_connectivity_cartesian) + ) # Assert that the result is True self.assertTrue(result) + + + def test_get_lonlat_face_edge_nodes_filled_value2(self): + # The face vertices order in counter-clockwise + # face_conn = [[0,1,2],[1,3,4,2]] + + #Each vertex is a 2D vector represent the longitude and latitude in degree. Call the node_lonlat_to_xyz to convert it to 3D vector + v0_deg = [10,10] + v1_deg = [15,15] + v2_deg = [5,15] + v3_deg = [15,45] + v4_deg = [5,45] + + # First convert them into radians + v0_rad = np.deg2rad(v0_deg) + v1_rad = np.deg2rad(v1_deg) + v2_rad = np.deg2rad(v2_deg) + v3_rad = np.deg2rad(v3_deg) + v4_rad = np.deg2rad(v4_deg) + + face_node_conn = np.array([[0, 1, 2, INT_FILL_VALUE],[1, 3, 4, 2]]) + n_face = 2 + n_max_face_edges = 4 + n_nodes_per_face = np.array([len(face) for face in face_node_conn]) + node_lon = np.array([v0_rad[0],v1_rad[0],v2_rad[0],v3_rad[0],v4_rad[0]]) + node_lat = np.array([v0_rad[1],v1_rad[1],v2_rad[1],v3_rad[1],v4_rad[1]]) + + # call the function to test + face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes( + face_node_conn, n_face, n_max_face_edges, node_lon, node_lat + ) + + # Define correct result + correct_result = np.array([ + [ + [ + [v0_rad[0], v0_rad[1]], + [v1_rad[0], v1_rad[1]] + ], + [ + [v1_rad[0], v1_rad[1]], + [v2_rad[0], v2_rad[1]] + ], + [ + [v2_rad[0], v2_rad[1]], + [v0_rad[0], v0_rad[1]] + ], + [ + [INT_FILL_VALUE, INT_FILL_VALUE], + [INT_FILL_VALUE, INT_FILL_VALUE] + ] + ], + [ + [ + [v1_rad[0], v1_rad[1]], + [v3_rad[0], v3_rad[1]] + ], + [ + [v3_rad[0], v3_rad[1]], + [v4_rad[0], v4_rad[1]] + ], + [ + [v4_rad[0], v4_rad[1]], + [v2_rad[0], v2_rad[1]] + ], + [ + [v2_rad[0], v2_rad[1]], + [v1_rad[0], v1_rad[1]] + ] + ] + ]) + + # Assert that the result is correct + self.assertEqual(face_edges_connectivity_lonlat.shape, correct_result.shape) diff --git a/uxarray/grid/geometry.py b/uxarray/grid/geometry.py index fe146990c..55b7d1eac 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -881,23 +881,35 @@ def _populate_bounds( intervals_tuple_list = [] intervals_name_list = [] + faces_edges_cartesian = _get_cartesian_face_edge_nodes( + grid.face_node_connectivity.values, + grid.n_face, + grid.n_max_face_edges, + grid.node_x.values, + grid.node_y.values, + grid.node_z.values, + ) + + faces_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes( + grid.face_node_connectivity.values, + grid.n_face, + grid.n_max_face_edges, + grid.node_lon.values, + grid.node_lat.values, + ) + for face_idx, face_nodes in enumerate(grid.face_node_connectivity): - face_edges_cartesian = _get_cartesian_face_edge_nodes( - grid.face_node_connectivity.values[face_idx], - grid.face_edge_connectivity.values[face_idx], - grid.edge_node_connectivity.values, - grid.node_x.values, - grid.node_y.values, - grid.node_z.values, - ) + face_edges_cartesian = faces_edges_cartesian[face_idx] - face_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes( - grid.face_node_connectivity.values[face_idx], - grid.face_edge_connectivity.values[face_idx], - grid.edge_node_connectivity.values, - grid.node_lon.values, - grid.node_lat.values, - ) + # Skip processing if the face is a dummy face + if np.any(face_edges_cartesian == INT_FILL_VALUE): + continue + + face_edges_lonlat_rad = faces_edges_lonlat_rad[face_idx] + + # Skip processing if the face is a dummy face + if np.any(face_edges_lonlat_rad == INT_FILL_VALUE): + continue is_GCA_list = ( is_face_GCA_list[face_idx] if is_face_GCA_list is not None else None diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index a488584d2..020aae1d9 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -189,123 +189,237 @@ def _newton_raphson_solver_for_gca_constLat( return np.append(y_new, constZ) -# TODO: Consider re-implementation in the future / better integration with API +def _swap_first_fill_value_with_last(arr): + """Swap the first occurrence of INT_FILL_VALUE in each sub-array with the + last value in the sub-array. + + Parameters: + ---------- + arr (np.ndarray): A 3D numpy array where the swap will be performed. + + Returns: + ------- + np.ndarray: The modified array with the swaps made. + """ + # Find the indices of the first INT_FILL_VALUE in each sub-array + mask = arr == INT_FILL_VALUE + reshaped_mask = mask.reshape(arr.shape[0], -1) + first_true_indices = np.argmax(reshaped_mask, axis=1) + + # If no INT_FILL_VALUE is found in a row, argmax will return 0, we need to handle this case + first_true_indices[~np.any(reshaped_mask, axis=1)] = -1 + + # Get the shape of the sub-arrays + subarray_shape = arr.shape[1:] + + # Calculate the 2D indices within each sub-array + valid_indices = first_true_indices != -1 + first_true_positions = np.unravel_index( + first_true_indices[valid_indices], subarray_shape + ) + + # Create an index array for the last value in each sub-array + last_indices = np.full((arr.shape[0],), subarray_shape[0] * subarray_shape[1] - 1) + last_positions = np.unravel_index(last_indices, subarray_shape) + + # Swap the first INT_FILL_VALUE with the last value in each sub-array + row_indices = np.arange(arr.shape[0]) + + # Advanced indexing to swap values + ( + arr[ + row_indices[valid_indices], first_true_positions[0], first_true_positions[1] + ], + arr[ + row_indices[valid_indices], + last_positions[0][valid_indices], + last_positions[1][valid_indices], + ], + ) = ( + arr[ + row_indices[valid_indices], + last_positions[0][valid_indices], + last_positions[1][valid_indices], + ], + arr[ + row_indices[valid_indices], first_true_positions[0], first_true_positions[1] + ], + ) + + return arr + + def _get_cartesian_face_edge_nodes( - face_nodes_ind, face_edges_ind, edge_nodes_grid, node_x, node_y, node_z + face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z ): """Construct an array to hold the edge Cartesian coordinates connectivity - for a face in a grid. + for multiple faces in a grid. Parameters ---------- - face_nodes_ind : np.ndarray, shape (n_nodes,) - The ith entry of Grid.face_node_connectivity, where n_nodes is the number of nodes in the face. - face_edges_ind : np.ndarray, shape (n_edges,) - The ith entry of Grid.face_edge_connectivity, where n_edges is the number of edges in the face. - edge_nodes_grid : np.ndarray, shape (n_edges, 2) - The entire Grid.edge_node_connectivity, where n_edges is the total number of edges in the grid. - node_x : np.ndarray, shape (n_nodes,) - The values of Grid.node_x, where n_nodes_total is the total number of nodes in the grid. - node_y : np.ndarray, shape (n_nodes,) - The values of Grid.node_y. - node_z : np.ndarray, shape (n_nodes,) - The values of Grid.node_z. + face_node_conn : np.ndarray + An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`. + n_face : int + The number of faces in the grid. Accessed through `grid.n_face`. + n_max_face_edges : int + The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`. + node_x : np.ndarray + An array of shape (n_nodes,) containing the x-coordinate values of the nodes. Accessed through `grid.node_x`. + node_y : np.ndarray + An array of shape (n_nodes,) containing the y-coordinate values of the nodes. Accessed through `grid.node_y`. + node_z : np.ndarray + An array of shape (n_nodes,) containing the z-coordinate values of the nodes. Accessed through `grid.node_z`. Returns ------- - face_edges : np.ndarray, shape (n_edges, 2, 3) - Face edge connectivity in Cartesian coordinates, where n_edges is the number of edges for the specific face. + face_edges_cartesian : np.ndarray + An array of shape (n_face, n_max_face_edges, 2, 3) containing the Cartesian coordinates of the edges + for each face. It might contain dummy values if the grid has holes. - Notes - ----- - - The function assumes that the inputs are well-formed and correspond to the same face. - - The output array contains the Cartesian coordinates for each edge of the face. + Examples + -------- + >>> face_node_conn = np.array([[0, 1, 2, 3, 4], [0, 1, 3, 4, INT_FILL_VALUE], [0, 1, 3, INT_FILL_VALUE, INT_FILL_VALUE]]) + >>> n_face = 3 + >>> n_max_face_edges = 5 + >>> node_x = np.array([0, 1, 1, 0, 1, 0]) + >>> node_y = np.array([0, 0, 1, 1, 2, 2]) + >>> node_z = np.array([0, 0, 0, 0, 1, 1]) + >>> _get_cartesian_face_edge_nodes(face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z) + array([[[[ 0, 0, 0], + [ 1, 0, 0]], + + [[ 1, 0, 0], + [ 1, 1, 0]], + + [[ 1, 1, 0], + [ 0, 1, 0]], + + [[ 0, 1, 0], + [ 1, 2, 1]], + + [[ 1, 2, 1], + [ 0, 0, 0]]], + + + [[[ 0, 0, 0], + [ 1, 0, 0]], + + [[ 1, 0, 0], + [ 0, 1, 0]], + + [[ 0, 1, 0], + [ 1, 2, 1]], + + [[ 1, 2, 1], + [ 0, 0, 0]], + + [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], + [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]], + + + [[[ 0, 0, 0], + [ 1, 0, 0]], + + [[ 1, 0, 0], + [ 0, 1, 0]], + + [[ 0, 1, 0], + [ 0, 0, 0]], + + [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], + [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]], + + [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], + [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]]]) """ + # Shift node connections to create edge connections + face_node_conn_shift = np.roll(face_node_conn, -1, axis=1) + + # Construct edge connections by combining original and shifted node connections + face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1) - # Create a mask that is True for all values not equal to INT_FILL_VALUE - mask = face_edges_ind != INT_FILL_VALUE + # swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array + face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn) - # Use the mask to select only the elements not equal to INT_FILL_VALUE - valid_edges = face_edges_ind[mask] - face_edges = edge_nodes_grid[valid_edges] + # Get the indices of the nodes from face_edge_conn + face_edge_conn_flat = face_edge_conn.reshape(-1) - # Ensure counter-clockwise order of edge nodes - # Start with the first two nodes - face_edges[0] = [face_nodes_ind[0], face_nodes_ind[1]] + valid_mask = face_edge_conn_flat != INT_FILL_VALUE - for idx in range(1, len(face_edges)): - if face_edges[idx][0] != face_edges[idx - 1][1]: - # Swap the node index in this edge if not in counter-clockwise order - face_edges[idx] = face_edges[idx][::-1] + # Get the valid node indices + valid_edges = face_edge_conn_flat[valid_mask] - # Fetch coordinates for each node in the face edges - cartesian_coordinates = np.array( - [ - [[node_x[node], node_y[node], node_z[node]] for node in edge] - for edge in face_edges - ] + # Create an array to hold the Cartesian coordinates of the edges + face_edges_cartesian = np.full( + (len(face_edge_conn_flat), 3), INT_FILL_VALUE, dtype=float ) - return cartesian_coordinates + # Fill the array with the Cartesian coordinates of the edges + face_edges_cartesian[valid_mask, 0] = node_x[valid_edges] + face_edges_cartesian[valid_mask, 1] = node_y[valid_edges] + face_edges_cartesian[valid_mask, 2] = node_z[valid_edges] + + return face_edges_cartesian.reshape(n_face, n_max_face_edges, 2, 3) def _get_lonlat_rad_face_edge_nodes( - face_nodes_ind, face_edges_ind, edge_nodes_grid, node_lon, node_lat + face_node_conn, n_face, n_max_face_edges, node_lon, node_lat ): - """Construct an array to hold the edge lat lon in radian connectivity for a - face in a grid. + """Construct an array to hold the edge latitude and longitude in radians + connectivity for multiple faces in a grid. Parameters ---------- - face_nodes_ind : np.ndarray, shape (n_nodes,) - The ith entry of Grid.face_node_connectivity, where n_nodes is the number of nodes in the face. - face_edges_ind : np.ndarray, shape (n_edges,) - The ith entry of Grid.face_edge_connectivity, where n_edges is the number of edges in the face. - edge_nodes_grid : np.ndarray, shape (n_edges, 2) - The entire Grid.edge_node_connectivity, where n_edges is the total number of edges in the grid. - node_lon : np.ndarray, shape (n_nodes,) - The values of Grid.node_lon. - node_lat : np.ndarray, shape (n_nodes,) - The values of Grid.node_lat, where n_nodes_total is the total number of nodes in the grid. + face_node_conn : np.ndarray + An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`. + n_face : int + The number of faces in the grid. Accessed through `grid.n_face`. + n_max_face_edges : int + The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`. + node_lon : np.ndarray + An array of shape (n_nodes,) containing the longitude values of the nodes in degrees. Accessed through `grid.node_lon`. + node_lat : np.ndarray + An array of shape (n_nodes,) containing the latitude values of the nodes in degrees. Accessed through `grid.node_lat`. + Returns ------- - face_edges : np.ndarray, shape (n_edges, 2, 3) - Face edge connectivity in Cartesian coordinates, where n_edges is the number of edges for the specific face. + face_edges_lonlat_rad : np.ndarray + An array of shape (n_face, n_max_face_edges, 2, 2) containing the latitude and longitude coordinates + in radians for the edges of each face. It might contain dummy values if the grid has holes. Notes ----- - - The function assumes that the inputs are well-formed and correspond to the same face. - - The output array contains the Latitude and longitude coordinates in radian for each edge of the face. + If the grid has holes, the function will return an entry of dummy value faces_edges_coordinates[i] filled with INT_FILL_VALUE. """ - # Create a mask that is True for all values not equal to INT_FILL_VALUE - mask = face_edges_ind != INT_FILL_VALUE + # Convert node coordinates to radians + node_lon_rad = np.deg2rad(node_lon) + node_lat_rad = np.deg2rad(node_lat) - # Use the mask to select only the elements not equal to INT_FILL_VALUE - valid_edges = face_edges_ind[mask] - face_edges = edge_nodes_grid[valid_edges] + # Shift node connections to create edge connections + face_node_conn_shift = np.roll(face_node_conn, -1, axis=1) - # Ensure counter-clockwise order of edge nodes - # Start with the first two nodes - face_edges[0] = [face_nodes_ind[0], face_nodes_ind[1]] + # Construct edge connections by combining original and shifted node connections + face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1) - for idx in range(1, len(face_edges)): - if face_edges[idx][0] != face_edges[idx - 1][1]: - # Swap the node index in this edge if not in counter-clockwise order - face_edges[idx] = face_edges[idx][::-1] + # swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array + face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn) - # Fetch coordinates for each node in the face edges - lonlat_coordinates = np.array( - [ - [ - [ - np.mod(np.deg2rad(node_lon[node]), 2 * np.pi), - np.deg2rad(node_lat[node]), - ] - for node in edge - ] - for edge in face_edges - ] + # Get the indices of the nodes from face_edge_conn + face_edge_conn_flat = face_edge_conn.reshape(-1) + + valid_mask = face_edge_conn_flat != INT_FILL_VALUE + + # Get the valid node indices + valid_edges = face_edge_conn_flat[valid_mask] + + # Create an array to hold the latitude and longitude in radians for the edges + face_edges_lonlat_rad = np.full( + (len(face_edge_conn_flat), 2), INT_FILL_VALUE, dtype=float ) - return lonlat_coordinates + # Fill the array with the latitude and longitude in radians for the edges + face_edges_lonlat_rad[valid_mask, 0] = node_lon_rad[valid_edges] + face_edges_lonlat_rad[valid_mask, 1] = node_lat_rad[valid_edges] + + return face_edges_lonlat_rad.reshape(n_face, n_max_face_edges, 2, 2)