Implements Grid.compute_face_node_angles()

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

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)

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

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",

uxarray/grid/utils.py

@@ -12,24 +12,63 @@ 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
     -------
     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):
     """