From 235249cc2ef6eedbfbc5a20a897735f8ee9de9c6 Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Wed, 4 Nov 2026 12:07:06 -0600 Subject: [PATCH 1/2] Fix calculus units for unit-sphere derivatives --- docs/user-guide/vector_calculus.ipynb | 50 ++++++++++++++++++++- test/core/test_calculus_units.py | 65 +++++++++++++++++++++++++++ uxarray/core/gradient.py | 6 +-- 3 files changed, 117 insertions(+), 4 deletions(-) create mode 100644 test/core/test_calculus_units.py diff --git a/docs/user-guide/vector_calculus.ipynb b/docs/user-guide/vector_calculus.ipynb index 78b74f1a5..c46dfd6b0 100644 --- a/docs/user-guide/vector_calculus.ipynb +++ b/docs/user-guide/vector_calculus.ipynb @@ -131,7 +131,55 @@ "id": "66caacb5", "metadata": {}, "source": [ - "> **Units.** By default, `gradient()`, `curl()`, and `divergence()` divide by `uxgrid.sphere_radius` so derivatives carry physical units (e.g. `[data units]/m`, `1/s` for velocity curl). Pass `scale_by_radius=False` to keep results on the unit sphere (per radian). If the grid has no `sphere_radius` attribute, the call falls back to unit-sphere output and emits a `UserWarning`." + "### Units and the unit-sphere convention\n", + "\n", + "UXarray computes spherical geometry and derivatives internally on a unit sphere. `uxgrid.sphere_radius` stores the physical radius for optional output scaling; it does not change that internal convention.\n", + "\n", + "For `gradient()`, `curl()`, and `divergence()`:\n", + "\n", + "- By default (`scale_by_radius=True`), a radius other than `1.0` divides the unit-sphere result by that radius, with output units `[data units]/m`. The physical radius is interpreted in meters.\n", + "- With `sphere_radius=1.0`, results retain per-radian units (`[data units]/rad`), even when `scale_by_radius=True`.\n", + "- With `scale_by_radius=False`, results always remain on the unit sphere, regardless of the stored radius.\n", + "- If the grid has no `sphere_radius` attribute, the default call falls back to unit-sphere output and emits a `UserWarning`.\n", + "\n", + "Leading and trailing whitespace is stripped from input units. Missing or blank units become `1`; ratio units are parenthesized, for example `(m/s)/m` (dimensionally equivalent to `1/s`), without symbolic simplification.\n", + "\n", + "The following example reuses the Gaussian field loaded above. A separate grid copy keeps the physical radius used by the rest of this notebook unchanged." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "calculus-units-example", + "metadata": {}, + "outputs": [], + "source": [ + "units_grid = uxds.uxgrid.copy()\n", + "q_units = ux.UxDataArray(\n", + " uxds[\"gaussian\"].values,\n", + " dims=[\"n_face\"],\n", + " uxgrid=units_grid,\n", + " attrs={\"units\": \" K s-1 \"},\n", + ")\n", + "\n", + "units_grid.sphere_radius = 1.0\n", + "grad_unit = q_units.gradient()[\"meridional_gradient\"]\n", + "assert grad_unit.attrs[\"units\"] == \"K s-1/rad\"\n", + "print(\"Radius 1:\", grad_unit.attrs[\"units\"])\n", + "\n", + "units_grid.sphere_radius = uxds.uxgrid.sphere_radius\n", + "grad_physical = q_units.gradient()[\"meridional_gradient\"]\n", + "grad_unscaled = q_units.gradient(scale_by_radius=False)[\"meridional_gradient\"]\n", + "assert grad_physical.attrs[\"units\"] == \"K s-1/m\"\n", + "assert grad_unscaled.attrs[\"units\"] == \"K s-1/rad\"\n", + "np.testing.assert_allclose(\n", + " grad_physical.values,\n", + " grad_unit.values / units_grid.sphere_radius,\n", + " equal_nan=True,\n", + ")\n", + "np.testing.assert_allclose(grad_unscaled.values, grad_unit.values, equal_nan=True)\n", + "print(\"Physical radius:\", grad_physical.attrs[\"units\"])\n", + "print(\"Scaling disabled:\", grad_unscaled.attrs[\"units\"])" ] }, { diff --git a/test/core/test_calculus_units.py b/test/core/test_calculus_units.py new file mode 100644 index 000000000..7da22a568 --- /dev/null +++ b/test/core/test_calculus_units.py @@ -0,0 +1,65 @@ +"""Regression coverage for calculus units and unit-sphere scaling (#1856).""" + +from contextlib import nullcontext + +import numpy as np +import numpy.testing as nt +import pytest + +import uxarray as ux + + +@pytest.mark.parametrize("method", ["gradient", "curl", "divergence"]) +@pytest.mark.parametrize("radius", [None, 1.0, 1.0 + 1e-12, 2.0]) +@pytest.mark.parametrize("scale_by_radius", [False, True]) +@pytest.mark.parametrize( + "units,base_units", + [(None, "1"), (" \t ", "1"), (" K s-1 ", "K s-1"), (" m/s \n", "(m/s)")], +) +def test_calculus_units_and_radius_scaling( + gridpath, datasetpath, method, radius, scale_by_radius, units, base_units +): + uxds = ux.open_dataset( + gridpath("mpas", "dyamond-30km", "gradient_grid_subset.nc"), + datasetpath("mpas", "dyamond-30km", "gradient_data_subset.nc"), + ) + if radius is None: + uxds.uxgrid._ds.attrs.pop("sphere_radius", None) + else: + uxds.uxgrid.sphere_radius = radius + + u = uxds["gaussian"].copy() + v = uxds["inverse_gaussian"].copy() + for component in (u, v): + component.attrs.pop("units", None) + if units is not None: + component.attrs["units"] = units + original_attrs = u.attrs.copy() + + def compute(scale): + if method == "gradient": + result = u.gradient(scale_by_radius=scale) + return [result["zonal_gradient"], result["meridional_gradient"]] + return [getattr(u, method)(v, scale_by_radius=scale)] + + unit_results = compute(False) + warning = ( + pytest.warns(UserWarning, match="result is left on the unit sphere") + if scale_by_radius and radius is None + else nullcontext() + ) + with warning: + results = compute(scale_by_radius) + + physically_scaled = scale_by_radius and radius is not None and radius != 1.0 + denominator = "m" if physically_scaled else "rad" + divisor = radius if physically_scaled else 1.0 + for result, unit_result in zip(results, unit_results): + assert result.attrs["units"] == f"{base_units}/{denominator}" + assert unit_result.attrs["units"] == f"{base_units}/rad" + finite = np.isfinite(unit_result.values) + assert finite.any() + assert np.any(np.abs(unit_result.values[finite]) > 0) + nt.assert_allclose(result.values, unit_result.values / divisor, equal_nan=True) + + assert u.attrs == original_attrs diff --git a/uxarray/core/gradient.py b/uxarray/core/gradient.py index cf14cb06c..5cbe8bdb8 100644 --- a/uxarray/core/gradient.py +++ b/uxarray/core/gradient.py @@ -231,12 +231,12 @@ def _compute_gradient( stacklevel=2, ) - base_units = str(data.attrs.get("units", "1")) + base_units = str(data.attrs.get("units", "1")).strip() if "/" in base_units: # put in (), e.g. (m/s)/m, not m/s/m. base_units = f"({base_units})" - if base_units.strip() == "": + if base_units == "": base_units = "1" # --> e.g., want result like "1/m", not "/m" or " /m" - was_scaled = scale_by_radius and has_sphere_radius + was_scaled = scale_by_radius and has_sphere_radius and uxgrid.sphere_radius != 1.0 denom_units = "m" if was_scaled else "rad" grad_units = f"{base_units}/{denom_units}" From c0d51e026d99562df39f524c82bff88f3c025aeb Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Tue, 10 Nov 2026 14:06:11 -0600 Subject: [PATCH 2/2] Fix docstring typo in gradient-based method docs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- uxarray/core/dataarray.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index a10f4ab1c..914d0731c 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1649,7 +1649,7 @@ def gradient(self, scale_by_radius: bool = True) -> UxDataset: the result carries physical units: [u units]/meter (e.g. result has units of 1/s if u represents velocity). When ``False`` (or if the grid's ``sphere_radius`` is set to 1.0), - the result caries per-radian units: [u units]/radian. + the result carries per-radian units: [u units]/radian. (e.g. result has units of (m/s)/rad if u represents velocity). CAUTION: this scaling does not account for any radius increases @@ -1839,7 +1839,7 @@ def divergence( the result carries physical units: [u units]/meter (e.g. result has units of 1/s if u represents velocity). When ``False`` (or if the grid's ``sphere_radius`` is set to 1.0), - the result caries per-radian units: [u units]/radian. + the result carries per-radian units: [u units]/radian. (e.g. result has units of (m/s)/rad if u represents velocity). CAUTION: this scaling does not account for any radius increases