From 40f2b7f9395762af91a0b10e756efae6bc5d8a80 Mon Sep 17 00:00:00 2001 From: nstarman Date: Mon, 17 Aug 2026 10:02:03 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(quantity):=20`=5F=5Farray=5F?= =?UTF-8?q?=5F`=20must=20not=20silently=20strip=20units?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Quantity.__array__` returned the bare value in whatever unit the quantity happened to carry, so a consumer received a number whose meaning depended on a unit it never saw: np.asarray(Q(1.5, "km")) -> 1.5 np.asarray(Q(1500, "m")) -> 1500 # the same length np.asarray(Q(90.0, "deg")) -> 90.0 # 57x, to a radian consumer `__array__` is reached implicitly, which is what makes this dangerous rather than merely surprising. `np.asarray` has always used it, and `jax.numpy.asarray` began doing so in jax 0.10 -- it raised `TypeError` before. Isolated: unxt 2.0.0 with jax 0.10 strips; unxt 2.0.1 with jax 0.7.2 raises. So the trigger is the jax upgrade, but the cause is here. Refuse a dimensionful quantity with `UnitConversionError`, naming `ustrip` as the explicit route. Dimensionless still converts -- there is nothing to lose. `UnitConversionError` rather than `TypeError` deliberately: callers write `try: np.asarray(x) except TypeError` and fall back, which would swallow this into a different silent path. unxt already knew the hazard: `base.py` documents `__array__` dropping units for astropy operands and works around it there. This closes the general case. One real consumer relied on the strip. `Dataset.unxt.quantify()` attached a Quantity to a *dimension* coordinate and let xarray's `pandas.Index` coercion discard it -- exactly the anti-pattern. It now skips dimension coordinates explicitly: same result, chosen rather than left to `__array__`. The guide's "Dimension Coordinates Cannot Hold Quantities" section documented the silent loss and now documents the refusal. Suggest backporting to versions/v2.0.x: the previous behaviour was the defect, not an interface anyone could correctly depend on. Co-Authored-By: Claude Opus 5 --- .../unxts.interop.xarray/docs/xarray-guide.md | 14 ++++--- .../unxts/interop/xarray/_src/conversion.py | 22 +++++++++- src/unxt/_src/quantity/mixins.py | 40 +++++++++++++++++-- tests/unit/test_quantity.py | 20 +++++++++- 4 files changed, 84 insertions(+), 12 deletions(-) diff --git a/packages/unxts.interop.xarray/docs/xarray-guide.md b/packages/unxts.interop.xarray/docs/xarray-guide.md index ccfa864b..77d16b97 100644 --- a/packages/unxts.interop.xarray/docs/xarray-guide.md +++ b/packages/unxts.interop.xarray/docs/xarray-guide.md @@ -601,7 +601,9 @@ print(stripped.data) ### Dimension Coordinates Cannot Hold Quantities -`xarray` backs every _dimension coordinate_ (one named like its dimension, shown with a `*` in the repr) with a `pandas.Index`. Building that index coerces the data to a plain `numpy` array, so a `Quantity` assigned to a dimension coordinate is silently unwrapped — its unit is lost. This is inherent to `xarray`'s indexing model, not something `unxts.interop.xarray` can override, and it affects every duck-array unit library (including `pint-xarray`) the same way. +`xarray` backs every _dimension coordinate_ (one named like its dimension, shown with a `*` in the repr) with a `pandas.Index`. Building that index coerces the data to a plain `numpy` array, so a dimension coordinate cannot hold a `Quantity`. This is inherent to `xarray`'s indexing model, not something `unxts.interop.xarray` can override, and it affects every duck-array unit library (including `pint-xarray`) the same way. + +Assigning one now **raises** rather than dropping the unit on the floor — `Quantity.__array__` refuses to hand a dimensionful value to a consumer that cannot see its unit. `quantify()` handles this for you by leaving dimension coordinates plain. **Workaround**: store the unitful values on a _non-dimension_ coordinate, keeping a plain index on the dimension itself: @@ -612,10 +614,12 @@ import xarray as xr data = [10.0, 20.0, 30.0] quantities = u.Quantity([1.0, 2.0, 3.0], "m") -# Dimension coordinate: ``x`` is unwrapped to a plain array, unit lost -da = xr.DataArray(data, dims=["x"], coords={"x": quantities}) -print(type(da.coords["x"].data).__name__) -# ndarray +# Dimension coordinate: refused, because the unit could not survive +try: + xr.DataArray(data, dims=["x"], coords={"x": quantities}) +except Exception as e: + print(type(e).__name__) +# UnitConversionError # Non-dimension coordinate: the Quantity (and its unit) is preserved da = xr.DataArray(data, dims=["i"], coords={"i": [0, 1, 2], "x": ("i", quantities)}) diff --git a/packages/unxts.interop.xarray/src/unxts/interop/xarray/_src/conversion.py b/packages/unxts.interop.xarray/src/unxts/interop/xarray/_src/conversion.py index ea2295a6..30bf3dbc 100644 --- a/packages/unxts.interop.xarray/src/unxts/interop/xarray/_src/conversion.py +++ b/packages/unxts.interop.xarray/src/unxts/interop/xarray/_src/conversion.py @@ -364,9 +364,18 @@ def attach_units(obj: DataArray, units: Mapping) -> DataArray: # DataArray/Dataset constructor keeps copying coordinate attrs -- an # xarray implementation detail, not a guarantee. Copying here makes # non-mutation of the caller's object structural instead. + # A *dimension* coordinate (named like its own dimension) is backed by a + # `pandas.Index`, which xarray builds by coercing to a plain array. A + # Quantity there cannot survive, so do not attach one: skipping is the + # same outcome xarray would reach anyway, but chosen here rather than + # left to `Quantity.__array__` to discard silently. See "Dimension + # Coordinates Cannot Hold Quantities" in the guide. + is_dim_coord = coord.dims == (name,) new_coords[name] = Variable( coord.dims, - coord.data if unit is None else _array_attach_units(coord.data, unit), + coord.data + if unit is None or is_dim_coord + else _array_attach_units(coord.data, unit), dict(coord.attrs), ) @@ -416,9 +425,18 @@ def attach_units(obj: Dataset, units: Mapping) -> Dataset: # DataArray/Dataset constructor keeps copying coordinate attrs -- an # xarray implementation detail, not a guarantee. Copying here makes # non-mutation of the caller's object structural instead. + # A *dimension* coordinate (named like its own dimension) is backed by a + # `pandas.Index`, which xarray builds by coercing to a plain array. A + # Quantity there cannot survive, so do not attach one: skipping is the + # same outcome xarray would reach anyway, but chosen here rather than + # left to `Quantity.__array__` to discard silently. See "Dimension + # Coordinates Cannot Hold Quantities" in the guide. + is_dim_coord = coord.dims == (name,) new_coords[name] = Variable( coord.dims, - coord.data if unit is None else _array_attach_units(coord.data, unit), + coord.data + if unit is None or is_dim_coord + else _array_attach_units(coord.data, unit), dict(coord.attrs), ) diff --git a/src/unxt/_src/quantity/mixins.py b/src/unxt/_src/quantity/mixins.py index 6343779b..e73f62b2 100644 --- a/src/unxt/_src/quantity/mixins.py +++ b/src/unxt/_src/quantity/mixins.py @@ -7,7 +7,11 @@ import equinox as eqx import numpy as np -from astropy.units import CompositeUnit +from astropy.units import ( + CompositeUnit, + UnitConversionError, + dimensionless_unscaled as one, +) from jax.typing import ArrayLike from jaxtyping import Array @@ -212,18 +216,46 @@ class NumPyCompatMixin: __array_namespace__: Callable[[], Any] def __array__(self, *args: object, **kw: object) -> np.ndarray: - """Return the array as a numpy array, stripping the units. + """Return a bare array -- but only where that loses nothing. + + A *dimensionful* quantity has no unambiguous array form. ``np.asarray`` + on ``Quantity(1.5, "km")`` would give ``1.5`` while the same length in + metres gives ``1500``, so the consumer reads a number whose meaning + depends on a unit it never saw. + + This matters because ``__array__`` is reached *implicitly*: ``np.asarray`` + has always used it, and ``jax.numpy.asarray`` does too as of jax 0.10 -- + where it previously raised. Returning the value here therefore turns a + loud failure into a wrong number, which is the one outcome a unit library + must not produce. Use `unxt.ustrip` to say which unit you meant. + + Dimensionless quantities convert, because there is nothing to lose. Examples -------- >>> from unxt import Quantity >>> import numpy as np - >>> q = Quantity(1.01, "m") - >>> np.array(q) + >>> np.array(Quantity(1.01, "")) array(1.01, dtype=float32) + A dimensionful one refuses, and names the way to be explicit: + + >>> try: + ... np.array(Quantity(1.01, "m")) + ... except Exception as e: + ... print(type(e).__name__) + UnitConversionError + """ + if self.unit != one: + msg = ( + f"cannot convert Quantity in {self.unit!r} to a bare array: the " + f"result would be a number whose meaning depends on a unit the " + f"caller never sees. Use `unxt.ustrip(, q)` to choose one, " + f"or `unxt.uconvert('', q)` if it really is dimensionless." + ) + raise UnitConversionError(msg) return np.asarray(uapi.ustrip(self.unit, self), *args, **kw) # TODO: why doesn't `__array_namespace__` supersede this? diff --git a/tests/unit/test_quantity.py b/tests/unit/test_quantity.py index 453ef1a9..f3bdfeba 100644 --- a/tests/unit/test_quantity.py +++ b/tests/unit/test_quantity.py @@ -106,7 +106,9 @@ def test_numpy_array_copy_kwarg_uses_array_protocol(): Without *args/**kwargs in __array__, NumPy's copy parameter would cause TypeError: __array__() got an unexpected keyword argument 'copy'. """ - q = u.Q(1.01, "m") + # Dimensionless: `__array__` refuses a dimensionful quantity, and the + # behaviour under test here is the `copy=` passthrough, not unit handling. + q = u.Q(1.01, "") # NumPy 2.0+ passes copy=True to __array__ by default arr = np.array(q) @@ -115,6 +117,22 @@ def test_numpy_array_copy_kwarg_uses_array_protocol(): assert np.isclose(arr, 1.01) +def test_numpy_array_refuses_a_dimensionful_quantity(): + """A bare array cannot carry a unit, so `__array__` must not invent one. + + Regression guard: `np.asarray` has always reached `__array__` implicitly and + `jax.numpy.asarray` began doing so in jax 0.10, where it previously raised. + Returning the value turns a unit error into a wrong number -- `Q(90, "deg")` + silently becoming `90.0` for a radian consumer is a 57x error no assertion + would notice. + """ + for spec in ("m", "km", "deg", "rad"): + with pytest.raises(UnitConversionError): + np.asarray(u.Q(1.0, spec)) + with pytest.raises(UnitConversionError): + jax_xp.asarray(u.Q(1.0, spec)) + + def test_uconvert(): """Test the ``u.Q.uconvert`` method.""" q = u.Q(1, "m")