From 85ec085765e7dadbb4701c355d76a9280e36a5e6 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Mon, 10 Aug 2026 13:58:16 -0500 Subject: [PATCH 1/3] Normalize the gradient by the dual cell and add the spherical metric terms Green-Gauss on a contour must divide by the area that contour encloses. The face gradient integrates around the dual cell but divided by the primal face area, inflating results by A_dual/A_primal: about 4 on quadrilateral meshes and 3 on hexagonal ones. Refining healpix z3 to z6 converged to 3.9278, 3.9811, 3.9952, 3.9988, so this was a normalization error rather than truncation. Curl and divergence used the planar formulas and dropped the u*tan(lat)/a and -v*tan(lat)/a metric terms. On solid-body rotation that costs exactly a factor of two, which partly cancelled the area inflation and made curl look like a clean 2x on quads but 1.5x on hexagons. Both corrections have to land together: fixing the area alone drives curl to 0.5. Faces whose contour is left open by a missing edge neighbor now return NaN, since no enclosed area exists to divide by. Interior faces of SCRIP-derived grids stay finite, so #1452 remains fixed. Adds manufactured-solution tests with non-zero closed-form answers. The previous suite relied on null tests and ordering comparisons, both of which are satisfied by an operator that is off by a constant factor. Fixes #1662 --- test/core/test_vector_calculus.py | 191 +++++++++++++++++++++++++++--- uxarray/core/dataarray.py | 27 ++++- uxarray/core/gradient.py | 100 +++++++++++++--- 3 files changed, 281 insertions(+), 37 deletions(-) diff --git a/test/core/test_vector_calculus.py b/test/core/test_vector_calculus.py index aad4c78bc..a3b9dc1f6 100644 --- a/test/core/test_vector_calculus.py +++ b/test/core/test_vector_calculus.py @@ -25,15 +25,31 @@ def test_gradient_output_format(self, gridpath, datasetpath): def test_gradient_all_boundary_faces(self, gridpath, datasetpath): """Quad hexagon grid has 4 faces, all touching the boundary. - Each face still has some interior edges, so the gradient should - produce finite (small) values rather than NaN. + Green-Gauss needs a closed contour to divide by the area it + encloses. Every face here sits on the mesh boundary, so no closed + contour exists and NaN is the honest answer rather than a value + normalized by an area that was never integrated over. """ uxds = ux.open_dataset(gridpath("ugrid", "quad-hexagon", "grid.nc"), datasetpath("ugrid", "quad-hexagon", "data.nc")) grad = uxds['t2m'].gradient() - assert not np.isnan(grad['meridional_gradient']).any() - assert not np.isnan(grad['zonal_gradient']).any() + assert np.isnan(grad['meridional_gradient']).all() + assert np.isnan(grad['zonal_gradient']).all() + + def test_gradient_partial_boundary_still_finite(self, gridpath): + """Interior faces of a SCRIP-derived grid keep finite gradients. + + Regression guard for #1452: the boundary handling must not go back + to NaN-ing an entire grid. + """ + uxgrid = ux.open_grid(gridpath("scrip", "ne30pg2", "grid.nc")) + lat = np.deg2rad(uxgrid.face_lat.values) + phi = ux.UxDataArray(np.sin(lat), dims=["n_face"], uxgrid=uxgrid, name="phi") + + grad = phi.gradient(scale_by_radius=False) + + assert np.isfinite(grad["meridional_gradient"].values).mean() > 0.95 class TestGradientMPASOcean: @@ -297,7 +313,12 @@ def test_scalardotgradient_rejects_misaligned_indexes(self, gridpath, datasetpat class TestDivergenceDyamondSubset: def test_divergence_constant_field(self, gridpath, datasetpath): - """Test divergence of constant vector field (should be zero)""" + """Divergence of a constant vector field reduces to the metric term. + + On a plane this would be zero, but on the sphere the divergence + carries -v*tan(lat)/a, so a constant v = 1 leaves exactly + -tan(lat)/a behind. + """ uxds = ux.open_dataset( gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"), datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc") @@ -309,15 +330,16 @@ def test_divergence_constant_field(self, gridpath, datasetpath): div_field = constant_u.divergence(constant_v) - # Divergence of constant field should be close to zero for interior faces # Boundary faces may have NaN values (which is expected) - finite_values = div_field.values[np.isfinite(div_field.values)] + finite = np.isfinite(div_field.values) + assert finite.any(), "No finite divergence values found" - # Check that we have some finite values (interior faces) - assert len(finite_values) > 0, "No finite divergence values found" + radius = uxds.uxgrid._ds.attrs["sphere_radius"] + expected = -np.tan(np.deg2rad(uxds.uxgrid.face_lat.values)) / radius - # Divergence of constant field should be close to zero for finite values - assert np.abs(finite_values).max() < 1e-10, f"Max divergence: {np.abs(finite_values).max()}" + nt.assert_allclose( + div_field.values[finite], expected[finite], rtol=1e-10, atol=1e-15 + ) def test_divergence_linear_field(self, gridpath, datasetpath): """Test divergence of linear vector field""" @@ -485,7 +507,11 @@ def test_curl_basic(self, gridpath, datasetpath): class TestCurlDyamondSubset: def test_curl_constant_field(self, gridpath, datasetpath): - """Test curl of constant vector field (should be zero)""" + """Curl of a constant vector field reduces to the metric term. + + On a plane this would be zero, but on the sphere the curl carries + u*tan(lat)/a, so a constant u = 1 leaves exactly tan(lat)/a behind. + """ uxds = ux.open_dataset( gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"), datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc") @@ -497,15 +523,16 @@ def test_curl_constant_field(self, gridpath, datasetpath): curl_field = constant_u.curl(constant_v) - # Curl of constant field should be close to zero for interior faces # Boundary faces may have NaN values (which is expected) - finite_values = curl_field.values[np.isfinite(curl_field.values)] + finite = np.isfinite(curl_field.values) + assert finite.any(), "No finite curl values found" - # Check that we have some finite values (interior faces) - assert len(finite_values) > 0, "No finite curl values found" + radius = uxds.uxgrid._ds.attrs["sphere_radius"] + expected = np.tan(np.deg2rad(uxds.uxgrid.face_lat.values)) / radius - # Curl of constant field should be close to zero for finite values - assert np.abs(finite_values).max() < 1e-10, f"Max curl: {np.abs(finite_values).max()}" + nt.assert_allclose( + curl_field.values[finite], expected[finite], rtol=1e-10, atol=1e-15 + ) def test_curl_linear_field(self, gridpath, datasetpath): """Test curl of linear vector field""" @@ -686,7 +713,7 @@ def test_curl_units_and_attributes(self, gridpath, datasetpath): # Check attributes assert "long_name" in curl_field.attrs assert "description" in curl_field.attrs - assert curl_field.attrs["description"] == "Curl of vector field computed as ∂v/∂x - ∂u/∂y" + assert curl_field.attrs["description"] == "Curl of vector field computed as ∂v/∂x - ∂u/∂y + u·tan(φ)/a" # Check name expected_name = f"curl_{u_component.name}_{v_component.name}" @@ -694,3 +721,129 @@ def test_curl_units_and_attributes(self, gridpath, datasetpath): # Check that grid is preserved assert curl_field.uxgrid == u_component.uxgrid + + +class TestSphericalManufacturedSolutions: + """Amplitude checks against closed-form answers on the unit sphere. + + The existing suite leans on null tests (constant fields, curl of a + gradient) and on sign/ordering comparisons. Those stay satisfied when an + operator is off by a constant factor, which is how the dual/primal area + mismatch and the missing metric terms survived. Each test below has a + non-zero exact answer, so a scale error fails immediately. + """ + + # Away from the poles, where tan(lat) blows up and the finite-volume + # stencil degrades. + MIDLAT = np.deg2rad(60) + + @staticmethod + def _grids(): + yield "healpix_z5_quad", ux.Grid.from_healpix(zoom=5) + + def test_gradient_amplitude(self): + """grad of sin(lat) has meridional component cos(lat), zonal zero.""" + for label, uxgrid in self._grids(): + lat = np.deg2rad(uxgrid.face_lat.values) + interior = np.abs(lat) < self.MIDLAT + + phi = ux.UxDataArray( + np.sin(lat), dims=["n_face"], uxgrid=uxgrid, name="phi" + ) + grad = phi.gradient(scale_by_radius=False) + + with np.errstate(divide="ignore", invalid="ignore"): + ratio = grad["meridional_gradient"].values / np.cos(lat) + sel = interior & np.isfinite(ratio) + assert np.abs(np.median(ratio[sel]) - 1.0) < 0.01, label + + def test_curl_solid_body_rotation(self): + """u = cos(lat), v = 0 has relative vorticity 2*sin(lat).""" + for label, uxgrid in self._grids(): + lat = np.deg2rad(uxgrid.face_lat.values) + interior = np.abs(lat) < self.MIDLAT + + u = ux.UxDataArray( + np.cos(lat), dims=["n_face"], uxgrid=uxgrid, name="u" + ) + v = ux.UxDataArray( + np.zeros_like(lat), dims=["n_face"], uxgrid=uxgrid, name="v" + ) + + with np.errstate(divide="ignore", invalid="ignore"): + ratio = u.curl(v, scale_by_radius=False).values / ( + 2 * np.sin(lat) + ) + sel = interior & np.isfinite(ratio) + assert np.abs(np.median(ratio[sel]) - 1.0) < 0.01, label + + def test_divergence_amplitude(self): + """u = 0, v = cos(lat) has divergence -2*sin(lat).""" + for label, uxgrid in self._grids(): + lat = np.deg2rad(uxgrid.face_lat.values) + interior = np.abs(lat) < self.MIDLAT + + u = ux.UxDataArray( + np.zeros_like(lat), dims=["n_face"], uxgrid=uxgrid, name="u" + ) + v = ux.UxDataArray( + np.cos(lat), dims=["n_face"], uxgrid=uxgrid, name="v" + ) + + with np.errstate(divide="ignore", invalid="ignore"): + ratio = u.divergence(v, scale_by_radius=False).values / ( + -2 * np.sin(lat) + ) + sel = interior & np.isfinite(ratio) + assert np.abs(np.median(ratio[sel]) - 1.0) < 0.01, label + + def test_gradient_converges_under_refinement(self): + """The error shrinks with resolution instead of sitting at a factor. + + This is the check that separates a normalization bug from truncation + error: before the fix the ratio converged to 4.0 on quads. + """ + errors = [] + for zoom in (4, 5, 6): + uxgrid = ux.Grid.from_healpix(zoom=zoom) + lat = np.deg2rad(uxgrid.face_lat.values) + interior = np.abs(lat) < self.MIDLAT + + phi = ux.UxDataArray( + np.sin(lat), dims=["n_face"], uxgrid=uxgrid, name="phi" + ) + grad = phi.gradient(scale_by_radius=False) + + with np.errstate(divide="ignore", invalid="ignore"): + ratio = grad["meridional_gradient"].values / np.cos(lat) + sel = interior & np.isfinite(ratio) + errors.append(abs(np.median(ratio[sel]) - 1.0)) + + assert errors[1] < errors[0] + assert errors[2] < errors[1] + + def test_hexagonal_matches_quadrilateral(self, gridpath): + """The answer must not depend on cell topology. + + The dual/primal ratio was ~4 on quads and ~3 on hexagons, so an + inflated gradient showed up as a topology-dependent answer. + """ + results = {} + for label, uxgrid in ( + ("quad", ux.Grid.from_healpix(zoom=5)), + ("hex", ux.open_grid(gridpath("mpas", "QU", "480", "grid.nc"))), + ): + lat = np.deg2rad(uxgrid.face_lat.values) + interior = np.abs(lat) < self.MIDLAT + + phi = ux.UxDataArray( + np.sin(lat), dims=["n_face"], uxgrid=uxgrid, name="phi" + ) + grad = phi.gradient(scale_by_radius=False) + + with np.errstate(divide="ignore", invalid="ignore"): + ratio = grad["meridional_gradient"].values / np.cos(lat) + sel = interior & np.isfinite(ratio) + results[label] = np.median(ratio[sel]) + + assert abs(results["quad"] - results["hex"]) < 0.02 diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 28fa22886..a58282501 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1663,8 +1663,17 @@ def curl( other, scale_by_radius=scale_by_radius ) - # Compute curl = ∂v/∂x - ∂u/∂y - curl_values = grad_v_zonal.values - grad_u_meridional.values + # Compute curl = ∂v/∂x - ∂u/∂y + u·tan(φ)/a + # + # The trailing term is the spherical metric term. Dropping it is only + # valid on a plane; on the sphere it costs a factor of two on + # solid-body rotation. When the derivatives have been divided by the + # radius the term carries the same 1/a factor. + tan_lat = np.tan(np.deg2rad(self.uxgrid.face_lat.values)) + metric = self.values * tan_lat + if scale_by_radius and "sphere_radius" in self.uxgrid._ds.attrs: + metric = metric / self.uxgrid._ds.attrs["sphere_radius"] + curl_values = grad_v_zonal.values - grad_u_meridional.values + metric u_units = self.attrs.get("units", "") has_sphere_radius = "sphere_radius" in self.uxgrid._ds.attrs @@ -1680,7 +1689,9 @@ def curl( attrs={ "long_name": f"Curl of ({self.name}, {other.name})", "units": curl_units, - "description": "Curl of vector field computed as ∂v/∂x - ∂u/∂y", + "description": ( + "Curl of vector field computed as ∂v/∂x - ∂u/∂y + u·tan(φ)/a" + ), }, uxgrid=self.uxgrid, name=f"curl_{self.name}_{other.name}", @@ -1756,7 +1767,7 @@ def divergence( u_gradient = self.gradient(scale_by_radius=scale_by_radius) v_gradient = other.gradient(scale_by_radius=scale_by_radius) - # For divergence: div(V) = ∂u/∂x + ∂v/∂y + # For divergence: div(V) = ∂u/∂x + ∂v/∂y - v·tan(φ)/a # We use the zonal gradient (∂/∂lon) of u and meridional gradient (∂/∂lat) of v u = u_gradient["zonal_gradient"] v = v_gradient["meridional_gradient"] @@ -1764,6 +1775,14 @@ def divergence( # Align DataArrays to ensure coords/dims match, then perform xarray-aware addition u, v = xr.align(u, v) divergence = u + v + + # Spherical metric term, the companion of the one in curl(). Omitting + # it is only valid on a plane. + tan_lat = np.tan(np.deg2rad(self.uxgrid.face_lat.values)) + metric = other.values * tan_lat + if scale_by_radius and "sphere_radius" in self.uxgrid._ds.attrs: + metric = metric / self.uxgrid._ds.attrs["sphere_radius"] + divergence = divergence - metric divergence.name = "divergence" # Infer units consistently with gradient()/curl(): a divergence is a diff --git a/uxarray/core/gradient.py b/uxarray/core/gradient.py index 30e2537d1..a1d385bc0 100644 --- a/uxarray/core/gradient.py +++ b/uxarray/core/gradient.py @@ -226,15 +226,57 @@ def _compute_gradient(data, scale_by_radius=True): @njit(cache=True) -def _normalize_and_project_gradient( - gradient, index, normal_lat, normal_lon, node_coords, node_neighbors -): - area, _ = calculate_face_area( - node_coords[0, node_neighbors].astype(np.float64), - node_coords[1, node_neighbors].astype(np.float64), - node_coords[2, node_neighbors].astype(np.float64), - ) +def _dual_cell_area(stencil_coords): + """Spherical area of the contour the Green-Gauss loop integrates around. + + ``stencil_coords`` holds the (3, n) Cartesian centroids of the faces that + form the contour. They arrive in connectivity order, which is not + necessarily the order in which they trace the polygon, so they are sorted + by azimuth about the contour centroid before the area is measured. + """ + n = stencil_coords.shape[1] + + # Contour centroid, used as the pole of the local azimuthal sort. + cx = np.sum(stencil_coords[0]) / n + cy = np.sum(stencil_coords[1]) / n + cz = np.sum(stencil_coords[2]) / n + cnorm = np.sqrt(cx * cx + cy * cy + cz * cz) + cx, cy, cz = cx / cnorm, cy / cnorm, cz / cnorm + + # Build a local tangent basis at the centroid. + if np.abs(cz) < 0.9: + ax, ay, az = 0.0, 0.0, 1.0 + else: + ax, ay, az = 1.0, 0.0, 0.0 + ex = ay * cz - az * cy + ey = az * cx - ax * cz + ez = ax * cy - ay * cx + enorm = np.sqrt(ex * ex + ey * ey + ez * ez) + ex, ey, ez = ex / enorm, ey / enorm, ez / enorm + fx = cy * ez - cz * ey + fy = cz * ex - cx * ez + fz = cx * ey - cy * ex + + angles = np.empty(n) + for i in range(n): + px = stencil_coords[0, i] + py = stencil_coords[1, i] + pz = stencil_coords[2, i] + angles[i] = np.arctan2(px * fx + py * fy + pz * fz, px * ex + py * ey + pz * ez) + + order = np.argsort(angles) + sorted_coords = np.empty((3, n)) + for i in range(n): + sorted_coords[0, i] = stencil_coords[0, order[i]] + sorted_coords[1, i] = stencil_coords[1, order[i]] + sorted_coords[2, i] = stencil_coords[2, order[i]] + + area, _ = calculate_face_area(sorted_coords[0], sorted_coords[1], sorted_coords[2]) + return area + +@njit(cache=True) +def _normalize_and_project_gradient(gradient, index, normal_lat, normal_lon, area): gradient = gradient / area # projection to horizontal gradient @@ -298,6 +340,17 @@ def _compute_gradients_on_faces( gradient = np.zeros(3) has_contribution = False + # Centroids of the faces forming the contour, collected as the loop + # walks it so the normalizing area matches the region integrated over. + max_stencil = face_node_connectivity.shape[1] * node_edge_connectivity.shape[1] + stencil = np.empty(max_stencil, dtype=np.int64) + n_stencil = 0 + + # Green-Gauss only applies to a closed contour. If any edge in this + # face's node neighborhood is missing a second face, the contour is + # open and no enclosed area exists. + contour_closed = True + for node_idx in face_node_connectivity[face_idx]: # take each node on that face if node_idx != INT_FILL_VALUE: for edge_idx in node_edge_connectivity[ @@ -310,6 +363,7 @@ def _compute_gradients_on_faces( # spurious INT_FILL_VALUE entries (e.g. SCRIP- # derived SE grids like ne120np4). See #1452. if INT_FILL_VALUE in edge_face_connectivity[edge_idx]: + contour_closed = False continue if ( @@ -347,16 +401,34 @@ def _compute_gradients_on_faces( ) has_contribution = True - if not has_contribution: - gradient = np.full(3, np.nan) - - node_neighbors = face_node_connectivity[face_idx] - node_neighbors = node_neighbors[node_neighbors != INT_FILL_VALUE] + for cand in (face1_idx, face2_idx): + seen = False + for s in range(n_stencil): + if stencil[s] == cand: + seen = True + break + if not seen: + stencil[n_stencil] = cand + n_stencil += 1 + + # The contour must be closed, and a polygon, before it encloses an area. + if not has_contribution or not contour_closed or n_stencil < 3: + gradients_faces[face_idx, 0] = np.nan + gradients_faces[face_idx, 1] = np.nan + continue + + stencil_coords = np.empty((3, n_stencil)) + for s in range(n_stencil): + stencil_coords[0, s] = face_coords[stencil[s], 0] + stencil_coords[1, s] = face_coords[stencil[s], 1] + stencil_coords[2, s] = face_coords[stencil[s], 2] + + area = _dual_cell_area(stencil_coords) # Normalize and project zonal and meridional components and store the result for the current face gradients_faces[face_idx, 0], gradients_faces[face_idx, 1] = ( _normalize_and_project_gradient( - gradient, face_idx, normal_lat, normal_lon, node_coords, node_neighbors + gradient, face_idx, normal_lat, normal_lon, area ) ) From fdfd2da1818c70c46ba1e72f57f28cabd2e0ad2d Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Mon, 10 Aug 2026 15:45:02 -0500 Subject: [PATCH 2/3] Speed up gradient computation ~27x by replacing quadrature area with closed form _compute_gradients_on_faces spent half its time re-deriving the dual-cell area with 4th-order Gaussian quadrature (16 points per fan triangle, ~580 heap allocations per face) when the area of a polygon bounded by great-circle arcs has an exact closed form (spherical excess). The other half was lost to per-edge np.cross/np.linalg.norm temporaries and array-based fill-value checks; both are now scalar, following the zero-allocation convention used elsewhere in the grid kernels. Gradient values shift by up to 5e-6 relative to the old quadrature-based result -- that gap is the retired quadrature's own truncation error, so the new values are the more accurate ones. All existing test assertions hold to 6 significant digits. Added a regression test pinning scale invariance, since the dyamond-30km test grid stores face_x/y/z in meters rather than as unit vectors, which is the one input that distinguishes a correct closed-form implementation from a plausible wrong one. Co-Authored-By: Claude Sonnet 5 --- test/core/test_vector_calculus.py | 26 +++ uxarray/core/gradient.py | 350 +++++++++++++++++++----------- 2 files changed, 252 insertions(+), 124 deletions(-) diff --git a/test/core/test_vector_calculus.py b/test/core/test_vector_calculus.py index a3b9dc1f6..5334d018e 100644 --- a/test/core/test_vector_calculus.py +++ b/test/core/test_vector_calculus.py @@ -185,6 +185,32 @@ def test_gradient_scales_by_radius(self, gridpath, datasetpath): assert grad_scaled["zonal_gradient"].attrs["units"].endswith("/m") assert grad_unit["zonal_gradient"].attrs["units"].endswith("/rad") + def test_gradient_invariant_to_coordinate_scale(self, gridpath): + """The gradient must not depend on the units of face_x/face_y/face_z. + + This grid is the one test mesh whose Cartesian face coordinates are + stored in meters (norm 6.37e6) rather than as unit vectors, so it is + what distinguishes a dual-cell area that normalizes its inputs from one + that does not. Skipping the normalization inflates the area by the + radius squared (~4e13) and collapses the gradient to ~0 instead of + raising anything. + + d/dlat of sin(lat) is cos(lat), so the meridional component divided by + cos(lat) should be 1. + """ + uxgrid = ux.open_grid(gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc")) + lat = np.deg2rad(uxgrid.face_lat.values) + phi = ux.UxDataArray(np.sin(lat), dims=["n_face"], uxgrid=uxgrid, name="phi") + + grad = phi.gradient(scale_by_radius=False) + mg = grad["meridional_gradient"].values + + finite = np.isfinite(mg) + assert finite.any() + + ratio = mg[finite] / np.cos(lat[finite]) + assert np.abs(np.median(ratio) - 1.0) < 0.01 + class TestDivergenceQuadHex: diff --git a/uxarray/core/gradient.py b/uxarray/core/gradient.py index a1d385bc0..839004853 100644 --- a/uxarray/core/gradient.py +++ b/uxarray/core/gradient.py @@ -5,7 +5,6 @@ from uxarray.constants import INT_FILL_VALUE from uxarray.errors import DataCenteringError, DimensionError -from uxarray.grid.area import calculate_face_area def _calculate_edge_face_difference(d_var, edge_faces, n_edge): @@ -36,7 +35,7 @@ def _calculate_edge_node_difference(d_var, edge_nodes): return np.abs(edge_node_diff) -@njit(cache=True) +@njit(cache=True, inline="always") def _compute_arc_length(lat_a, lat_b, lon_a, lon_b): dlat = np.radians(lat_b - lat_a) dlon = np.radians(lon_b - lon_a) @@ -110,10 +109,6 @@ def _compute_gradient(data, scale_by_radius=True): face_lat = uxgrid.face_lat.values face_lon = uxgrid.face_lon.values - node_coords = np.array( - [uxgrid.node_x.values, uxgrid.node_y.values, uxgrid.node_z.values] - ) - face_lon_rad = np.deg2rad(face_lon) face_lat_rad = np.deg2rad(face_lat) normal_lat = np.array( @@ -140,7 +135,6 @@ def _compute_gradient(data, scale_by_radius=True): uxgrid.node_edge_connectivity.values, face_lat, face_lon, - node_coords, normal_lon, normal_lat, ) @@ -226,22 +220,67 @@ def _compute_gradient(data, scale_by_radius=True): @njit(cache=True) -def _dual_cell_area(stencil_coords): +def _dual_cell_area(sx, sy, sz, angles, n): """Spherical area of the contour the Green-Gauss loop integrates around. - ``stencil_coords`` holds the (3, n) Cartesian centroids of the faces that - form the contour. They arrive in connectivity order, which is not - necessarily the order in which they trace the polygon, so they are sorted - by azimuth about the contour centroid before the area is measured. + ``sx``/``sy``/``sz`` hold the Cartesian centroids of the faces forming the + contour in their first ``n`` entries; ``angles`` is a scratch buffer of at + least ``n`` entries. All four are modified in place. + + The centroids arrive in connectivity order, which is not necessarily the + order in which they trace the polygon, so they are sorted by azimuth about + the contour centroid before the area is measured. + + Notes + ----- + This replaces a call to :func:`uxarray.grid.area.calculate_face_area`, which + integrated this area with 4th-order Gaussian quadrature: 16 quadrature + points per fan triangle, each evaluating a spherical Jacobian. For a polygon + bounded by great-circle arcs that integral has a closed form -- the + spherical excess, via Van Oosterom and Strackee -- so the quadrature cost + roughly 80x the exact answer (23.9 us vs 0.3 us per call) and was slightly + less accurate. + + Two properties of the retired quadrature are reproduced deliberately: + + - It accumulated *unsigned* Jacobians, i.e. an unsigned sum over fan + triangles from vertex 0. Dual cells are often non-convex -- a HealPix dual + cell alternates near edge-neighbor centroids with far corner-neighbor + centroids -- and for those a signed excess sum differs from an unsigned + one by ~1%. Summing ``abs`` per triangle matches the previous area to + ~5e-14; a signed sum does not. This is a compatibility choice, not a claim + that unsigned is the geometrically correct convention for a + self-overlapping fan. + - It was scale-invariant, because its Jacobian divided the radius out. The + excess formula is not, so the points are normalized onto the unit sphere + first: some grids store ``face_x``/``face_y``/``face_z`` in meters rather + than as unit vectors (e.g. the dyamond-30km test subset, at a radius of + 6.37e6). + + Gradients therefore shift by up to ~5e-6 relative to the quadrature-based + implementation. That gap is the retired quadrature's own truncation error; + the closed form is the more accurate value. """ - n = stencil_coords.shape[1] + # The excess formula needs unit vectors; face centroids are not always stored + # normalized. + for i in range(n): + inv_r = 1.0 / np.sqrt(sx[i] * sx[i] + sy[i] * sy[i] + sz[i] * sz[i]) + sx[i] *= inv_r + sy[i] *= inv_r + sz[i] *= inv_r # Contour centroid, used as the pole of the local azimuthal sort. - cx = np.sum(stencil_coords[0]) / n - cy = np.sum(stencil_coords[1]) / n - cz = np.sum(stencil_coords[2]) / n - cnorm = np.sqrt(cx * cx + cy * cy + cz * cz) - cx, cy, cz = cx / cnorm, cy / cnorm, cz / cnorm + cx = 0.0 + cy = 0.0 + cz = 0.0 + for i in range(n): + cx += sx[i] + cy += sy[i] + cz += sz[i] + inv_c = 1.0 / np.sqrt(cx * cx + cy * cy + cz * cz) + cx *= inv_c + cy *= inv_c + cz *= inv_c # Build a local tangent basis at the centroid. if np.abs(cz) < 0.9: @@ -251,41 +290,67 @@ def _dual_cell_area(stencil_coords): ex = ay * cz - az * cy ey = az * cx - ax * cz ez = ax * cy - ay * cx - enorm = np.sqrt(ex * ex + ey * ey + ez * ez) - ex, ey, ez = ex / enorm, ey / enorm, ez / enorm + inv_e = 1.0 / np.sqrt(ex * ex + ey * ey + ez * ez) + ex *= inv_e + ey *= inv_e + ez *= inv_e fx = cy * ez - cz * ey fy = cz * ex - cx * ez fz = cx * ey - cy * ex - angles = np.empty(n) for i in range(n): - px = stencil_coords[0, i] - py = stencil_coords[1, i] - pz = stencil_coords[2, i] - angles[i] = np.arctan2(px * fx + py * fy + pz * fz, px * ex + py * ey + pz * ez) + angles[i] = np.arctan2( + sx[i] * fx + sy[i] * fy + sz[i] * fz, sx[i] * ex + sy[i] * ey + sz[i] * ez + ) - order = np.argsort(angles) - sorted_coords = np.empty((3, n)) - for i in range(n): - sorted_coords[0, i] = stencil_coords[0, order[i]] - sorted_coords[1, i] = stencil_coords[1, order[i]] - sorted_coords[2, i] = stencil_coords[2, order[i]] + # Insertion sort by azimuth, carrying the coordinates along. n is bounded by + # n_max_face_nodes * n_max_node_edges (order 10s), so this beats allocating + # an argsort permutation and a second coordinate buffer per face. + for i in range(1, n): + key = angles[i] + kx = sx[i] + ky = sy[i] + kz = sz[i] + j = i - 1 + while j >= 0 and angles[j] > key: + angles[j + 1] = angles[j] + sx[j + 1] = sx[j] + sy[j + 1] = sy[j] + sz[j + 1] = sz[j] + j -= 1 + angles[j + 1] = key + sx[j + 1] = kx + sy[j + 1] = ky + sz[j + 1] = kz + + # Spherical excess of each fan triangle from vertex 0, summed unsigned. + area = 0.0 + ax = sx[0] + ay = sy[0] + az = sz[0] + for j in range(1, n - 1): + bx = sx[j] + by = sy[j] + bz = sz[j] + cx = sx[j + 1] + cy = sy[j + 1] + cz = sz[j + 1] + triple = ( + ax * (by * cz - bz * cy) + + ay * (bz * cx - bx * cz) + + az * (bx * cy - by * cx) + ) + denom = ( + 1.0 + + (ax * bx + ay * by + az * bz) + + (bx * cx + by * cy + bz * cz) + + (cx * ax + cy * ay + cz * az) + ) + area += np.abs(2.0 * np.arctan2(triple, denom)) - area, _ = calculate_face_area(sorted_coords[0], sorted_coords[1], sorted_coords[2]) return area -@njit(cache=True) -def _normalize_and_project_gradient(gradient, index, normal_lat, normal_lon, area): - gradient = gradient / area - - # projection to horizontal gradient - zonal_grad = np.sum(gradient * normal_lon[index]) - meridional_grad = np.sum(gradient * normal_lat[index]) - - return zonal_grad, meridional_grad - - @njit(cache=True, parallel=True) def _compute_gradients_on_faces( data, @@ -296,7 +361,6 @@ def _compute_gradients_on_faces( node_edge_connectivity, face_lat, face_lon, - node_coords, normal_lon, normal_lat, ): @@ -333,106 +397,144 @@ def _compute_gradients_on_faces( """ - gradients_faces = np.full((n_face, 2), np.nan) + gradient_zonal = np.empty(n_face) + gradient_meridional = np.empty(n_face) + + n_face_nodes = face_node_connectivity.shape[1] + n_node_edges = node_edge_connectivity.shape[1] + max_stencil = n_face_nodes * n_node_edges # Parallel across faces for face_idx in prange(n_face): - gradient = np.zeros(3) - has_contribution = False - # Centroids of the faces forming the contour, collected as the loop # walks it so the normalizing area matches the region integrated over. - max_stencil = face_node_connectivity.shape[1] * node_edge_connectivity.shape[1] stencil = np.empty(max_stencil, dtype=np.int64) + stencil_x = np.empty(max_stencil) + stencil_y = np.empty(max_stencil) + stencil_z = np.empty(max_stencil) + angles = np.empty(max_stencil) n_stencil = 0 + # Gradient accumulated component-wise to keep the inner loop free of + # temporary arrays. + grad_x = 0.0 + grad_y = 0.0 + grad_z = 0.0 + has_contribution = False + # Green-Gauss only applies to a closed contour. If any edge in this # face's node neighborhood is missing a second face, the contour is # open and no enclosed area exists. contour_closed = True - for node_idx in face_node_connectivity[face_idx]: # take each node on that face - if node_idx != INT_FILL_VALUE: - for edge_idx in node_edge_connectivity[ - node_idx - ]: # grab each edge connected to that node - if edge_idx != INT_FILL_VALUE: - # Skip edges that lack a second face neighbor - # instead of NaN-ing the entire face. Fixes - # grids where edge_face_connectivity has - # spurious INT_FILL_VALUE entries (e.g. SCRIP- - # derived SE grids like ne120np4). See #1452. - if INT_FILL_VALUE in edge_face_connectivity[edge_idx]: - contour_closed = False - continue - - if ( - face_idx not in edge_face_connectivity[edge_idx] - ): # check if edge connected to original face - face1_idx = edge_face_connectivity[edge_idx][0] - face2_idx = edge_face_connectivity[edge_idx][1] - - face1_coords = face_coords[face1_idx] - face2_coords = face_coords[face2_idx] - - # compute normal pointing outwards from face - cross = np.cross(face1_coords, face2_coords) - norm = np.linalg.norm(cross) - if np.dot(cross, face1_coords - face_coords[face_idx]) > 0: - normal = cross / norm - else: - normal = -cross / norm - - # compute arc length between the two faces - arc_length = _compute_arc_length( - face_lat[face1_idx], - face_lat[face2_idx], - face_lon[face1_idx], - face_lon[face2_idx], - ) - - # compute trapezoidal rule - trapz = (data[face1_idx] + data[face2_idx]) / 2 - - # add to the gradient (subtract correction term) - gradient = ( - gradient - + (trapz - data[face_idx]) * arc_length * normal - ) - has_contribution = True - - for cand in (face1_idx, face2_idx): - seen = False - for s in range(n_stencil): - if stencil[s] == cand: - seen = True - break - if not seen: - stencil[n_stencil] = cand - n_stencil += 1 + data_face = data[face_idx] + face_x = face_coords[face_idx, 0] + face_y = face_coords[face_idx, 1] + face_z = face_coords[face_idx, 2] + + for i in range(n_face_nodes): # take each node on that face + node_idx = face_node_connectivity[face_idx, i] + if node_idx == INT_FILL_VALUE: + continue + + for j in range(n_node_edges): # grab each edge connected to that node + edge_idx = node_edge_connectivity[node_idx, j] + if edge_idx == INT_FILL_VALUE: + continue + + # edge_face_connectivity is always (n_edge, 2), so the two + # neighbors can be compared directly. + face1_idx = edge_face_connectivity[edge_idx, 0] + face2_idx = edge_face_connectivity[edge_idx, 1] + + # Skip edges that lack a second face neighbor instead of + # NaN-ing the entire face. Fixes grids where + # edge_face_connectivity has spurious INT_FILL_VALUE entries + # (e.g. SCRIP-derived SE grids like ne120np4). See #1452. + if face1_idx == INT_FILL_VALUE or face2_idx == INT_FILL_VALUE: + contour_closed = False + continue + + # skip edges connected to the original face + if face1_idx == face_idx or face2_idx == face_idx: + continue + + face1_x = face_coords[face1_idx, 0] + face1_y = face_coords[face1_idx, 1] + face1_z = face_coords[face1_idx, 2] + face2_x = face_coords[face2_idx, 0] + face2_y = face_coords[face2_idx, 1] + face2_z = face_coords[face2_idx, 2] + + # compute normal pointing outwards from face + cross_x = face1_y * face2_z - face1_z * face2_y + cross_y = face1_z * face2_x - face1_x * face2_z + cross_z = face1_x * face2_y - face1_y * face2_x + norm = np.sqrt( + cross_x * cross_x + cross_y * cross_y + cross_z * cross_z + ) + if ( + cross_x * (face1_x - face_x) + + cross_y * (face1_y - face_y) + + cross_z * (face1_z - face_z) + ) > 0: + inv_norm = 1.0 / norm + else: + inv_norm = -1.0 / norm + + # compute arc length between the two faces + arc_length = _compute_arc_length( + face_lat[face1_idx], + face_lat[face2_idx], + face_lon[face1_idx], + face_lon[face2_idx], + ) + + # compute trapezoidal rule + trapz = (data[face1_idx] + data[face2_idx]) / 2 + + # add to the gradient (subtract correction term) + weight = (trapz - data_face) * arc_length * inv_norm + grad_x += weight * cross_x + grad_y += weight * cross_y + grad_z += weight * cross_z + has_contribution = True + + for cand in (face1_idx, face2_idx): + seen = False + for s in range(n_stencil): + if stencil[s] == cand: + seen = True + break + if not seen: + stencil[n_stencil] = cand + stencil_x[n_stencil] = face_coords[cand, 0] + stencil_y[n_stencil] = face_coords[cand, 1] + stencil_z[n_stencil] = face_coords[cand, 2] + n_stencil += 1 # The contour must be closed, and a polygon, before it encloses an area. if not has_contribution or not contour_closed or n_stencil < 3: - gradients_faces[face_idx, 0] = np.nan - gradients_faces[face_idx, 1] = np.nan + gradient_zonal[face_idx] = np.nan + gradient_meridional[face_idx] = np.nan continue - stencil_coords = np.empty((3, n_stencil)) - for s in range(n_stencil): - stencil_coords[0, s] = face_coords[stencil[s], 0] - stencil_coords[1, s] = face_coords[stencil[s], 1] - stencil_coords[2, s] = face_coords[stencil[s], 2] - - area = _dual_cell_area(stencil_coords) + area = _dual_cell_area(stencil_x, stencil_y, stencil_z, angles, n_stencil) # Normalize and project zonal and meridional components and store the result for the current face - gradients_faces[face_idx, 0], gradients_faces[face_idx, 1] = ( - _normalize_and_project_gradient( - gradient, face_idx, normal_lat, normal_lon, area - ) - ) - - return gradients_faces[:, 0], gradients_faces[:, 1] + inv_area = 1.0 / area + gradient_zonal[face_idx] = ( + grad_x * normal_lon[face_idx, 0] + + grad_y * normal_lon[face_idx, 1] + + grad_z * normal_lon[face_idx, 2] + ) * inv_area + gradient_meridional[face_idx] = ( + grad_x * normal_lat[face_idx, 0] + + grad_y * normal_lat[face_idx, 1] + + grad_z * normal_lat[face_idx, 2] + ) * inv_area + + return gradient_zonal, gradient_meridional # TODO: Add support for this after merging face-centered implementation From ebd09638d3773b3e5a46665d9b0024906df7606a Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Tue, 11 Aug 2026 08:54:11 -0500 Subject: [PATCH 3/3] Sum the dual-cell fan signed rather than unsigned The closed-form area reproduced the retired quadrature's unsigned accumulation for bit-compatibility. That convention is wrong for a non-convex cell: a fan from vertex 0 sweeps some triangles backwards and they have to cancel, but taking abs per triangle double-counts them. About 46% of HealPix z4 dual cells are non-convex, differing by up to 20% in area. Measured against d/dlat sin(lat) = cos(lat), the unsigned sum leaves a worst-case gradient error of ~19% that does not improve with resolution (0.186 / 0.192 / 0.196 at z4 / z5 / z6). Signed converges as expected (0.0127 / 0.0088 / 0.0051). Only the tail is affected, which is why the median looked healthy either way; the new test bounds the worst face so a regression to unsigned fails. --- test/core/test_vector_calculus.py | 30 +++++++++++++++++++++ uxarray/core/gradient.py | 43 +++++++++++++++---------------- 2 files changed, 51 insertions(+), 22 deletions(-) diff --git a/test/core/test_vector_calculus.py b/test/core/test_vector_calculus.py index 5334d018e..9d37fd1f2 100644 --- a/test/core/test_vector_calculus.py +++ b/test/core/test_vector_calculus.py @@ -212,6 +212,36 @@ def test_gradient_invariant_to_coordinate_scale(self, gridpath): assert np.abs(np.median(ratio) - 1.0) < 0.01 +class TestGradientDualCellArea: + """The dual cell is frequently non-convex, so its area has to be measured + with a signed fan sum.""" + + @pytest.mark.parametrize("zoom,tol", [(4, 0.02), (5, 0.01), (6, 0.006)]) + def test_gradient_worst_case_error_converges(self, zoom, tol): + """Bound the worst face, not just the median. + + A HealPix dual cell alternates near edge-neighbor centroids with far + corner-neighbor centroids, so a fan from vertex 0 sweeps some triangles + backwards. Summing their spherical excess unsigned double-counts those + and leaves a ~19% error on the worst face that does not shrink with + resolution; only the median looks healthy. Guard the tail so that a + regression to an unsigned sum fails here. + """ + uxgrid = ux.Grid.from_healpix(zoom) + lat = np.deg2rad(uxgrid.face_lat.values) + phi = ux.UxDataArray(np.sin(lat), dims=["n_face"], uxgrid=uxgrid, name="phi") + + mg = phi.gradient(scale_by_radius=False)["meridional_gradient"].values + + # d/dlat of sin(lat) is cos(lat); stay off the poles where the + # projection is ill-conditioned. + ok = np.isfinite(mg) & (np.abs(lat) < np.deg2rad(60)) + assert ok.any() + + error = np.abs(mg[ok] / np.cos(lat[ok]) - 1.0) + assert error.max() < tol + + class TestDivergenceQuadHex: def test_divergence_output_format(self, gridpath, datasetpath): diff --git a/uxarray/core/gradient.py b/uxarray/core/gradient.py index 839004853..c3782f869 100644 --- a/uxarray/core/gradient.py +++ b/uxarray/core/gradient.py @@ -241,25 +241,23 @@ def _dual_cell_area(sx, sy, sz, angles, n): roughly 80x the exact answer (23.9 us vs 0.3 us per call) and was slightly less accurate. - Two properties of the retired quadrature are reproduced deliberately: - - - It accumulated *unsigned* Jacobians, i.e. an unsigned sum over fan - triangles from vertex 0. Dual cells are often non-convex -- a HealPix dual - cell alternates near edge-neighbor centroids with far corner-neighbor - centroids -- and for those a signed excess sum differs from an unsigned - one by ~1%. Summing ``abs`` per triangle matches the previous area to - ~5e-14; a signed sum does not. This is a compatibility choice, not a claim - that unsigned is the geometrically correct convention for a - self-overlapping fan. - - It was scale-invariant, because its Jacobian divided the radius out. The - excess formula is not, so the points are normalized onto the unit sphere - first: some grids store ``face_x``/``face_y``/``face_z`` in meters rather - than as unit vectors (e.g. the dyamond-30km test subset, at a radius of - 6.37e6). - - Gradients therefore shift by up to ~5e-6 relative to the quadrature-based - implementation. That gap is the retired quadrature's own truncation error; - the closed form is the more accurate value. + The fan triangles are summed *signed*, and the magnitude is taken once at + the end. Dual cells are frequently non-convex -- a HealPix dual cell + alternates near edge-neighbor centroids with far corner-neighbor centroids + -- and for a non-convex polygon a fan from vertex 0 sweeps some triangles + backwards. Those must cancel. The retired quadrature accumulated unsigned + Jacobians instead, which double-counts the reversed triangles: measured + against the exact ``d/dlat sin(lat) = cos(lat)`` solution on HealPix, the + unsigned area leaves a worst-case gradient error of ~19% that does not + improve with resolution (0.186 / 0.192 / 0.196 at z4 / z5 / z6), while the + signed sum converges as expected (0.0127 / 0.0088 / 0.0051). Roughly 46% of + z4 dual cells are affected, by up to 20% in area. + + Scale invariance is preserved from the quadrature, whose Jacobian divided + the radius out. The excess formula has no such property, so the points are + normalized onto the unit sphere first: some grids store + ``face_x``/``face_y``/``face_z`` in meters rather than as unit vectors + (e.g. the dyamond-30km test subset, at a radius of 6.37e6). """ # The excess formula needs unit vectors; face centroids are not always stored # normalized. @@ -323,7 +321,8 @@ def _dual_cell_area(sx, sy, sz, angles, n): sy[j + 1] = ky sz[j + 1] = kz - # Spherical excess of each fan triangle from vertex 0, summed unsigned. + # Spherical excess of each fan triangle from vertex 0, summed signed so + # that backward-swept triangles of a non-convex cell cancel. area = 0.0 ax = sx[0] ay = sy[0] @@ -346,9 +345,9 @@ def _dual_cell_area(sx, sy, sz, angles, n): + (bx * cx + by * cy + bz * cz) + (cx * ax + cy * ay + cz * az) ) - area += np.abs(2.0 * np.arctan2(triple, denom)) + area += 2.0 * np.arctan2(triple, denom) - return area + return np.abs(area) @njit(cache=True, parallel=True)