NF: add support for Python operators to ArraySequence

Author: MarcCote

nibabel/streamlines/array_sequence.py

@@ -53,6 +53,35 @@ def update_seq(self, arr_seq):
         arr_seq._lengths = np.array(self.lengths)
 
 
+def _define_operators(cls):
+    """ Decorator which adds support for some Python operators. """
+    def _wrap(cls, op, name=None, inplace=False, unary=False):
+        name = name or op
+        if unary:
+            setattr(cls, name, lambda self: self._op(op))
+        else:
+            setattr(cls, name,
+                    lambda self, value: self._op(op, value, inplace=inplace))
+
+    for op in ["__iadd__", "__isub__", "__imul__", "__idiv__",
+               "__ifloordiv__", "__itruediv__", "__ior__"]:
+        _wrap(cls, op, inplace=True)
+
+    for op in ["__add__", "__sub__", "__mul__", "__div__",
+               "__floordiv__", "__truediv__", "__or__"]:
+        op_ = "__i{}__".format(op.strip("_"))
+        _wrap(cls, op_, name=op)
+
+    for op in ["__eq__", "__ne__", "__lt__", "__le__", "__gt__", "__ge__"]:
+        _wrap(cls, op)
+
+    for op in ["__neg__"]:
+        _wrap(cls, op, unary=True)
+
+    return cls
+
+
+@_define_operators
 class ArraySequence(object):
     """ Sequence of ndarrays having variable first dimension sizes.
 
@@ -120,6 +149,23 @@ def data(self):
         """ Elements in this array sequence. """
         return self._data
 
+    def _check_shape(self, arrseq):
+        """ Check whether this array sequence is compatible with another. """
+        msg = "cannot perform operation - array sequences have different"
+        if len(self._lengths) != len(arrseq._lengths):
+            msg += " lengths: {} vs. {}."
+            raise ValueError(msg.format(len(self._lengths), len(arrseq._lengths)))
+
+        if self.total_nb_rows != arrseq.total_nb_rows:
+            msg += " amount of data: {} vs. {}."
+            raise ValueError(msg.format(self.total_nb_rows, arrseq.total_nb_rows))
+
+        if self.common_shape != arrseq.common_shape:
+            msg += " common shape: {} vs. {}."
+            raise ValueError(msg.format(self.common_shape, arrseq.common_shape))
+
+        return True
+
     def _get_next_offset(self):
         """ Offset in ``self._data`` at which to write next rowelement """
         if len(self._offsets) == 0:
@@ -377,6 +423,37 @@ def __setitem__(self, idx, elements):
         for o1, l1, o2, l2 in zip(offsets, lengths, elements._offsets, elements._lengths):
             data[o1:o1 + l1] = elements._data[o2:o2 + l2]
 
+    def _op(self, op, value=None, inplace=False):
+        """ Applies some operator to this arraysequence.
+
+        This handles both unary and binary operators with a scalar or another
+        array sequence. Operations are performed directly on the underlying
+        data, or a copy of it, which depends on the value of `inplace`.
+
+        Parameters
+        ----------
+        op : str
+            Name of the Python operator (e.g., `"__add__"`).
+        value : scalar or :class:`ArraySequence`, optional
+            If None, the operator is assumed to be unary.
+            Otherwise, that value is used in the binary operation.
+        inplace: bool, optional
+            If False, the operation is done on a copy of this array sequence.
+            Otherwise, this array sequence gets modified directly.
+        """
+        seq = self if inplace else self.copy()
+
+        if is_array_sequence(value) and seq._check_shape(value):
+            for o1, l1, o2, l2 in zip(seq._offsets, seq._lengths, value._offsets, value._lengths):
+                seq._data[o1:o1 + l1] = getattr(seq._data[o1:o1 + l1], op)(value._data[o2:o2 + l2])
+
+        else:
+            args = [] if value is None else [value]  # Dealing with unary and binary ops.
+            for o1, l1 in zip(seq._offsets, seq._lengths):
+                seq._data[o1:o1 + l1] = getattr(seq._data[o1:o1 + l1], op)(*args)
+
+        return seq
+
     def __iter__(self):
         if len(self._lengths) != len(self._offsets):
             raise ValueError("ArraySequence object corrupted:"

nibabel/streamlines/tests/test_array_sequence.py

@@ -277,6 +277,67 @@ def test_arraysequence_getitem(self):
         check_arr_seq_view(seq_view, SEQ_DATA['seq'])
         check_arr_seq(seq_view, [d[:, 2] for d in SEQ_DATA['data'][::-2]])
 
+    def test_arraysequence_operators(self):
+        for op in ["__add__", "__sub__", "__mul__", "__floordiv__", "__truediv__",
+                   "__eq__", "__ne__", "__lt__", "__le__", "__gt__", "__ge__"]:
+            # Test math operators with a scalar.
+            for scalar in [42, 0.5, True]:
+                seq = getattr(SEQ_DATA['seq'], op)(scalar)
+                assert_true(seq is not SEQ_DATA['seq'])
+                check_arr_seq(seq, [getattr(d, op)(scalar) for d in SEQ_DATA['data']])
+
+            # Test math operators with another ArraySequence.
+            seq = getattr(SEQ_DATA['seq'], op)(SEQ_DATA['seq'])
+            assert_true(seq is not SEQ_DATA['seq'])
+            check_arr_seq(seq, [getattr(d, op)(d) for d in SEQ_DATA['data']])
+
+            # Test math operators with ArraySequence views.
+            orig = SEQ_DATA['seq'][::2]
+            seq = getattr(orig, op)(orig)
+            assert_true(seq is not orig)
+            check_arr_seq(seq, [getattr(d, op)(d) for d in SEQ_DATA['data'][::2]])
+
+        # Test in-place operators.
+        for op in ["__iadd__", "__isub__", "__imul__", "__ifloordiv__", "__itruediv__"]:
+            # Test in-place math operators with a scalar.
+            for scalar in [42, 0.5, True]:
+                seq = seq_orig = SEQ_DATA['seq'].copy()
+                seq = getattr(seq, op)(scalar)
+                assert_true(seq is seq_orig)
+                check_arr_seq(seq, [getattr(d.copy(), op)(scalar) for d in SEQ_DATA['data']])
+
+            # Test in-place math operators with another ArraySequence.
+            seq = seq_orig = SEQ_DATA['seq'].copy()
+            seq = getattr(seq, op)(SEQ_DATA['seq'])
+            assert_true(seq is seq_orig)
+            check_arr_seq(seq, [getattr(d.copy(), op)(d) for d in SEQ_DATA['data']])
+
+            # Test in-place math operators with ArraySequence views.
+            seq = seq_orig = SEQ_DATA['seq'].copy()[::2]
+            seq = getattr(seq, op)(seq)
+            assert_true(seq is seq_orig)
+            check_arr_seq(seq, [getattr(d.copy(), op)(d) for d in SEQ_DATA['data'][::2]])
+
+            # Operations between array sequences of different lengths.
+            seq = SEQ_DATA['seq'].copy()
+            assert_raises(ValueError, getattr(seq, op), SEQ_DATA['seq'][::2])
+
+            # Operations between array sequences with different amount of data.
+            seq1 = ArraySequence(np.arange(10).reshape(5, 2))
+            seq2 = ArraySequence(np.arange(15).reshape(5, 3))
+            assert_raises(ValueError, getattr(seq1, op), seq2)
+
+            # Operations between array sequences with different common shape.
+            seq1 = ArraySequence(np.arange(12).reshape(2, 2, 3))
+            seq2 = ArraySequence(np.arange(8).reshape(2, 2, 2))
+            assert_raises(ValueError, getattr(seq1, op), seq2)
+
+        # Unary operators
+        for op in ["__neg__"]:
+            seq = getattr(SEQ_DATA['seq'], op)()
+            assert_true(seq is not SEQ_DATA['seq'])
+            check_arr_seq(seq, [getattr(d, op)() for d in SEQ_DATA['data']])
+
     def test_arraysequence_setitem(self):
         # Set one item
         seq = SEQ_DATA['seq'] * 0