diff --git a/docs/api.rst b/docs/api.rst index 6b3986cbb..71b947304 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -514,6 +514,7 @@ Mathematical Operators UxDataArray.divergence UxDataArray.gradient UxDataArray.integrate + UxDataArray.scalardotgradient Aggregations diff --git a/docs/user-guide/vector_calculus.ipynb b/docs/user-guide/vector_calculus.ipynb index af411e816..8e47f8078 100644 --- a/docs/user-guide/vector_calculus.ipynb +++ b/docs/user-guide/vector_calculus.ipynb @@ -20,7 +20,8 @@ "1. Gradient of face-centered scalar fields (zonal and meridional components)\n", "2. Curl of vector fields (including curl of a gradient ≈ 0 and a synthetic vortex)\n", "3. Divergence of vector fields (including Laplacian as ∇·∇φ, divergence of a vortex ≈ 0, and radial expansion > 0)\n", - "4. Vector calculus identities verification\n" + "4. Scalar dot gradient (advection term **v** · ∇q)\n", + "5. Vector calculus identities verification\n" ] }, { @@ -681,12 +682,78 @@ "(radial_vectors + radial_div).opts(shared_axes=False)" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Scalar Dot Gradient (**v** · ∇q)\n", + "\n", + "### Background\n", + "\n", + "Many geophysical applications need the projection of a vector field **v** = (u, v) onto the gradient of a scalar field q. This quantity appears, for example, as the horizontal advection term in transport equations:\n", + "\n", + "$$\n", + "\\mathbf{v} \\cdot \\nabla q = u\\,\\frac{\\partial q}{\\partial x} + v\\,\\frac{\\partial q}{\\partial y}\n", + "$$\n", + "\n", + "### Usage\n", + "\n", + "The `UxDataArray.scalardotgradient(v, q)` method computes this dot product directly. The data variable it is called on (`self`) is treated as the **zonal** component `u`, `v` is the **meridional** component, and `q` is the scalar field whose gradient is taken. All three inputs must be face-centered and share the same grid and dimensions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Treat the Gaussian and its inverse as a vector field (u, v)\n", + "u_field = uxds[\"gaussian\"]\n", + "v_field = uxds[\"inverse_gaussian\"]\n", + "\n", + "# Scalar field whose gradient is advected\n", + "q_field = uxds[\"gaussian\"]\n", + "\n", + "# v . grad(q) = u * dq/dx + v * dq/dy\n", + "vdotgradq = u_field.scalardotgradient(v_field, q_field)\n", + "vdotgradq" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verify against an explicit gradient-based computation\n", + "grad_q = q_field.gradient()\n", + "expected = u_field * grad_q[\"zonal_gradient\"] + v_field * grad_q[\"meridional_gradient\"]\n", + "\n", + "# Boundary faces yield NaN gradients, so compare only finite values\n", + "finite = np.isfinite(vdotgradq.values) & np.isfinite(expected.values)\n", + "max_diff = np.abs(vdotgradq.values[finite] - expected.values[finite]).max()\n", + "print(f\"Max difference vs. explicit u*dq/dx + v*dq/dy: {max_diff:.2e}\")\n", + "print(f\"Matches explicit computation: {max_diff < 1e-12}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Visualize the scalar dot gradient field\n", + "vdotgradq.plot(cmap=\"RdBu_r\", aspect=1).opts(\n", + " title=\"Scalar Dot Gradient v · ∇q\", colorbar=True\n", + ")" + ] + }, { "cell_type": "markdown", "id": "35", "metadata": {}, "source": [ - "## 4. Vector Calculus Identities\n", + "## 5. Vector Calculus Identities\n", "\n", "Let's verify some fundamental vector calculus identities using our computed fields:\n", "\n", diff --git a/test/core/test_vector_calculus.py b/test/core/test_vector_calculus.py index a991a258b..e6009dafc 100644 --- a/test/core/test_vector_calculus.py +++ b/test/core/test_vector_calculus.py @@ -193,6 +193,79 @@ def test_divergence_basic(self, gridpath, datasetpath): assert np.isfinite(div_field.values).any() +class TestScalarDotGradientMPASOcean: + + def test_scalardotgradient_uses_known_gradient_components( + self, gridpath, datasetpath, monkeypatch + ): + """Test scalar dot gradient against independently supplied gradients.""" + uxds = ux.open_dataset( + gridpath("mpas", "QU", "480", "grid.nc"), + datasetpath("mpas", "QU", "480", "data.nc"), + ) + + n_face = uxds.uxgrid.n_face + dims = ["n_face"] + scalar = ux.UxDataArray( + np.zeros(n_face), dims=dims, uxgrid=uxds.uxgrid, name="scalar" + ) + u_component = ux.UxDataArray( + np.full(n_face, 2.0), dims=dims, uxgrid=uxds.uxgrid, name="u" + ) + v_component = ux.UxDataArray( + np.full(n_face, -0.5), dims=dims, uxgrid=uxds.uxgrid, name="v" + ) + + def mock_gradient(self): + return ux.UxDataset( + { + "zonal_gradient": ux.UxDataArray( + np.full(n_face, 3.0), dims=dims, uxgrid=self.uxgrid + ), + "meridional_gradient": ux.UxDataArray( + np.full(n_face, -4.0), dims=dims, uxgrid=self.uxgrid + ), + }, + uxgrid=self.uxgrid, + ) + + monkeypatch.setattr(ux.UxDataArray, "gradient", mock_gradient) + + result = u_component.scalardotgradient(v_component, scalar) + + expected = np.full(n_face, 8.0) + nt.assert_allclose(result.values, expected, rtol=0.0, atol=0.0) + + assert isinstance(result, ux.UxDataArray) + assert result.name == "scalar_dot_gradient" + assert result.attrs["long_name"] == "scalar dot gradient" + assert result.sizes == u_component.sizes + + def test_scalardotgradient_rejects_misaligned_indexes(self, gridpath, datasetpath): + """Test scalar dot gradient fails instead of silently realigning faces.""" + uxds = ux.open_dataset( + gridpath("mpas", "QU", "480", "grid.nc"), + datasetpath("mpas", "QU", "480", "data.nc"), + ) + + scalar = uxds["bottomDepth"] + u_component = ux.UxDataArray( + np.ones(scalar.size), + dims=["n_face"], + coords={"n_face": np.arange(scalar.size)}, + uxgrid=uxds.uxgrid, + ) + v_component = ux.UxDataArray( + np.ones(scalar.size), + dims=["n_face"], + coords={"n_face": np.arange(scalar.size) + 1}, + uxgrid=uxds.uxgrid, + ) + + with pytest.raises(ValueError): + u_component.scalardotgradient(v_component, scalar) + + class TestDivergenceDyamondSubset: def test_divergence_constant_field(self, gridpath, datasetpath): diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index b75199cc7..46141384b 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1614,6 +1614,65 @@ def divergence(self, other: "UxDataArray", **kwargs) -> "UxDataArray": return divergence_da + def scalardotgradient(self, v: "UxDataArray", q: "UxDataArray") -> "UxDataArray": + """ + Compute the dot product between a vector field and the gradient of a scalar field. + + Parameters + ---------- + v : UxDataArray + The meridional component of the vector field. ``self`` is treated as + the zonal component. + q : UxDataArray + Scalar field whose gradient is dotted with the vector field. + + Returns + ------- + scalar_dot_gradient : UxDataArray + Dot product ``self * dq/dx + v * dq/dy``. + """ + if not isinstance(v, UxDataArray): + raise TypeError("v must be a UxDataArray") + + if not isinstance(q, UxDataArray): + raise TypeError("q must be a UxDataArray") + + if self.uxgrid != v.uxgrid or self.uxgrid != q.uxgrid: + raise ValueError("All UxDataArrays must have the same grid") + + if self.dims != v.dims or self.dims != q.dims: + raise ValueError("All UxDataArrays must have the same dimensions") + + if self.ndim > 1: + raise ValueError( + "Scalar dot gradient currently requires 1D face-centered data. " + "Consider selecting a single slice before computing." + ) + + if not (self._face_centered() and v._face_centered() and q._face_centered()): + raise ValueError( + "Computing the scalar dot gradient is only supported for face-centered data variables." + ) + + # Validate coordinate alignment up-front so a misaligned input fails + # before the (potentially expensive) gradient call. + u_aligned, v_aligned, q_aligned = xr.align(self, v, q, join="exact", copy=False) + + q_gradient = q_aligned.gradient() + q_zonal = q_gradient["zonal_gradient"] + q_meridional = q_gradient["meridional_gradient"] + + scalar_dot_gradient = (u_aligned * q_zonal) + (v_aligned * q_meridional) + scalar_dot_gradient.name = "scalar_dot_gradient" + scalar_dot_gradient.attrs.update( + { + "long_name": "scalar dot gradient", + "description": "Dot product u * (dq/dx) + v * (dq/dy).", + } + ) + + return UxDataArray(scalar_dot_gradient, uxgrid=self.uxgrid) + def difference(self, destination: str | None = "edge"): """Computes the absolute difference of a data variable.