diff --git a/dascore/core/coords.py b/dascore/core/coords.py index cf1a2636f..7d79e0608 100644 --- a/dascore/core/coords.py +++ b/dascore/core/coords.py @@ -3,8 +3,9 @@ from __future__ import annotations import abc +import numbers from collections.abc import Sized -from contextlib import suppress +from contextlib import contextmanager, suppress from functools import cache from operator import gt, lt from typing import Any, TypeVar @@ -23,8 +24,9 @@ import dascore as dc from dascore.compat import array, is_array from dascore.constants import _AGG_FUNCS, DIM_REDUCE_DOCS, dascore_styles -from dascore.exceptions import CoordError, ParameterError +from dascore.exceptions import CoordError, ParameterError, UnitError from dascore.units import ( + PintError, Quantity, Unit, convert_units, @@ -58,6 +60,112 @@ step_type = TypeVar("step_type") +# Types which coords know how to combine with in array operations. +_ARRAY_OP_TYPES = (np.ndarray, np.generic, numbers.Number, Quantity, list, tuple) + +# Ufuncs whose operands must all share the same units. These also return +# outputs in those units, unless the output is boolean. +_UNIT_MATCHED_UFUNCS = frozenset( + { + np.add, + np.subtract, + np.mod, + np.fmod, + np.remainder, + np.maximum, + np.minimum, + np.fmax, + np.fmin, + np.hypot, + np.greater, + np.greater_equal, + np.less, + np.less_equal, + np.equal, + np.not_equal, + np.positive, + np.negative, + np.absolute, + np.fabs, + np.rint, + np.floor, + np.ceil, + np.trunc, + } +) + +# Numpy functions which reduce an array to a single value. Time-like coords +# need special handling for these (see _reduce_time_like). +_REDUCING_ARRAY_FUNCS = frozenset( + { + np.mean, + np.nanmean, + np.median, + np.nanmedian, + np.std, + np.nanstd, + np.sum, + np.nansum, + np.min, + np.nanmin, + np.max, + np.nanmax, + } +) + + +def _map_nested(func, obj): + """Apply func to each non-container element of a nested structure.""" + if isinstance(obj, tuple | list): + return type(obj)(_map_nested(func, x) for x in obj) + if isinstance(obj, dict): + return {i: _map_nested(func, v) for i, v in obj.items()} + return func(obj) + + +@contextmanager +def _unit_error_context(func, units): + """Raise a dascore UnitError when pint can't perform an operation.""" + try: + yield + except PintError as ex: + name = getattr(func, "__name__", func) + msg = f"{name} failed for coordinate with units of {units}. {ex}" + raise UnitError(msg) from ex + + +def _to_magnitude(obj, units=None): + """Strip units from a quantity, first converting to units if provided.""" + if not isinstance(obj, Quantity): + return obj + return obj.magnitude if units is None else obj.to(units).magnitude + + +def _wrap_array_op_output(out, units=None): + """ + Convert the output of an array operation back to a coordinate. + + Scalars (eg reductions) keep their units but are not coordinates, and + boolean arrays are left alone since they are masks, not coordinates. + """ + if isinstance(out, tuple | list): # Eg np.divmod or np.array_split. + return type(out)(_wrap_array_op_output(x, units) for x in out) + if isinstance(out, Quantity): + # Dimensionless units can still have a scale (eg m/cm) so the + # magnitude has to be converted before the units are dropped. + if out.units.dimensionless: + out, units = out.to("dimensionless").magnitude, None + else: + out, units = out.magnitude, out.units + if not is_array(out) or np.ndim(out) == 0: + # Time-like values (eg datetime64) can't have units attached. + if units is None or dtype_time_like(np.asarray(out).dtype): + return out + return out * units + if np.issubdtype(out.dtype, np.bool_): + return out + return get_coord(data=out, units=units) + def ensure_consistent_dtype(value, name, dtype): """Ensure the values are consistent with dtype.""" @@ -247,6 +355,14 @@ class BaseCoord(DascoreBaseModel, abc.ABC): Coordinates should usually be created with [get_coords](`dascore.core.coords.get_coord`) rather than using the class directly. + + Notes + ----- + Coordinates support python operators, numpy ufuncs, and numpy functions, + each of which returns a new coordinate whose units reflect the operation + performed. Operations which return a single value (eg `np.mean`) return + a scalar and operations which return booleans (eg `np.greater`) return + arrays. See the [coordinate tutorial](/tutorial/coords.qmd) for details. """ units: UnitQuantity = None @@ -487,6 +603,158 @@ def __array__(self, dtype=None, copy=False): """Numpy method for getting array data with `np.array(coord)`.""" return self.data + def _to_operand(self, obj, units=None): + """ + Convert an input of an array operation to an array or quantity. + + Parameters + ---------- + obj + The object to convert. Coords and quantities are unpacked into + their values and units, anything else is passed through. + units + If provided, the units assumed for values which have none (eg + the 1 in `coord + 1`). + """ + if isinstance(obj, BaseCoord): + data, obj_units = obj.data, obj.units + elif isinstance(obj, Quantity): + data, obj_units = obj.magnitude, obj.units + else: + data, obj_units = obj, None + # Time-like coords operate on raw values; pint knows nothing of + # numpy's datetime64/timedelta64. + if dtype_time_like(self.dtype): + return data + obj_units = obj_units if obj_units is not None else units + return data if obj_units is None else data * obj_units + + def _get_op_units(self, inputs): + """Get the units which apply to operands which have none.""" + if self.units is not None: + return self.units + others = (getattr(x, "units", None) for x in inputs) + return next((x for x in others if x is not None), None) + + def _operate(self, ufunc, *inputs): + """Apply a ufunc, deferring to other types when they aren't known.""" + if any(not isinstance(x, (BaseCoord, *_ARRAY_OP_TYPES)) for x in inputs): + return NotImplemented + return ufunc(*inputs) + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + """Implement numpy's ufunc protocol (eg np.sqrt(coord), coord + 1).""" + # Coord data is read-only, so numpy raises for operations which + # write into it, but ufunc.at ignores the flag and must be rejected. + if method == "at": + msg = ( + "Coordinates are immutable, so ufunc.at is not supported. " + "Apply the operation to the coordinate's values instead." + ) + raise ParameterError(msg) + if any(not isinstance(x, (BaseCoord, *_ARRAY_OP_TYPES)) for x in inputs): + return NotImplemented + # Ufuncs which require operands share units also preserve them. + matched = ufunc in _UNIT_MATCHED_UFUNCS + units = self._get_op_units(inputs) + operands = [self._to_operand(x, units if matched else None) for x in inputs] + with _unit_error_context(ufunc, self.unit_str): + if method != "__call__": + # Pint doesn't implement reduce/accumulate/outer, so units + # are handled here, which only works if they don't change. + if not matched and units is not None: + msg = ( + f"The units resulting from {ufunc.__name__}.{method} " + f"are ambiguous for a coordinate with units of " + f"{self.unit_str}. Use the coordinate's values instead." + ) + raise UnitError(msg) + operands = [_to_magnitude(x, units) for x in operands] + # When operands are quantities pint performs the unit algebra + # (eg m * m -> m ** 2) and raises on invalid ops (eg m + s). + out = getattr(ufunc, method)(*operands, **kwargs) + return _wrap_array_op_output(out, units if matched else None) + + def __array_function__(self, func, types, args, kwargs): + """Implement numpy's array protocol (eg np.concatenate([coord1])).""" + if any(not issubclass(x, (BaseCoord, *_ARRAY_OP_TYPES)) for x in types): + return NotImplemented + # Numpy can't reduce absolute times so dascore's logic is used. + if dtype_time_like(self.dtype) and func in _REDUCING_ARRAY_FUNCS: + out = _reduce_time_like(func, self.data) + return _wrap_array_op_output(out[0] if out.size == 1 else out, self.units) + args = _map_nested(self._to_operand, args) + kwargs = _map_nested(self._to_operand, kwargs) + with _unit_error_context(func, self.unit_str): + out = func(*args, **kwargs) + return _wrap_array_op_output(out) + + def __add__(self, other): + return self._operate(np.add, self, other) + + def __radd__(self, other): + return self._operate(np.add, other, self) + + def __sub__(self, other): + return self._operate(np.subtract, self, other) + + def __rsub__(self, other): + return self._operate(np.subtract, other, self) + + def __mul__(self, other): + return self._operate(np.multiply, self, other) + + def __rmul__(self, other): + return self._operate(np.multiply, other, self) + + def __truediv__(self, other): + return self._operate(np.divide, self, other) + + def __rtruediv__(self, other): + return self._operate(np.divide, other, self) + + def __floordiv__(self, other): + return self._operate(np.floor_divide, self, other) + + def __rfloordiv__(self, other): + return self._operate(np.floor_divide, other, self) + + def __mod__(self, other): + return self._operate(np.mod, self, other) + + def __rmod__(self, other): + return self._operate(np.mod, other, self) + + def __pow__(self, other): + return self._operate(np.power, self, other) + + def __rpow__(self, other): + return self._operate(np.power, other, self) + + def __neg__(self): + return self._operate(np.negative, self) + + def __pos__(self): + return self._operate(np.positive, self) + + def __abs__(self): + return self._operate(np.absolute, self) + + # Note: __eq__ (and __ne__) are not defined here; they compare + # coordinates, not their values, since coords are pydantic models. + + def __gt__(self, other): + return self._operate(np.greater, self, other) + + def __ge__(self, other): + return self._operate(np.greater_equal, self, other) + + def __lt__(self, other): + return self._operate(np.less, self, other) + + def __le__(self, other): + return self._operate(np.less_equal, self, other) + @cached_method def min(self): """Return min value.""" diff --git a/dascore/units.py b/dascore/units.py index 8ce1e35f7..1145685ff 100644 --- a/dascore/units.py +++ b/dascore/units.py @@ -11,6 +11,9 @@ import pandas as pd import pint from pint import DimensionalityError, Quantity, UndefinedUnitError, Unit + +# Re-exported so dascore code can catch any error raised by pint. +from pint.errors import PintError # noqa: F401 from platformdirs import user_cache_path import dascore as dc diff --git a/docs/tutorial/coords.qmd b/docs/tutorial/coords.qmd index 31cc8ffea..28825c31d 100644 --- a/docs/tutorial/coords.qmd +++ b/docs/tutorial/coords.qmd @@ -208,6 +208,55 @@ assert coord.get_next_index(1) == 1 assert coord.get_next_index(2.000001) == 3 ``` +## Coordinate Arithmetic + +Coordinates behave much like numpy arrays; python operators, numpy ufuncs, and +numpy functions all work on them and return new coordinates. Units are tracked +through the operation, so multiplying two coordinates with units of `m` returns +a coordinate with units of `m**2`. + +```{python} +import numpy as np + +from dascore.core import get_coord + +coord = get_coord(data=np.array([1., 4., 9.]), units="m") + +# Values without units are assumed to be in the coordinate's units. +assert (coord + 1).units == coord.units + +# But units are otherwise tracked through the operation. +assert (coord * coord).unit_str == "m ** 2" +assert np.sqrt(coord).unit_str == "m ** 0.5" + +# Coordinates (or quantities) with other units are first converted. +from dascore.units import get_quantity + +assert np.allclose((coord + 100 * get_quantity("cm")).values, coord.values + 1) +``` + +Two exceptions to the "operations return coordinates" rule keep results useful: +operations which return a single value (eg `np.mean`) return a scalar, and +operations which return booleans (eg `np.greater`) return arrays which can be +used to index other arrays. + +```{python} +# Reductions return a scalar (with units, when the coord has them). +print(np.mean(coord)) + +# Comparisons return boolean arrays. +print(np.greater(coord, 2 * get_quantity("m"))) +``` + +Comparison operators work as well, so `coord > 2 * get_quantity("m")` is the +same as the call above. + +Operations which don't make sense for the coordinate's units, such as adding +seconds to meters, raise a [`UnitError`](`dascore.exceptions.UnitError`). So do +ufunc methods whose units can't be determined, such as `np.multiply.reduce`. +Coordinate values are read-only, so numpy raises a `ValueError` for operations +which would write into them. + # CoordManager The [`CoordManager`](`dascore.core.coordmanager.CoordManager`) handles a group of coordinates and provides methods for updating managed data arrays. diff --git a/tests/test_core/test_coords.py b/tests/test_core/test_coords.py index 9031795c6..b72c3eb2c 100644 --- a/tests/test_core/test_coords.py +++ b/tests/test_core/test_coords.py @@ -26,7 +26,7 @@ CoordSummary, get_coord, ) -from dascore.exceptions import CoordError, ParameterError +from dascore.exceptions import CoordError, ParameterError, UnitError from dascore.units import get_quantity, percent from dascore.utils.misc import all_close, register_func from dascore.utils.time import dtype_time_like, is_datetime64, is_timedelta64, to_float @@ -2006,6 +2006,268 @@ def test_not_implemented_in_baseclass(self, evenly_sampled_coord): BaseCoord.change_length(coord, 10) +class TestCoordinateArithmetic: + """Tests for treating coordinates like arrays (see issue #566).""" + + @pytest.fixture() + def coord(self): + """A simple coordinate with units for testing arithmetic.""" + return get_coord(data=np.array([2.0, 4.0, 8.0]), units="m") + + @pytest.fixture() + def unitless_coord(self, coord): + """The same coordinate without units.""" + return get_coord(data=coord.values) + + @pytest.mark.parametrize( + "op", + [ + lambda x, y: x + y, + lambda x, y: x - y, + lambda x, y: x * y, + lambda x, y: x / y, + lambda x, y: x // y, + lambda x, y: x % y, + lambda x, y: x**y, + ], + ) + @pytest.mark.parametrize("reverse", [False, True]) + def test_binary_ops_match_arrays(self, unitless_coord, op, reverse): + """Binary operators should return coords with the array's values.""" + coord, other = unitless_coord, 3.0 + args = (other, coord) if reverse else (coord, other) + out = op(*args) + expected = op(*((other, coord.values) if reverse else (coord.values, other))) + assert isinstance(out, BaseCoord) + np.testing.assert_allclose(out.values, expected) + + @pytest.mark.parametrize( + "op", [lambda x: -x, lambda x: +x, lambda x: abs(x), np.sign] + ) + def test_unary_ops(self, coord, op): + """Unary operators should also return coordinates.""" + out = op(coord) + assert isinstance(out, BaseCoord) + np.testing.assert_allclose(out.values, op(coord.values)) + + def test_scalars_assume_coord_units(self, coord): + """Scalars added to a coord are assumed to have the coord's units.""" + out = coord + 1 + assert out.units == coord.units + np.testing.assert_allclose(out.values, coord.values + 1) + + def test_units_converted(self, coord): + """Quantities/coords with other units should be converted.""" + out = coord + 100 * get_quantity("cm") + assert out.units == coord.units + np.testing.assert_allclose(out.values, coord.values + 1) + out = coord + get_coord(data=coord.values * 100, units="cm") + np.testing.assert_allclose(out.values, coord.values * 2) + + def test_incompatible_units_raise(self, coord): + """Adding coords with incompatible units should raise.""" + with pytest.raises(UnitError): + coord + get_coord(data=coord.values, units="s") + + @pytest.mark.parametrize( + "op, units", + [ + (lambda x: x * x, "m ** 2"), + (lambda x: x**2, "m ** 2"), + (lambda x: np.sqrt(x), "m ** 0.5"), + (lambda x: x / 2, "m"), + (lambda x: x + x, "m"), + (lambda x: 1 / x, "1 / m"), + ], + ) + def test_units_track_operation(self, coord, op, units): + """Units should reflect the operation performed, not simply persist.""" + out = op(coord) + assert out.units == get_quantity(units) + + def test_dimensionless_output_drops_units(self, coord): + """Operations which cancel units should return a unitless coord.""" + out = coord / coord + assert isinstance(out, BaseCoord) + assert out.units is None + + @pytest.mark.parametrize("op", [np.exp, np.sin, lambda x: 2**x]) + def test_ops_requiring_dimensionless_raise(self, coord, op): + """Ops which need dimensionless inputs raise for coords with units.""" + with pytest.raises(UnitError): + op(coord) + + def test_no_units_coord(self): + """Coords without units should behave like plain arrays.""" + coord = get_coord(data=np.array([1.0, 2.0, 3.0])) + out = np.exp(coord * 2) + assert isinstance(out, BaseCoord) + assert out.units is None + np.testing.assert_allclose(out.values, np.exp(coord.values * 2)) + + def test_time_coords(self, random_patch): + """Time coords should support arithmetic with timedeltas.""" + coord = random_patch.get_coord("time") + out = coord + dc.to_timedelta64(1) + assert isinstance(out, BaseCoord) + assert np.all(out.values == coord.values + dc.to_timedelta64(1)) + # Subtracting two time coords should yield a timedelta coord. + diff = coord - coord.min() + assert dtype_time_like(diff.dtype) + + def test_array_operands(self, coord): + """Arrays should work on either side of the operation.""" + array = np.array([1.0, 2.0, 3.0]) + out1, out2 = coord + array, array + coord + assert isinstance(out1, BaseCoord) and isinstance(out2, BaseCoord) + np.testing.assert_allclose(out1.values, out2.values) + + def test_comparison_ufunc_returns_array(self, coord): + """Boolean output are masks, not coordinates.""" + out = np.greater(coord, 4 * get_quantity("m")) + assert isinstance(out, np.ndarray) + np.testing.assert_array_equal(out, coord.values > 4) + + def test_ufunc_method(self, coord): + """Ufunc methods (eg accumulate) should also return coords.""" + out = np.add.accumulate(coord) + assert isinstance(out, BaseCoord) + assert out.units == coord.units + np.testing.assert_allclose(out.values, np.add.accumulate(coord.values)) + + def test_reduction_returns_quantity(self, coord): + """Reductions to scalars return quantities rather than coords.""" + out = np.linalg.norm(coord) + assert not isinstance(out, BaseCoord) + assert out.units == coord.units + assert np.isclose(out.magnitude, np.linalg.norm(coord.values)) + + def test_array_function_nested_inputs(self, coord): + """Array functions should handle coords nested in args and kwargs.""" + expected = np.concatenate([coord.values, coord.values]) + for arg in (tuple([coord, coord]), [coord, coord]): + out = np.concatenate(arg) + assert isinstance(out, BaseCoord) + assert out.units == coord.units + np.testing.assert_allclose(out.values, expected) + out = np.mean(a=coord) + assert np.isclose(out.magnitude, np.mean(coord.values)) + + def test_unsupported_type_raises_type_error(self, coord): + """Unsupported operands should raise a normal TypeError.""" + with pytest.raises(TypeError): + coord + "bob" + assert coord._operate(np.add, coord, "bob") is NotImplemented + + def test_multiple_outputs(self, unitless_coord): + """Ufuncs with multiple outputs should return multiple coords.""" + out = np.divmod(unitless_coord, 2) + assert isinstance(out, tuple) and len(out) == 2 + assert all(isinstance(x, BaseCoord) for x in out) + expected = np.divmod(unitless_coord.values, 2) + for coord_out, array_out in zip(out, expected): + np.testing.assert_allclose(coord_out.values, array_out) + + def test_unsupported_op_with_units_raises(self, coord): + """Ops pint doesn't support should still raise a sensible error.""" + with pytest.raises(TypeError): + np.divmod(coord, 2) + + def test_unitless_reduction(self, unitless_coord): + """Reductions of unitless coords should return plain scalars.""" + out = np.mean(unitless_coord) + assert np.isclose(out, np.mean(unitless_coord.values)) + + def test_ufunc_unknown_type_defers(self, coord): + """The ufunc protocol should defer on types it doesn't know.""" + out = coord.__array_ufunc__(np.add, "__call__", coord, "bob") + assert out is NotImplemented + + def test_units_of_other_operands_used(self, coord, unitless_coord): + """A coord without units should not discard other operands' units.""" + out1, out2 = unitless_coord * coord, coord * unitless_coord + assert out1.units == out2.units == get_quantity("m") + np.testing.assert_allclose(out1.values, out2.values) + # Values without units are then assumed to be in the other's units. + assert (unitless_coord + coord).units == get_quantity("m") + + def test_scaled_dimensionless_output(self, coord): + """Units which cancel but have a scale should scale the values.""" + other = get_coord(data=coord.values * 100, units="cm") + out = coord / other + assert out.units is None + np.testing.assert_allclose(out.values, np.ones_like(coord.values)) + + def test_ufunc_method_converts_units(self, coord): + """Ufunc methods should convert units rather than ignore them.""" + other = get_coord(data=coord.values * 100, units="cm") + out = np.add.outer(coord, other) + assert out.units == coord.units + np.testing.assert_allclose(out.values, np.add.outer(*[coord.values] * 2)) + + def test_ambiguous_ufunc_method_units_raise(self, coord): + """Ufunc methods which change units aren't supported with units.""" + with pytest.raises(UnitError, match="ambiguous"): + np.multiply.reduce(coord) + # But they work fine when there are no units to get wrong. + out = np.multiply.reduce(get_coord(data=coord.values)) + assert np.isclose(out, np.multiply.reduce(coord.values)) + + def test_ufunc_at_raises(self, coord): + """Ufunc.at ignores numpy's read-only flag so it must be rejected.""" + with pytest.raises(ParameterError, match="immutable"): + np.add.at(coord, [0], 1) + np.testing.assert_allclose(coord.values, [2.0, 4.0, 8.0]) + + @pytest.mark.parametrize( + "op", + [ + lambda x: np.copyto(x, 0), + lambda x: np.put(x, [0], 1), + lambda x: np.add(x, 1, out=x.values), + ], + ) + def test_operations_writing_to_coord_raise(self, unitless_coord, op): + """Coord data is read-only, so numpy rejects writing into it.""" + with pytest.raises(ValueError, match="read-only"): + op(unitless_coord) + np.testing.assert_allclose(unitless_coord.values, [2.0, 4.0, 8.0]) + + def test_array_function_returning_list(self, unitless_coord): + """Array functions which return lists should return coords.""" + out = np.array_split(unitless_coord, 3) + assert isinstance(out, list) + assert all(isinstance(x, BaseCoord) for x in out) + + @pytest.mark.parametrize( + "op", + [ + lambda x, y: x > y, + lambda x, y: x >= y, + lambda x, y: x < y, + lambda x, y: x <= y, + ], + ) + def test_comparisons_return_masks(self, coord, op): + """Comparison operators should return boolean arrays.""" + out = op(coord, 4 * get_quantity("m")) + assert isinstance(out, np.ndarray) + np.testing.assert_array_equal(out, op(coord.values, 4)) + + @pytest.mark.parametrize("func", [np.mean, np.median, np.max, np.std]) + def test_time_coord_reductions(self, random_patch, func): + """Reductions on time coords should use dascore's time-aware logic.""" + coord = random_patch.get_coord("time") + out = func(coord) + assert dtype_time_like(np.asarray(out).dtype) + assert not isinstance(out, BaseCoord) + + def test_array_function_not_implemented_for_unknown_types(self, coord): + """Unknown types in the array function protocol should defer.""" + out = coord.__array_function__(np.mean, (str,), (coord,), {}) + assert out is NotImplemented + + class TestIssues: """Tests for special issues related to coords.""" diff --git a/tests/test_io/test_dasdae/test_dasdae.py b/tests/test_io/test_dasdae/test_dasdae.py index 31cbf823b..c21a2fb02 100644 --- a/tests/test_io/test_dasdae/test_dasdae.py +++ b/tests/test_io/test_dasdae/test_dasdae.py @@ -260,7 +260,7 @@ def test_roundtrip_datetime_coord(self, tmp_path_factory, random_patch): """Ensure a patch with an attached datetime coord works.""" path = tmp_path_factory.mktemp("roundtrip_datetme_coord") / "out.h5" dist = random_patch.get_coord("distance") - dt = dc.to_datetime64(np.zeros_like(dist)) + dt = dc.to_datetime64(np.zeros_like(dist.values)) dt[0] = dc.to_datetime64("2017-09-17") new = random_patch.update_coords(dt=("distance", dt)) new.io.write(path, "dasdae") @@ -271,7 +271,7 @@ def test_roundtrip_nullish_datetime_coord(self, tmp_path_factory, random_patch): """Ensure a patch with an attached datetime coord with nulls works.""" path = tmp_path_factory.mktemp("roundtrip_datetime_coord") / "out.h5" dist = random_patch.get_coord("distance") - dt = dc.to_datetime64(np.zeros_like(dist)) + dt = dc.to_datetime64(np.zeros_like(dist.values)) dt[~dt.astype(bool)] = np.datetime64("nat") dt[0] = dc.to_datetime64("2017-09-17") dt[-4] = dc.to_datetime64("2020-01-03")