From cc26ade36cf67796a190d84e52f78cf0d80d8a11 Mon Sep 17 00:00:00 2001 From: Tias Guns Date: Sat, 18 Jul 2026 15:09:04 +0100 Subject: [PATCH 1/3] ufuncs for ndvararray, simpler than thought --- cpmpy/expressions/variables.py | 291 ++++++++++++--------------------- tests/test_expressions.py | 42 +++++ 2 files changed, 146 insertions(+), 187 deletions(-) diff --git a/cpmpy/expressions/variables.py b/cpmpy/expressions/variables.py index 88f348859..728c5bd2d 100644 --- a/cpmpy/expressions/variables.py +++ b/cpmpy/expressions/variables.py @@ -57,7 +57,8 @@ """ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Callable, Iterable +import operator import threading import warnings # for deprecation warning from functools import reduce @@ -68,6 +69,32 @@ from .core import Expression, ExprLike, BoolExprLike, ListLike, BoolVal from .utils import is_num, is_int, is_boolexpr, get_bounds +# NumPy ufunc -> Python operator (single source of truth for NDVarArray.__array_ufunc__) +_ND_BINOPS: dict[np.ufunc, Callable[[Any, Any], Any]] = { + np.add: operator.add, + np.subtract: operator.sub, + np.multiply: operator.mul, + np.floor_divide: operator.floordiv, + np.true_divide: operator.truediv, + np.mod: operator.mod, + np.remainder: operator.mod, + np.power: operator.pow, + np.equal: operator.eq, + np.not_equal: operator.ne, + np.less: operator.lt, + np.less_equal: operator.le, + np.greater: operator.gt, + np.greater_equal: operator.ge, + np.bitwise_and: operator.and_, + np.bitwise_or: operator.or_, + np.bitwise_xor: operator.xor, +} +_ND_UNARYOPS: dict[np.ufunc, Callable[[Any], Any]] = { + np.negative: operator.neg, + np.absolute: operator.abs, + np.invert: operator.invert, +} + _BV_PREFIX = "BV" _IV_PREFIX = "IV" _VAR_ERR = f"Variable names starting with {_IV_PREFIX} or {_BV_PREFIX} are reserved for internal use only, chose a different name" @@ -270,7 +297,7 @@ def cpm_array(arr: ListLike[ExprLike|Any]) -> NDVarArray: elif not arr.flags['FORC']: # Ensure the array is contiguous arr = np.ascontiguousarray(arr) - order = 'F' if arr.flags['F_CONTIGUOUS'] else 'C' + order: Literal['C', 'F'] = 'F' if arr.flags['F_CONTIGUOUS'] else 'C' return NDVarArray(shape=arr.shape, dtype=arr.dtype, buffer=arr, order=order) @@ -445,20 +472,16 @@ class NDVarArray(np.ndarray): Do not create this object directly, use one of the functions in this module ``_has_subexpr`` caches :meth:`has_subexpr` (``None`` = not computed yet; ``True`` / ``False`` = cached). + + Operator overloading follows NumPy's ndarray-subclass rule: all logic lives in + :meth:`__array_ufunc__` (not in ``__add__`` / ``__eq__`` / ...). Python operators + and ``np.add`` / ``np.equal`` / ... therefore share one path. """ _has_subexpr: Optional[bool] = None # will be overwritten in instance, here for type hinting - def __init__(self, shape: int|np.integer|tuple[int|np.integer, ...], **kwargs: Any) -> None: - # bit ugly, but np.int and np.bool do not play well with > overloading - if np.issubdtype(self.dtype, np.integer): - self.astype(int) - elif np.issubdtype(self.dtype, np.bool_): - self.astype(bool) - + def __array_finalize__(self, obj: Optional[np.ndarray]) -> None: + # views / copies / ufunc outputs: __init__ is not called self._has_subexpr = None - # no need to call ndarray __init__ method as specified in the np.ndarray documentation: - # "No ``__init__`` method is needed because the array is fully initialized - # after the ``__new__`` method." def has_subexpr(self) -> bool: """True if :meth:`flat` has an :class:`Expression` that is not a variable (:class:`_NumVarImpl`) or :class:`~cpmpy.expressions.core.BoolVal`.""" @@ -516,38 +539,36 @@ def __getitem__(self, index): # TODO: any typing would have to be compatible wi return super().__getitem__(index) - """ - make the given array the first dimension in the returned array - """ - def __axis(self, axis): - - arr = self - - # correct type and value checks - if not isinstance(axis,int): - raise TypeError("Axis keyword argument in .sum() should always be an integer") - if axis >= arr.ndim: - raise ValueError("Axis out of range") - - if axis < 0: - axis += arr.ndim - - # Change the array to make the selected axis the first dimension - if axis > 0: - iter_axis = list(range(arr.ndim)) - iter_axis.remove(axis) - iter_axis.insert(0, axis) - arr = arr.transpose(iter_axis) - - return arr + def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any) -> Any: + """ + NumPy ufunc entry point for symbolic element-wise ops. - def sum(self, axis=None, out=None): + Supported ``__call__`` ufuncs map to :mod:`operator` and build Expression trees + via :meth:`_elementwise`. Unsupported ufuncs raise :class:`TypeError` (never + fall through to object-dtype ufuncs that boolify comparisons). + """ + if kwargs.get("out") is not None: + raise TypeError("NDVarArray does not support the 'out=' argument for ufuncs") + if method != "__call__": + # .prod(axis=...) used to call multiply.reduce; we no longer rely on that + raise TypeError( + f"NDVarArray does not support np.{ufunc.__name__}.{method}; " + "use Python operators or cp.*" + ) + if ufunc in _ND_UNARYOPS: + return self._elementwise(_ND_UNARYOPS[ufunc], inputs[0]) + if ufunc in _ND_BINOPS: + return self._elementwise(_ND_BINOPS[ufunc], *inputs) + raise TypeError( + f"NDVarArray does not support np.{ufunc.__name__}; use Python operators or cp.*" + ) + + def sum(self, axis=None, dtype=None, out=None, keepdims=False, **kwargs): """ overwrite np.sum(NDVarArray) as people might use it """ - - if out is not None: - raise NotImplementedError() + if out is not None or dtype is not None or keepdims: + raise NotImplementedError("out/dtype/keepdims not supported for NDVarArray operators") if axis is None: # simple case where we want the sum over the whole array return cp.sum(self) @@ -555,53 +576,53 @@ def sum(self, axis=None, out=None): return cpm_array(np.apply_along_axis(cp.sum, axis=axis, arr=self)) - def prod(self, axis=None, out=None): + def prod(self, axis=None, dtype=None, out=None, keepdims=False, **kwargs): """ overwrite np.prod(NDVarArray) as people might use it """ - - if out is not None: - raise NotImplementedError() + if out is not None or dtype is not None or keepdims: + raise NotImplementedError("out/dtype/keepdims not supported for NDVarArray operators") if axis is None: # simple case where we want the product over the whole array return reduce(lambda a, b: a * b, self.flatten()) - # TODO: is there a better way? This does pairwise multiplication still - return cpm_array(np.multiply.reduce(self, axis=axis)) + return cpm_array(np.apply_along_axis( + lambda a: reduce(lambda x, y: x * y, a), axis=axis, arr=self + )) - def max(self, axis=None, out=None): + def max(self, axis=None, dtype=None, out=None, keepdims=False, **kwargs): """ overwrite np.max(NDVarArray) as people might use it """ - if out is not None: - raise NotImplementedError() + if out is not None or dtype is not None or keepdims: + raise NotImplementedError("out/dtype/keepdims not supported for NDVarArray operators") if axis is None: # simple case where we want the maximum over the whole array return cp.max(self) return cpm_array(np.apply_along_axis(cp.max, axis=axis, arr=self)) - def min(self, axis=None, out=None): + def min(self, axis=None, dtype=None, out=None, keepdims=False, **kwargs): """ overwrite np.min(NDVarArray) as people might use it """ - if out is not None: - raise NotImplementedError() + if out is not None or dtype is not None or keepdims: + raise NotImplementedError("out/dtype/keepdims not supported for NDVarArray operators") if axis is None: # simple case where we want the minimum over the whole array return cp.min(self) return cpm_array(np.apply_along_axis(cp.min, axis=axis, arr=self)) - def any(self, axis=None, out=None): + def any(self, axis=None, dtype=None, out=None, keepdims=False, **kwargs): """ overwrite np.any(NDVarArray) """ if any(not is_boolexpr(x) for x in self.flat): raise TypeError("Cannot call .any() in an array not consisting only of bools") - if out is not None: - raise NotImplementedError() + if out is not None or dtype is not None or keepdims: + raise NotImplementedError("out/dtype/keepdims not supported for NDVarArray operators") if axis is None: # simple case where we want a disjunction over the whole array return cp.any(self) @@ -609,15 +630,15 @@ def any(self, axis=None, out=None): return cpm_array(np.apply_along_axis(cp.any, axis=axis, arr=self)) - def all(self, axis=None, out=None): + def all(self, axis=None, dtype=None, out=None, keepdims=False, **kwargs): """ - overwrite np.any(NDVarArray) + overwrite np.all(NDVarArray) """ if any(not is_boolexpr(x) for x in self.flat): - raise TypeError("Cannot call .any() in an array not consisting only of bools") + raise TypeError("Cannot call .all() in an array not consisting only of bools") - if out is not None: - raise NotImplementedError() + if out is not None or dtype is not None or keepdims: + raise NotImplementedError("out/dtype/keepdims not supported for NDVarArray operators") if axis is None: # simple case where we want a conjunction over the whole array return cp.all(self) @@ -633,137 +654,33 @@ def get_bounds(self) -> tuple[np.ndarray, np.ndarray]: return np.asarray(lbs).reshape(self.shape), \ np.asarray(ubs).reshape(self.shape) - # VECTORIZED master function (delegate) - def _vectorized(self, other: ExprLike|Iterable|Any, attr: str, **kwargs) -> NDVarArray: - """ - NumPy-broadcast ``other`` to ``self.shape``, then apply ``attr`` element-wise. - - Args: - other (ExprLike|Iterable|Any): The other operand. - Typically an array/list of Expressions, or a single Expression, or a constant (or anything np compatible) - attr (str): The attribute to vectorize (e.g. ``__eq__``, ``__add__``, ``implies``, etc.) - **kwargs: Extra keyword arguments passed to each element call (e.g. ``simplify`` for ``implies``). - - Returns: - NDVarArray: The vectorized result. - """ - if not isinstance(other, np.ndarray): - other = np.asarray(other) + @staticmethod + def _elementwise(op: Callable[..., Any], *inputs: Any) -> NDVarArray: + """Broadcast ``inputs`` and apply ``op`` element-wise; return an :class:`NDVarArray`.""" + arrays = [ + np.asarray(x, dtype=object) if not np.isscalar(x) else np.array(x, dtype=object) + for x in inputs + ] try: - broadcast_shape = np.broadcast_shapes(other.shape, self.shape) + arrays = list(np.broadcast_arrays(*arrays)) except ValueError as e: + shapes = " ".join(str(a.shape) for a in arrays) raise ValueError( - f"operands could not be broadcast together with shapes {self.shape} {other.shape}" + f"operands could not be broadcast together with shapes {shapes}" ) from e - if broadcast_shape != self.shape: - raise ValueError( - f"other with shape {other.shape} could not be broadcast to self with shape {self.shape}" - ) - other = np.broadcast_to(other, self.shape) - # s.__eq__(o) <-> getattr(s, '__eq__')(o) - flat_res = cpm_array([ - # unwrap numpy 'generic' scalar to is Python int item, so int <= np.int64 does not return NotImplemented - getattr(s, attr)(o.item() if isinstance(o, np.generic) else o, **kwargs) - for s, o in zip(self.flat, other.flat)]) - # typing is wrong, reshape does return NDVarArray - return flat_res.reshape(self.shape) # type: ignore - - # VECTORIZED comparisons - def __eq__(self, other): - return self._vectorized(other, '__eq__') - - def __ne__(self, other): - return self._vectorized(other, '__ne__') - - def __lt__(self, other): - return self._vectorized(other, '__lt__') - - def __le__(self, other): - return self._vectorized(other, '__le__') - - def __gt__(self, other): - return self._vectorized(other, '__gt__') - - def __ge__(self, other): - return self._vectorized(other, '__ge__') - - # VECTORIZED math operators - # only 'abs' 'neg' and binary ones - # '~' not needed, gets translated to ==0 and that is already handled - def __abs__(self): - return cpm_array([abs(s) for s in self]) - - def __neg__(self): - return cpm_array([-s for s in self]) - - def __add__(self, other): - return self._vectorized(other, '__add__') - - def __radd__(self, other): - return self._vectorized(other, '__radd__') - - def __sub__(self, other): - return self._vectorized(other, '__sub__') - - def __rsub__(self, other): - return self._vectorized(other, '__rsub__') - - def __mul__(self, other): - return self._vectorized(other, '__mul__') - - def __rmul__(self, other): - return self._vectorized(other, '__rmul__') - - def __truediv__(self, other): - return self._vectorized(other, '__truediv__') - - def __rtruediv__(self, other): - return self._vectorized(other, '__rtruediv__') - - def __floordiv__(self, other): - return self._vectorized(other, '__floordiv__') - - def __rfloordiv__(self, other): - return self._vectorized(other, '__rfloordiv__') - - def __mod__(self, other): - return self._vectorized(other, '__mod__') - - def __rmod__(self, other): - return self._vectorized(other, '__rmod__') - - def __pow__(self, other, modulo=None): - assert (modulo is None), "Power operator: modulo not supported" - return self._vectorized(other, '__pow__') - - def __rpow__(self, other, modulo=None): - assert (modulo is None), "Power operator: modulo not supported" - return self._vectorized(other, '__rpow__') - - # VECTORIZED Bool operators - def __invert__(self): - return cpm_array([~s for s in self]) - - def __and__(self, other): - return self._vectorized(other, '__and__') - - def __rand__(self, other): - return self._vectorized(other, '__rand__') - - def __or__(self, other): - return self._vectorized(other, '__or__') - - def __ror__(self, other): - return self._vectorized(other, '__ror__') - - def __xor__(self, other): - return self._vectorized(other, '__xor__') - - def __rxor__(self, other): - return self._vectorized(other, '__rxor__') - - def implies(self, other: BoolExprLike|Iterable[BoolExprLike], simplify=False) -> NDVarArray: - return self._vectorized(other, 'implies', simplify=simplify) + out = np.empty(arrays[0].shape, dtype=object) + flat_in = [a.flat for a in arrays] + flat_out: list[Any] = [] + for elems in zip(*flat_in): + args = [e.item() if isinstance(e, np.generic) else e for e in elems] + flat_out.append(op(*args)) + out.flat[:] = flat_out + return out.view(NDVarArray) + + def implies(self, other: BoolExprLike|Iterable[BoolExprLike], simplify: bool = False) -> NDVarArray: + def _impl(a: Any, b: Any) -> Any: + return a.implies(b, simplify=simplify) + return self._elementwise(_impl, self, other) #in __contains__(self, value) Check membership # CANNOT meaningfully overwrite, python always returns True/False diff --git a/tests/test_expressions.py b/tests/test_expressions.py index f753b1357..7285d8462 100644 --- a/tests/test_expressions.py +++ b/tests/test_expressions.py @@ -280,11 +280,53 @@ def test_numpy_expr_mul(self): for idx in np.ndindex(expr.shape): assert str(expr[idx]) == str(ref[idx]) + def test_full_broadcast_mul(self): + # ufunc-style: (2,1) * (2,3) -> (2,3) + x = intvar(0, 10, shape=(2, 1), name="x") + w = np.array([[1, 2, 3], [4, 5, 6]]) + expr = x * w + assert expr.shape == (2, 3) + assert isinstance(expr, NDVarArray) + def test_incompatible_broadcast_raises(self): x = intvar(0, 10, shape=(3, 4), name="x") with pytest.raises(ValueError, match="broadcast"): x * np.array([1, 2]) + def test_np_equal_returns_comparisons(self): + x = intvar(0, 5, shape=3, name="x") + w = np.array([1, 2, 3]) + eq = np.equal(x, w) + assert isinstance(eq, NDVarArray) + assert all(isinstance(e, Comparison) for e in eq) + assert str(eq[0]) == "x[0] == 1" + + def test_reflected_add(self): + x = intvar(0, 5, shape=3, name="x") + r = 1 + x + assert isinstance(r, NDVarArray) + assert str(r[0]) == "1 + (x[0])" + + def test_unsupported_ufunc_raises(self): + x = intvar(0, 5, shape=3, name="x") + with pytest.raises(TypeError, match="does not support"): + np.sin(x) + + def test_views_have_has_subexpr(self): + x = intvar(0, 5, shape=4, name="x") + y = x[1:] + assert isinstance(y, NDVarArray) + assert y.has_subexpr() is False + z = (x + 1)[1:] + assert z.has_subexpr() is True + + def test_np_sum_keepdims_accepted(self): + x = intvar(0, 5, shape=3, name="x") + # must not TypeError on unexpected kwargs; keepdims itself is unsupported + with pytest.raises(NotImplementedError, match="out/dtype/keepdims"): + np.sum(x, keepdims=True) + + class TestArrayExpressions: def test_scalar_expr_with_ndarray(self): From af5816bf04ad7ff18acbb1a1bcc176082b8af565 Mon Sep 17 00:00:00 2001 From: Tias Guns Date: Sat, 18 Jul 2026 15:17:33 +0100 Subject: [PATCH 2/3] boolvar/intvar with np.view --- cpmpy/expressions/variables.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/cpmpy/expressions/variables.py b/cpmpy/expressions/variables.py index 728c5bd2d..cf42acd95 100644 --- a/cpmpy/expressions/variables.py +++ b/cpmpy/expressions/variables.py @@ -172,10 +172,9 @@ def boolvar(shape: int|np.integer|tuple[int|np.integer, ...] = 1, names = _gen_var_names(name, shape) # create np.array 'data' representation of the decision variables - data = np.array([_BoolVarImpl(name=n) for n in names]) - # insert into custom ndarray - r = NDVarArray(shape, dtype=object, buffer=data) - r._has_subexpr = False # A bit ugly (acces to private field) but otherwise np.ndarray constructor complains if we pass it as an argument to NDVarArray + data = np.array([_BoolVarImpl(name=n) for n in names]).reshape(shape) + r = data.view(NDVarArray) + r._has_subexpr = False # A bit ugly (acces to private field) but otherwise np.ndarray constructor complains if we pass it as an argument to NDVarArray return r @@ -258,10 +257,9 @@ def intvar(lb: int, ub: int, shape: int|np.integer|tuple[int|np.integer, ...] = names = _gen_var_names(name, shape) # create np.array 'data' representation of the decision variables - data = np.array([_IntVarImpl(lb, ub, name=n) for n in names]) # repeat new instances - # insert into custom ndarray - r = NDVarArray(shape, dtype=object, buffer=data) - r._has_subexpr = False # A bit ugly (acces to private field) but otherwise np.ndarray constructor complains if we pass it as an argument to NDVarArray + data = np.array([_IntVarImpl(lb, ub, name=n) for n in names]).reshape(shape) # repeat new instances + r = data.view(NDVarArray) + r._has_subexpr = False # A bit ugly (acces to private field) but otherwise np.ndarray constructor complains if we pass it as an argument to NDVarArray return r @@ -296,9 +294,8 @@ def cpm_array(arr: ListLike[ExprLike|Any]) -> NDVarArray: arr = np.array(arr) elif not arr.flags['FORC']: # Ensure the array is contiguous arr = np.ascontiguousarray(arr) - - order: Literal['C', 'F'] = 'F' if arr.flags['F_CONTIGUOUS'] else 'C' - return NDVarArray(shape=arr.shape, dtype=arr.dtype, buffer=arr, order=order) + + return arr.view(NDVarArray) class NullShapeError(Exception): From 91766d2c56d254e06d30355eae6c44f020977d4a Mon Sep 17 00:00:00 2001 From: Tias Guns Date: Sat, 18 Jul 2026 16:35:37 +0100 Subject: [PATCH 3/3] cleanup and some added tests --- cpmpy/expressions/variables.py | 40 ++++++++++++++++++---------------- tests/test_expressions.py | 19 ++++++++++++++++ 2 files changed, 40 insertions(+), 19 deletions(-) diff --git a/cpmpy/expressions/variables.py b/cpmpy/expressions/variables.py index 038ad1fd2..99b62f56f 100644 --- a/cpmpy/expressions/variables.py +++ b/cpmpy/expressions/variables.py @@ -70,7 +70,7 @@ from .utils import is_num, is_int, is_boolexpr, get_bounds # NumPy ufunc -> Python operator (single source of truth for NDVarArray.__array_ufunc__) -_ND_BINOPS: dict[np.ufunc, Callable[[Any, Any], Any]] = { +_ND_UFUNCS: dict[np.ufunc, Callable[..., Any]] = { np.add: operator.add, np.subtract: operator.sub, np.multiply: operator.mul, @@ -88,8 +88,10 @@ np.bitwise_and: operator.and_, np.bitwise_or: operator.or_, np.bitwise_xor: operator.xor, -} -_ND_UNARYOPS: dict[np.ufunc, Callable[[Any], Any]] = { + np.logical_and: operator.and_, + np.logical_or: operator.or_, + np.logical_xor: operator.xor, + np.logical_not: operator.invert, np.negative: operator.neg, np.absolute: operator.abs, np.invert: operator.invert, @@ -174,7 +176,7 @@ def boolvar(shape: int|np.integer|tuple[int|np.integer, ...] = 1, # create np.array 'data' representation of the decision variables data = np.array([_BoolVarImpl(name=n) for n in names]).reshape(shape) r = data.view(NDVarArray) - r._has_subexpr = False # A bit ugly (acces to private field) but otherwise np.ndarray constructor complains if we pass it as an argument to NDVarArray + r._has_subexpr = False # set after view; __array_finalize__ left it None return r @@ -259,7 +261,7 @@ def intvar(lb: int, ub: int, shape: int|np.integer|tuple[int|np.integer, ...] = # create np.array 'data' representation of the decision variables data = np.array([_IntVarImpl(lb, ub, name=n) for n in names]).reshape(shape) # repeat new instances r = data.view(NDVarArray) - r._has_subexpr = False # A bit ugly (acces to private field) but otherwise np.ndarray constructor complains if we pass it as an argument to NDVarArray + r._has_subexpr = False # set after view; __array_finalize__ left it None return r @@ -554,21 +556,23 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: via :meth:`_elementwise`. Unsupported ufuncs raise :class:`TypeError` (never fall through to object-dtype ufuncs that boolify comparisons). """ - if kwargs.get("out") is not None: - raise TypeError("NDVarArray does not support the 'out=' argument for ufuncs") + if kwargs: + raise TypeError( + f"NDVarArray does not support ufunc keyword arguments {sorted(kwargs)}; " + "use Python operators or cp.*" + ) if method != "__call__": # .prod(axis=...) used to call multiply.reduce; we no longer rely on that raise TypeError( f"NDVarArray does not support np.{ufunc.__name__}.{method}; " "use Python operators or cp.*" ) - if ufunc in _ND_UNARYOPS: - return self._elementwise(_ND_UNARYOPS[ufunc], inputs[0]) - if ufunc in _ND_BINOPS: - return self._elementwise(_ND_BINOPS[ufunc], *inputs) - raise TypeError( - f"NDVarArray does not support np.{ufunc.__name__}; use Python operators or cp.*" - ) + op = _ND_UFUNCS.get(ufunc) + if op is None: + raise TypeError( + f"NDVarArray does not support np.{ufunc.__name__}; use Python operators or cp.*" + ) + return self._elementwise(op, *inputs) def sum(self, axis=None, dtype=None, out=None, keepdims=False, **kwargs): """ @@ -676,12 +680,10 @@ def _elementwise(op: Callable[..., Any], *inputs: Any) -> NDVarArray: f"operands could not be broadcast together with shapes {shapes}" ) from e out = np.empty(arrays[0].shape, dtype=object) - flat_in = [a.flat for a in arrays] - flat_out: list[Any] = [] - for elems in zip(*flat_in): + for i, elems in enumerate(zip(*(a.flat for a in arrays))): + # unwrap numpy scalars so Expression ops get Python ints/bools args = [e.item() if isinstance(e, np.generic) else e for e in elems] - flat_out.append(op(*args)) - out.flat[:] = flat_out + out.flat[i] = op(*args) return out.view(NDVarArray) def implies(self, other: BoolExprLike|Iterable[BoolExprLike], simplify: bool = False) -> NDVarArray: diff --git a/tests/test_expressions.py b/tests/test_expressions.py index 32be8c031..1b1526b35 100644 --- a/tests/test_expressions.py +++ b/tests/test_expressions.py @@ -301,16 +301,35 @@ def test_np_equal_returns_comparisons(self): assert all(isinstance(e, Comparison) for e in eq) assert str(eq[0]) == "x[0] == 1" + def test_np_logical_and(self): + b = boolvar(shape=3, name="b") + expr = np.logical_and(b, ~b) + assert isinstance(expr, NDVarArray) + assert str(expr[0]) == "(b[0]) and (~b[0])" + def test_reflected_add(self): x = intvar(0, 5, shape=3, name="x") r = 1 + x assert isinstance(r, NDVarArray) assert str(r[0]) == "1 + (x[0])" + def test_expr_left_of_array(self): + # Expression.__lt__/__add__ reverse onto NDVarArray (ndarray ufuncs) + x = intvar(0, 5, name="x") + arr = intvar(0, 5, shape=3, name="arr") + lt = x < arr + assert isinstance(lt, NDVarArray) + assert str(lt[0]) == "(arr[0]) > (x)" + s = x + arr + assert isinstance(s, NDVarArray) + assert str(s[0]) == "(x) + (arr[0])" + def test_unsupported_ufunc_raises(self): x = intvar(0, 5, shape=3, name="x") with pytest.raises(TypeError, match="does not support"): np.sin(x) + with pytest.raises(TypeError, match="keyword arguments"): + np.add(x, 1, out=np.empty(3, dtype=object)) def test_views_have_has_subexpr(self): x = intvar(0, 5, shape=4, name="x")