diff --git a/docs/api.rst b/docs/api.rst index 2addaabd8..bc92a7513 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -192,6 +192,7 @@ Methods Grid.copy Grid.calculate_total_face_area Grid.compute_face_areas + Grid.compute_face_node_angles Grid.construct_face_centers Grid.get_ball_tree Grid.get_kd_tree diff --git a/test/grid/geometry/test_angles.py b/test/grid/geometry/test_angles.py new file mode 100644 index 000000000..f818ac576 --- /dev/null +++ b/test/grid/geometry/test_angles.py @@ -0,0 +1,67 @@ +""" +Purpose: tests related to angle calculations on a grid +""" + +import numpy as np +import xarray as xr + +import uxarray as ux + + +def test_face_node_angles_triangle(): + """ensure Grid.compute_face_node_angles() works as expected for a simple ~30,60,90 triangle.""" + # make a tiny triangle with known angles (90,60,30 degrees): + # (n1) + # | %% + # | %% + # (n0) ------ (n2) + node_lon = [0, 0, np.sqrt(3)] + node_lat = [0, 1, 0] + face_node_connectivity = [[0, 1, 2]] + grid = ux.Grid.from_topology(node_lon, node_lat, face_node_connectivity) + angles_rad = grid.compute_face_node_angles() + angles_deg = grid.compute_face_node_angles(degrees=True) + assert np.allclose(np.rad2deg(angles_rad), angles_deg) + angles_uxarr = grid.compute_face_node_angles(as_uxarray=True) + assert isinstance(angles_rad, xr.DataArray) + assert isinstance(angles_uxarr, ux.UxDataArray) + assert np.all(angles_rad == angles_uxarr) + angle_at_n0 = angles_deg.isel(n_face=0, n_max_face_nodes=0) + angle_at_n1 = angles_deg.isel(n_face=0, n_max_face_nodes=1) + angle_at_n2 = angles_deg.isel(n_face=0, n_max_face_nodes=2) + assert angle_at_n0 == 90 # this turns out to be exact... from arctan2(any_value, 0) + assert np.abs(angle_at_n1 - 60) < 1e-2 # basically 60 degrees + assert np.abs(angle_at_n2 - 30) < 1e-2 # basically 30 degrees + # on a unit sphere, spherical excess == face area, via Girard's theorem. + spherical_excess = angles_rad.sum('n_max_face_nodes') - np.pi + face_areas = grid.compute_face_areas() + assert np.allclose(spherical_excess, face_areas, atol=0, rtol=1e-12) + +def test_face_node_angles_hexagons_and_pentagons(): + """ensure face_node_angles works as expected on grids with hexagons and pentagons""" + grid = ux.tutorial.open_grid('quad-hexagon') # has multiple faces, all hexagons. + angles_deg = grid.compute_face_node_angles(degrees=True) + # every hexagon in this grid is close to regular (all 120 degree angles): + regular_hex_deviation = angles_deg - 120 + assert np.max(np.abs(regular_hex_deviation)) < 4.0 + # generalized spherical excess formula uses (n - 2) * np.pi; n==6 for all of these faces + angles = grid.compute_face_node_angles() # (need to use radians for this formula) + spherical_excess = angles.sum('n_max_face_nodes') - (6 - 2) * np.pi + face_areas = grid.compute_face_areas() + assert np.allclose(spherical_excess, face_areas, atol=0, rtol=1e-10) + + # now test a grid which has pentagons too, + # to ensure the implementation works even when the number of nodes per face varies. + grid = ux.tutorial.open_grid('mpas-QU-480') + # not all close to regular so don't try to check that. + # ensure nan values wherever n_max_face_nodes dimension is larger than n_nodes_per_face + angles = grid.compute_face_node_angles() + assert not np.all(grid.n_nodes_per_face == grid.n_max_face_nodes) + should_have_nans = angles.where(grid.n_nodes_per_face < grid.n_max_face_nodes, drop=True) + assert should_have_nans.size > 0 + should_be_nans = should_have_nans.isel(n_max_face_nodes = -1) + assert np.all(np.isnan(should_be_nans)) + # generalized spherical excess formula uses (n_nodes_per_face - 2) * np.pi + spherical_excess = angles.sum('n_max_face_nodes') - (grid.n_nodes_per_face - 2) * np.pi + face_areas = grid.compute_face_areas() + assert np.allclose(spherical_excess, face_areas, atol=0, rtol=1e-9) diff --git a/uxarray/grid/angles.py b/uxarray/grid/angles.py new file mode 100644 index 000000000..678f94cb4 --- /dev/null +++ b/uxarray/grid/angles.py @@ -0,0 +1,77 @@ +""" +Purpose: angle calculations on a grid +""" + +import numpy as np +from numba import njit, prange + +from uxarray.grid.utils import _numba_norm3, _small_angle_of_2_vectors + + +@njit(cache=True, parallel=True) +def _compute_face_node_angles_convex( + node_x, + node_y, + node_z, + face_node_connectivity, + n_nodes_per_face, +): + """ + Calculate the angles at each node for each face, assuming convex faces + and a spherical geometry (these assumptions occur throughout uxarray). + + Parameters + ---------- + node_x : np.ndarray with shape (n_nodes,) + X coordinates of the nodes. + node_y : np.ndarray with shape (n_nodes,) + Y coordinates of the nodes. + node_z : np.ndarray with shape (n_nodes,) + Z coordinates of the nodes. + face_node_connectivity : np.ndarray with shape (n_faces, n_max_face_nodes) + Connectivity array defining which nodes form each face. + n_nodes_per_face : np.ndarray with shape (n_faces,) + Number of nodes for each face. + + Returns + ------- + np.ndarray with shape (n_faces, n_max_face_nodes) + Angles at each node of each face. + INT_FILL_VALUE elements from face_node_connectivity correspond with np.nan in the result. + """ + n_faces, _n_max_face_nodes = face_node_connectivity.shape + result = np.full(face_node_connectivity.shape, np.nan, dtype=np.float64) + for i in prange(n_faces): + n_nodes = n_nodes_per_face[i] + for j in range(n_nodes): + ihere = face_node_connectivity[i, j] + iprev = face_node_connectivity[i, (j - 1) % n_nodes] + inext = face_node_connectivity[i, (j + 1) % n_nodes] + xhere = node_x[ihere] + yhere = node_y[ihere] + zhere = node_z[ihere] + v1 = (node_x[iprev] - xhere, node_y[iprev] - yhere, node_z[iprev] - zhere) + v2 = (node_x[inext] - xhere, node_y[inext] - yhere, node_z[inext] - zhere) + # Spherical geometry: project onto tangent plane at the current node + normal = (xhere, yhere, zhere) + normal_norm = _numba_norm3(normal) # |normal| + normal = ( + normal[0] / normal_norm, + normal[1] / normal_norm, + normal[2] / normal_norm, + ) + # v1 -= np.dot(v1, normal) * normal + v1_dot_normal = v1[0] * normal[0] + v1[1] * normal[1] + v1[2] * normal[2] + v2_dot_normal = v2[0] * normal[0] + v2[1] * normal[1] + v2[2] * normal[2] + v1 = ( + v1[0] - v1_dot_normal * normal[0], + v1[1] - v1_dot_normal * normal[1], + v1[2] - v1_dot_normal * normal[2], + ) + v2 = ( + v2[0] - v2_dot_normal * normal[0], + v2[1] - v2_dot_normal * normal[1], + v2[2] - v2_dot_normal * normal[2], + ) + result[i, j] = _small_angle_of_2_vectors(v1, v2) + return result diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 2f1bb582d..f72aae45f 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -19,6 +19,7 @@ from uxarray.cross_sections import GridCrossSectionAccessor from uxarray.errors import DataCenteringError, DimensionError, GridInvalidError from uxarray.formatting_html import grid_repr +from uxarray.grid.angles import _compute_face_node_angles_convex from uxarray.grid.area import _get_all_face_area_from_coords from uxarray.grid.bounds import _populate_face_bounds from uxarray.grid.connectivity import ( @@ -1972,6 +1973,53 @@ def copy(self): source_dims_dict=self._source_dims_dict, ) + def compute_face_node_angles( + self, + *, + degrees: bool = False, + as_uxarray: bool = False, + ) -> xr.DataArray | UxDataArray: + """Compute the angles at each node of each face in the grid. + Assumes convex faces and a spherical geometry (consistent with other uxarray methods). + + Parameters + ---------- + degrees : bool, defaults to False + Whether to return angles in degrees (if True) or radians (if False). + as_uxarray : bool, defaults to False + Whether to return a uxarray.DataArray (if True) instead of an xarray.DataArray (if False). + If True, equivalent to uxarray.DataArray(self.compute_face_node_angles(..., as_uxarray=False), uxgrid=self). + + Returns + ------- + face_node_angles : xr.DataArray or uxarray.UxDataArray (if as_uxarray=True) + The internal angles at each node, for each face in the grid. + Has 'n_face' and 'n_max_face_nodes' dimensions, with same size as in self. + For faces with fewer than n_max_face_nodes, fill value is np.nan. + """ + from uxarray.conventions.ugrid import FACE_DIM, N_MAX_FACE_NODES_DIM + + result = _compute_face_node_angles_convex( + self.node_x.values, + self.node_y.values, + self.node_z.values, + self.face_node_connectivity.values, + self.n_nodes_per_face.values, + ) + result = xr.DataArray( + data=result, + dims=[FACE_DIM, N_MAX_FACE_NODES_DIM], + name="face_node_angles", + attrs={"description": "Internal angles at each node of each face."}, + ) + if degrees: + result = np.rad2deg(result) + if as_uxarray: + from uxarray.core.dataarray import UxDataArray + + result = UxDataArray(result, uxgrid=self) + return result + def calculate_total_face_area( self, quadrature_rule: str = "triangular", diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index d70962273..6169d5892 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -12,9 +12,9 @@ def _small_angle_of_2_vectors(u, v): Parameters ---------- - u : numpy.ndarray + u : numpy.ndarray or iterable of length 3 The first 3D vector. - v : numpy.ndarray + v : numpy.ndarray or iterable of length 3 The second 3D vector. Returns @@ -22,14 +22,53 @@ def _small_angle_of_2_vectors(u, v): float The smallest angle between `u` and `v` in radians. """ - v_norm_times_u = np.linalg.norm(v) * u - u_norm_times_v = np.linalg.norm(u) * v - vec_minus = v_norm_times_u - u_norm_times_v - vec_sum = v_norm_times_u + u_norm_times_v - angle_u_v_rad = 2 * np.arctan2(np.linalg.norm(vec_minus), np.linalg.norm(vec_sum)) + # don't convert to numpy array if not already numpy array. + # The formula is: angle = 2 * arctan2(| |v|*u - |u|*v |, | |v|*u + |u|*v |) + v_norm = _numba_norm3(v) + u_norm = _numba_norm3(u) + v_norm_times_u = (v_norm * u[0], v_norm * u[1], v_norm * u[2]) + u_norm_times_v = (u_norm * v[0], u_norm * v[1], u_norm * v[2]) + vec_minus = ( + v_norm_times_u[0] - u_norm_times_v[0], + v_norm_times_u[1] - u_norm_times_v[1], + v_norm_times_u[2] - u_norm_times_v[2], + ) + vec_sum = ( + v_norm_times_u[0] + u_norm_times_v[0], + v_norm_times_u[1] + u_norm_times_v[1], + v_norm_times_u[2] + u_norm_times_v[2], + ) + norm_vec_minus = _numba_norm3(vec_minus) + norm_vec_sum = _numba_norm3(vec_sum) + angle_u_v_rad = 2 * np.arctan2(norm_vec_minus, norm_vec_sum) return angle_u_v_rad +# TODO: move _numba_norm3 to a higher-level utils file. For more details, see issue #1648. +@njit(cache=True) +def _numba_norm3(u): + """ + Compute the Euclidean norm of a 3D vector. + Implementation is currently equivalent to np.linalg.norm: + sqrt(u[0]**2 + u[1]**2 + u[2]**2) + + Does NOT internally convert u to a list or numpy array; + utilizing tuples in numba instead of many tiny lists/arrays + can improve performance significantly. + + Parameters + ---------- + u : iterable of length 3, possibly a numpy array + The 3D vector. + + Returns + ------- + float + The Euclidean norm of the vector `u`. + """ + return (u[0] ** 2 + u[1] ** 2 + u[2] ** 2) ** 0.5 + + @njit(cache=True) def _angle_of_2_vectors(u, v): """