o Address comments, make one test file for vector calculus

Author: rajeeja

test/core/test_gradient.py

@@ -5,7 +5,144 @@
 import numpy.testing as nt
 
 
-class TestQuadHex:
+# TODO: pytest fixtures
+
+
+class TestGradientQuadHex:
+
+    def test_gradient_output_format(self, gridpath, datasetpath):
+        """Tests the output format of gradient functionality"""
+        uxds = ux.open_dataset(gridpath("ugrid", "quad-hexagon", "grid.nc"), datasetpath("ugrid", "quad-hexagon", "data.nc"))
+
+        grad_ds = uxds['t2m'].gradient()
+
+        assert isinstance(grad_ds, ux.UxDataset)
+        assert "zonal_gradient" in grad_ds
+        assert "meridional_gradient" in grad_ds
+        assert "gradient" in grad_ds.attrs
+        assert uxds['t2m'].sizes == grad_ds.sizes
+
+    def test_gradient_all_boundary_faces(self, gridpath, datasetpath):
+        """Quad hexagon grid has 4 faces, each of which are on the boundary, so the expected gradients are zero for both components"""
+        uxds = ux.open_dataset(gridpath("ugrid", "quad-hexagon", "grid.nc"), datasetpath("ugrid", "quad-hexagon", "data.nc"))
+
+        grad = uxds['t2m'].gradient()
+
+        assert np.isnan(grad['meridional_gradient']).all()
+        assert np.isnan(grad['zonal_gradient']).all()
+
+
+class TestGradientMPASOcean:
+
+    def test_gradient(self, gridpath, datasetpath):
+        uxds = ux.open_dataset(gridpath("mpas", "QU", "480", "grid.nc"), datasetpath("mpas", "QU", "480", "data.nc"))
+
+        grad = uxds['bottomDepth'].gradient()
+
+        # There should be some boundary faces
+        assert np.isnan(grad['meridional_gradient']).any()
+        assert np.isnan(grad['zonal_gradient']).any()
+
+        # Not every face is on the boundary, ensure there are valid values
+        assert not np.isnan(grad['meridional_gradient']).all()
+        assert not np.isnan(grad['zonal_gradient']).all()
+
+
+class TestGradientDyamondSubset:
+
+    center_fidx = 153
+    left_fidx   = 100
+    right_fidx  = 164
+    top_fidx    = 154
+    bottom_fidx = 66
+
+    def test_lat_field(self, gridpath, datasetpath):
+        """Gradient of a latitude field. All vectors should be pointing east."""
+        uxds =  ux.open_dataset(
+            gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+            datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+        )
+        grad = uxds['face_lat'].gradient()
+        zg, mg = grad.zonal_gradient, grad.meridional_gradient
+        assert mg.max() > zg.max()
+
+        assert mg.min() > zg.max()
+
+
+    def test_lon_field(self, gridpath, datasetpath):
+        """Gradient of a longitude field. All vectors should be pointing north."""
+        uxds =  ux.open_dataset(
+            gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+            datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+        )
+        grad = uxds['face_lon'].gradient()
+        zg, mg = grad.zonal_gradient, grad.meridional_gradient
+        assert zg.max() > mg.max()
+
+        assert zg.min() > mg.max()
+
+    def test_gaussian_field(self, gridpath, datasetpath):
+        """Gradient of a gaussian field. All vectors should be pointing toward the center"""
+        uxds =  ux.open_dataset(
+            gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+            datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+        )
+        grad = uxds['gaussian'].gradient()
+        zg, mg = grad.zonal_gradient, grad.meridional_gradient
+        mag = np.hypot(zg, mg)
+        angle = np.arctan2(mg, zg)
+
+        # Ensure a valid range for min/max
+        assert zg.min() < 0
+        assert zg.max() > 0
+        assert mg.min() < 0
+        assert mg.max() > 0
+
+        # The Magnitude at the center is less than the corners
+        assert mag[self.center_fidx] < mag[self.left_fidx]
+        assert mag[self.center_fidx] < mag[self.right_fidx]
+        assert mag[self.center_fidx] < mag[self.top_fidx]
+        assert mag[self.center_fidx] < mag[self.bottom_fidx]
+
+        # Pointing Towards Center
+        assert angle[self.left_fidx] < 0
+        assert angle[self.right_fidx] > 0
+        assert angle[self.top_fidx] < 0
+        assert angle[self.bottom_fidx] > 0
+
+
+
+    def test_inverse_gaussian_field(self, gridpath, datasetpath):
+        """Gradient of an inverse gaussian field. All vectors should be pointing outward from the center."""
+        uxds =  ux.open_dataset(
+            gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+            datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+        )
+        grad = uxds['inverse_gaussian'].gradient()
+        zg, mg = grad.zonal_gradient, grad.meridional_gradient
+        mag = np.hypot(zg, mg)
+        angle = np.arctan2(mg, zg)
+
+        # Ensure a valid range for min/max
+        assert zg.min() < 0
+        assert zg.max() > 0
+        assert mg.min() < 0
+        assert mg.max() > 0
+
+        # The Magnitude at the center is less than the corners
+        assert mag[self.center_fidx] < mag[self.left_fidx]
+        assert mag[self.center_fidx] < mag[self.right_fidx]
+        assert mag[self.center_fidx] < mag[self.top_fidx]
+        assert mag[self.center_fidx] < mag[self.bottom_fidx]
+
+        # Pointing Away from Center
+        assert angle[self.left_fidx] > 0
+        assert angle[self.right_fidx] < 0
+        assert angle[self.top_fidx] > 0
+        assert angle[self.bottom_fidx] < 0
+
+
+class TestDivergenceQuadHex:
 
     def test_divergence_output_format(self, gridpath, datasetpath):
         """Tests the output format of divergence functionality"""
@@ -39,7 +176,7 @@ def test_divergence_input_validation(self, gridpath, datasetpath):
         # This would require creating data with different dims, so we'll skip for now
 
 
-class TestMPASOcean:
+class TestDivergenceMPASOcean:
 
     def test_divergence_basic(self, gridpath, datasetpath):
         """Basic test of divergence computation"""
@@ -56,7 +193,7 @@ def test_divergence_basic(self, gridpath, datasetpath):
         assert np.isfinite(div_field.values).any()
 
 
-class TestDyamondSubset:
+class TestDivergenceDyamondSubset:
 
     def test_divergence_constant_field(self, gridpath, datasetpath):
         """Test divergence of constant vector field (should be zero)"""

test/core/test_divergence.py test/core/test_vector_calculus.pytest/core/test_divergence.py renamed to test/core/test_vector_calculus.py

@@ -1397,24 +1397,23 @@ def divergence(self, other: "UxDataArray", **kwargs) -> "UxDataArray":
         # For divergence: div(V) = ∂u/∂x + ∂v/∂y
         # In spherical coordinates: div(V) = (1/cos(lat)) * ∂u/∂lon + ∂v/∂lat
         # We use the zonal gradient (∂/∂lon) of u and meridional gradient (∂/∂lat) of v
-        divergence_values = (
-            u_gradient["zonal_gradient"].values
-            + v_gradient["meridional_gradient"].values
-        )
-
-        # Create the divergence UxDataArray
-        divergence_da = UxDataArray(
-            data=divergence_values,
-            name="divergence",
-            dims=self.dims,
-            uxgrid=self.uxgrid,
-            attrs={
+        u = u_gradient["zonal_gradient"]
+        v = v_gradient["meridional_gradient"]
+
+        # Align DataArrays to ensure coords/dims match, then perform xarray-aware addition
+        u, v = xr.align(u, v)
+        divergence = u + v
+        divergence.name = "divergence"
+        divergence.attrs.update(
+            {
                 "divergence": True,
                 "units": "1/s" if "units" not in kwargs else kwargs["units"],
-            },
-            coords=self.coords,
+            }
         )
 
+        # Wrap result as a UxDataArray while preserving uxgrid and coords
+        divergence_da = UxDataArray(divergence, uxgrid=self.uxgrid)
+
         return divergence_da
 
     def difference(self, destination: Optional[str] = "edge"):