Fix 566: coordinate arithmetic

dascore/core/coords.py

@@ -430,10 +430,109 @@ def __str__(self):
 
     __repr__ = __str__
 
+    __array_priority__ = 1000.0
+
     def __array__(self, dtype=None, copy=False):
         """Numpy method for getting array data with `np.array(coord)`."""
         return self.data
 
+    def _get_coord_output(self, data, units=None):
+        """Return output from operations as a coordinate when possible."""
+        if isinstance(data, BaseCoord):
+            return data
+        if hasattr(data, "magnitude") and hasattr(data, "units"):
+            return get_coord(data=data.magnitude, units=data.units)
+        return get_coord(data=data, units=units)
+
+    def _binary_coord_op(self, operator, other, reversed=False):
+        """Apply a binary operator and return a new coordinate."""
+        other_data = other.data if isinstance(other, BaseCoord) else other
+        # Addition/subtraction treat scalars as values in current units.
+        if hasattr(other_data, "units") and operator in (np.add, np.subtract):
+            other_data = convert_units(
+                other_data.magnitude,
+                self.units,
+                other_data.units,
+            )
+        lhs, rhs = (other_data, self.data) if reversed else (self.data, other_data)
+        out = operator(lhs, rhs)
+        units = (
+            self.units if not np.issubdtype(np.asarray(out).dtype, np.bool_) else None
+        )
+        return self._get_coord_output(out, units=units)
+
+    def __add__(self, other):
+        return self._binary_coord_op(np.add, other)
+
+    def __sub__(self, other):
+        return self._binary_coord_op(np.subtract, other)
+
+    def __mul__(self, other):
+        return self._binary_coord_op(np.multiply, other)
+
+    def __truediv__(self, other):
+        return self._binary_coord_op(np.divide, other)
+
+    def __floordiv__(self, other):
+        return self._binary_coord_op(np.floor_divide, other)
+
+    def __pow__(self, other):
+        return self._binary_coord_op(np.power, other)
+
+    def __mod__(self, other):
+        return self._binary_coord_op(np.mod, other)
+
+    __radd__ = __add__
+
+    def __rsub__(self, other):
+        return self._binary_coord_op(np.subtract, other, reversed=True)
+
+    __rmul__ = __mul__
+
+    def __rtruediv__(self, other):
+        return self._binary_coord_op(np.divide, other, reversed=True)
+
+    def __rfloordiv__(self, other):
+        return self._binary_coord_op(np.floor_divide, other, reversed=True)
+
+    def __rpow__(self, other):
+        return self._binary_coord_op(np.power, other, reversed=True)
+
+    def __rmod__(self, other):
+        return self._binary_coord_op(np.mod, other, reversed=True)
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        """Support numpy ufunc operations and return coordinate outputs."""
+        method_func = ufunc if method == "__call__" else getattr(ufunc, method)
+        converted = [x.data if isinstance(x, BaseCoord) else x for x in inputs]
+        out = method_func(*converted, **kwargs)
+        units = (
+            self.units if not np.issubdtype(np.asarray(out).dtype, np.bool_) else None
+        )
+        return self._get_coord_output(out, units=units)
+
+    def __array_function__(self, func, types, args, kwargs):
+        """Support NumPy array-function protocol for coordinates."""
+        if not any(issubclass(t, BaseCoord) for t in types):
+            return NotImplemented
+
+        def _convert(obj):
+            if isinstance(obj, BaseCoord):
+                return obj.data
+            if isinstance(obj, tuple):
+                return tuple(_convert(x) for x in obj)
+            if isinstance(obj, list):
+                return [_convert(x) for x in obj]
+            if isinstance(obj, dict):
+                return {k: _convert(v) for k, v in obj.items()}
+            return obj
+
+        out = func(*_convert(args), **_convert(kwargs))
+        units = (
+            self.units if not np.issubdtype(np.asarray(out).dtype, np.bool_) else None
+        )
+        return self._get_coord_output(out, units=units)
+
     @cached_method
     def min(self):
         """Return min value."""

tests/test_core/test_coords.py

@@ -1895,6 +1895,89 @@ def test_not_implemented_in_baseclass(self, evenly_sampled_coord):
             BaseCoord.change_length(coord, 10)
 
 
+
+
+class TestCoordinateArithmetic:
+    """Tests for coordinate arithmetic behavior (issue #566)."""
+
+    @pytest.mark.parametrize(
+        "func, expected, other",
+        [
+            (lambda coord, val: coord + val, lambda vals, val: vals + val, 2),
+            (lambda coord, val: coord - val, lambda vals, val: vals - val, 2),
+            (lambda coord, val: coord * val, lambda vals, val: vals * val, 2),
+            (lambda coord, val: coord / val, lambda vals, val: vals / val, 2),
+            (lambda coord, val: coord // val, lambda vals, val: vals // val, 2),
+            (lambda coord, val: coord**val, lambda vals, val: vals**val, 2),
+            (lambda coord, val: coord % val, lambda vals, val: vals % val, 3),
+            (lambda coord, val: val + coord, lambda vals, val: val + vals, 2),
+            (lambda coord, val: val - coord, lambda vals, val: val - vals, 10),
+            (lambda coord, val: val * coord, lambda vals, val: val * vals, 2),
+            (lambda coord, val: val / coord, lambda vals, val: val / vals, 10),
+            (lambda coord, val: val // coord, lambda vals, val: val // vals, 10),
+            (lambda coord, val: val**coord, lambda vals, val: val**vals, 2),
+            (lambda coord, val: val % coord, lambda vals, val: val % vals, 10),
+        ],
+    )
+    def test_all_dunder_binary_ops(self, func, expected, other):
+        """Ensure all new arithmetic dunder paths return coordinates."""
+        coord = get_coord(data=[2, 4, 8], units="m")
+
+        out = func(coord, other)
+        assert isinstance(out, BaseCoord)
+        assert out.units == coord.units
+        np.testing.assert_allclose(out.values, expected(coord.values, other))
+
+    def test_numpy_ufunc_call_returns_coord(self):
+        """Ensure numpy ufunc __call__ dispatch returns coordinates."""
+        coord = get_coord(data=[1, 4, 9], units="m")
+        out = np.sqrt(coord)
+
+        assert isinstance(out, BaseCoord)
+        assert out.units == coord.units
+        np.testing.assert_allclose(out.values, np.array([1.0, 2.0, 3.0]))
+
+    def test_numpy_ufunc_accumulate_returns_coord(self):
+        """Ensure numpy ufunc method dispatch (e.g. accumulate) returns coordinates."""
+        coord = get_coord(data=[1, 2, 3], units="m")
+        out = np.add.accumulate(coord)
+
+        assert isinstance(out, BaseCoord)
+        assert out.units == coord.units
+        np.testing.assert_array_equal(out.values, np.array([1, 3, 6]))
+
+    def test_numpy_array_function_returns_coord(self):
+        """Ensure numpy array functions return coordinates where possible."""
+        coord = get_coord(data=[3, 4], units="m")
+        out = np.linalg.norm(coord)
+
+        assert isinstance(out, BaseCoord)
+        assert out.units == coord.units
+        np.testing.assert_allclose(out.values, np.array([5.0]))
+
+    def test_tuple_list_dict_conversions(self):
+        """Ensure tuple/list/dict conversion paths are exercised."""
+        coord1 = get_coord(data=[1, 2], units="m")
+        coord2 = get_coord(data=[3, 4], units="m")
+
+        out_tuple = np.concatenate((coord1, coord2))
+        assert isinstance(out_tuple, BaseCoord)
+        np.testing.assert_array_equal(out_tuple.values, np.array([1, 2, 3, 4]))
+
+        out_list = np.concatenate([coord1, coord2])
+        assert isinstance(out_list, BaseCoord)
+        np.testing.assert_array_equal(out_list.values, np.array([1, 2, 3, 4]))
+
+        out_kwargs = np.mean(a=coord1)
+        assert isinstance(out_kwargs, BaseCoord)
+        np.testing.assert_allclose(out_kwargs.values, np.array([1.5]))
+
+    def test_array_function_returns_not_implemented_for_non_coord_types(self):
+        """Ensure non-coordinate array_function dispatch returns NotImplemented."""
+        coord = get_coord(data=[1, 2, 3], units="m")
+        out = coord.__array_function__(np.mean, (np.ndarray,), (coord,), {})
+        assert out is NotImplemented
+
 class TestIssues:
     """Tests for special issues related to coords."""