o Divergence implementation, first pass

Author: rajeeja

test/core/test_divergence.py

@@ -0,0 +1,167 @@
+import numpy as np
+import pytest
+
+import uxarray as ux
+import numpy.testing as nt
+
+
+class TestQuadHex:
+
+    def test_divergence_output_format(self, gridpath, datasetpath):
+        """Tests the output format of divergence functionality"""
+        uxds = ux.open_dataset(gridpath("ugrid", "quad-hexagon", "grid.nc"), datasetpath("ugrid", "quad-hexagon", "data.nc"))
+
+        # Create two components for vector field (using same data for simplicity in test)
+        u_component = uxds['t2m']
+        v_component = uxds['t2m']
+        
+        div_da = u_component.divergence(v_component)
+
+        assert isinstance(div_da, ux.UxDataArray)
+        assert div_da.name == "divergence"
+        assert "divergence" in div_da.attrs
+        assert u_component.sizes == div_da.sizes
+
+    def test_divergence_input_validation(self, gridpath, datasetpath):
+        """Tests input validation for divergence method"""
+        uxds = ux.open_dataset(gridpath("ugrid", "quad-hexagon", "grid.nc"), datasetpath("ugrid", "quad-hexagon", "data.nc"))
+
+        u_component = uxds['t2m']
+        
+        # Test with non-UxDataArray
+        with pytest.raises(TypeError, match="other must be a UxDataArray"):
+            u_component.divergence(np.array([1, 2, 3]))
+        
+        # Test with different grids (create a simple test case)
+        # This would require creating another grid, so we'll skip for now
+        
+        # Test with different dimensions
+        # This would require creating data with different dims, so we'll skip for now
+
+
+class TestMPASOcean:
+
+    def test_divergence_basic(self, gridpath, datasetpath):
+        """Basic test of divergence computation"""
+        uxds = ux.open_dataset(gridpath("mpas", "QU", "480", "grid.nc"), datasetpath("mpas", "QU", "480", "data.nc"))
+
+        # Use the same field for both components (not physically meaningful but tests the method)
+        u_component = uxds['bottomDepth']
+        v_component = uxds['bottomDepth']
+        
+        div_field = u_component.divergence(v_component)
+
+        # Check that we get finite values where expected
+        assert not np.isnan(div_field.values).all()
+        assert np.isfinite(div_field.values).any()
+
+
+class TestDyamondSubset:
+
+    def test_divergence_constant_field(self, gridpath, datasetpath):
+        """Test divergence of constant vector field (should be zero)"""
+        uxds = ux.open_dataset(
+            gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+            datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+        )
+        
+        # Create constant fields
+        constant_u = uxds['face_lat'] * 0 + 1.0  # Constant field = 1
+        constant_v = uxds['face_lat'] * 0 + 1.0  # Constant field = 1
+        
+        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)]
+        
+        # Check that we have some finite values (interior faces)
+        assert len(finite_values) > 0, "No finite divergence values found"
+        
+        # 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()}"
+
+    def test_divergence_linear_field(self, gridpath, datasetpath):
+        """Test divergence of linear vector field"""
+        uxds = ux.open_dataset(
+            gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+            datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+        )
+        
+        # Create linear fields: u = x, v = y (in spherical coordinates: u = lon, v = lat)
+        u_component = uxds['face_lon']  # Linear in longitude
+        v_component = uxds['face_lat']  # Linear in latitude
+        
+        div_field = u_component.divergence(v_component)
+        
+        # Check that we have some finite values (interior faces)
+        finite_mask = np.isfinite(div_field.values)
+        finite_values = div_field.values[finite_mask]
+        
+        assert len(finite_values) > 0, "No finite divergence values found"
+        
+        # For linear fields, divergence should be finite where computable
+        # Boundary faces may have NaN values (which is expected)
+        assert not np.isnan(finite_values).any(), "Found NaN in finite values"
+
+    def test_divergence_radial_field(self, gridpath, datasetpath):
+        """Test divergence of radial vector field (should be positive)"""
+        uxds = ux.open_dataset(
+            gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+            datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+        )
+        
+        # Create a radial field pointing outward from center
+        # Use the inverse gaussian as a proxy for radial distance
+        radial_distance = uxds['inverse_gaussian']
+        
+        # Create components proportional to position (simplified radial field)
+        u_component = radial_distance * uxds['face_lon']
+        v_component = radial_distance * uxds['face_lat']
+        
+        div_field = u_component.divergence(v_component)
+        
+        # Check that we have some finite values (interior faces)
+        finite_mask = np.isfinite(div_field.values)
+        finite_values = div_field.values[finite_mask]
+        
+        assert len(finite_values) > 0, "No finite divergence values found"
+        
+        # Boundary faces may have NaN values (which is expected)
+        assert not np.isnan(finite_values).any(), "Found NaN in finite values"
+        
+        # Most finite values should be positive for an expanding field
+        positive_values = finite_values > 0
+        assert positive_values.sum() > len(finite_values) * 0.3  # At least 30% positive
+
+    def test_divergence_curl_identity(self, gridpath, datasetpath):
+        """Test that divergence of curl is zero (vector calculus identity)"""
+        uxds = ux.open_dataset(
+            gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"),
+            datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc")
+        )
+        
+        # Create a scalar potential field
+        potential = uxds['gaussian']
+        
+        # Compute gradient to get a conservative vector field
+        grad = potential.gradient()
+        u_component = grad['zonal_gradient']
+        v_component = grad['meridional_gradient']
+        
+        # Compute divergence of this gradient (should be the Laplacian)
+        div_grad = u_component.divergence(v_component)
+        
+        # Check that we have some finite values (interior faces)
+        finite_mask = np.isfinite(div_grad.values)
+        finite_values = div_grad.values[finite_mask]
+        
+        assert len(finite_values) > 0, "No finite divergence values found"
+        
+        # This tests the Laplacian computation via div(grad)
+        # Boundary faces may have NaN values (which is expected)
+        assert not np.isnan(finite_values).any(), "Found NaN in finite values"
+        
+        # The Laplacian of a Gaussian should have both positive and negative values
+        assert (finite_values > 0).any(), "No positive Laplacian values found"
+        assert (finite_values < 0).any(), "No negative Laplacian values found"

uxarray/core/dataarray.py

@@ -19,6 +19,7 @@
     _calculate_edge_node_difference,
     _compute_gradient,
 )
+from uxarray.core.vector_calculus import _calculate_divergence
 from uxarray.core.utils import _map_dims_to_ugrid
 from uxarray.core.zonal import _compute_non_conservative_zonal_mean
 from uxarray.cross_sections import UxDataArrayCrossSectionAccessor
@@ -1266,6 +1267,81 @@ def gradient(self, **kwargs) -> UxDataset:
             coords=self.coords,
         )
 
+    def divergence(self, other: "UxDataArray", **kwargs) -> "UxDataArray":
+        """
+        Computes the divergence of the vector field defined by this UxDataArray and other.
+
+        Parameters
+        ----------
+        other : UxDataArray
+            The second component of the vector field. This UxDataArray represents the first component.
+        **kwargs
+            Additional keyword arguments (reserved for future use).
+
+        Returns
+        -------
+        divergence : UxDataArray
+            UxDataArray containing the divergence of the vector field.
+
+        Notes
+        -----
+        The divergence is computed using the finite volume method. For a vector field V = (u, v),
+        where u and v are the components represented by this UxDataArray and other respectively,
+        the divergence is calculated as div(V) = ∂u/∂x + ∂v/∂y.
+
+        The implementation uses edge-centered gradients and face-centered divergence calculation
+        following the discrete divergence theorem.
+
+        Example
+        -------
+        >>> u_component = uxds["u_wind"]  # First component of vector field
+        >>> v_component = uxds["v_wind"]  # Second component of vector field
+        >>> div_field = u_component.divergence(v_component)
+        """
+        from uxarray import UxDataArray
+
+        if not isinstance(other, UxDataArray):
+            raise TypeError("other must be a UxDataArray")
+
+        if self.uxgrid != other.uxgrid:
+            raise ValueError("Both UxDataArrays must have the same grid")
+
+        if self.dims != other.dims:
+            raise ValueError("Both UxDataArrays must have the same dimensions")
+
+        if self.ndim > 1:
+            raise ValueError(
+                "Divergence currently requires 1D face-centered data. Consider "
+                "reducing the dimension by selecting data across leading dimensions (e.g., `.isel(time=0)`, "
+                "`.sel(lev=500)`, or `.mean('time')`)."
+            )
+
+        if not (self._face_centered() and other._face_centered()):
+            raise ValueError(
+                "Computing the divergence is only supported for face-centered data variables."
+            )
+
+        # Compute gradients of both components
+        u_grad_zonal, u_grad_meridional = _compute_gradient(self)
+        v_grad_zonal, v_grad_meridional = _compute_gradient(other)
+
+        # 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_grad_zonal.values + v_grad_meridional.values
+
+        # Create the divergence UxDataArray
+        divergence_da = UxDataArray(
+            data=divergence_values,
+            name="divergence",
+            dims=self.dims,
+            uxgrid=self.uxgrid,
+            attrs={"divergence": True, "units": "1/s" if "units" not in kwargs else kwargs["units"]},
+            coords=self.coords,
+        )
+
+        return divergence_da
+
     def difference(self, destination: Optional[str] = "edge"):
         """Computes the absolute difference of a data variable.