Merge pull request #408 from bsc-wdc/in_place_magic_methods_array

Added in-place operations

Author: FernandoVN98

dislib/data/array.py

@@ -224,6 +224,123 @@ def __sub__(self, other):
         return Array(blocks, self._top_left_shape, self._reg_shape,
                      self.shape, self._sparse)
 
+    def __isub__(self, other):
+        if self.shape[1] != other.shape[1] or (other.shape[0] != 1 and
+                                               (other.shape[0] != self.shape[0]
+                                               or other._reg_shape[0] !=
+                                               self._reg_shape[0])):
+            raise NotImplementedError("Reverse subtraction not "
+                                      "implemented for the "
+                                      "given objects")
+
+        if other.shape[0] == self.shape[0] and other._reg_shape[0]\
+                == self._reg_shape[0]:
+            if other.shape[0] == self.shape[0] and \
+                    other._reg_shape[0] == self._reg_shape[0]:
+                if self._reg_shape[1] != other._reg_shape[1]:
+                    raise ValueError("incorrect block sizes for the requested "
+                                     f"subtraction ("
+                                     f"{self._reg_shape[0]}"", "
+                                     f"{self._reg_shape[1]} !="
+                                     f"{other._reg_shape[0]}"", "
+                                     f"{other._reg_shape[1]})")
+
+                if self._top_left_shape != other._top_left_shape:
+                    raise ValueError("Incompatible block sizes of the "
+                                     "top left block of the matrices"
+                                     "b._top_left_shape != b._top_left_shape")
+            # matrix - matrix
+            blocks = []
+
+            for hblock, others in zip(self._iterator("rows"),
+                                      other._iterator("rows")):
+                out_blocks = [object() for _ in range(hblock._n_blocks[1])]
+                _combine_blocks(hblock._blocks, others._blocks,
+                                Array._subtract, out_blocks)
+                blocks.append(out_blocks)
+            compss_delete_object(self._blocks)
+            self._blocks = blocks
+
+            return self
+        # matrix - vector
+        blocks = []
+
+        for hblock in self._iterator("rows"):
+            out_blocks = [object() for _ in range(hblock._n_blocks[1])]
+            _combine_blocks(hblock._blocks, other._blocks,
+                            Array._subtract, out_blocks)
+            blocks.append(out_blocks)
+
+        compss_delete_object(self._blocks)
+        self._blocks = blocks
+
+        return self
+
+    def __add__(self, other):
+        if self.shape[1] != other.shape[1] or other.shape[0] != 1:
+            raise NotImplementedError("Addition not implemented for the "
+                                      "given objects")
+
+        # matrix + vector
+        blocks = []
+
+        for hblock in self._iterator("rows"):
+            out_blocks = [object() for _ in range(hblock._n_blocks[1])]
+            _combine_blocks(hblock._blocks, other._blocks,
+                            Array._add, out_blocks)
+            blocks.append(out_blocks)
+
+        return Array(blocks, self._top_left_shape, self._reg_shape,
+                     self.shape, self._sparse)
+
+    def __iadd__(self, other):
+        if self.shape[1] != other.shape[1] or (other.shape[0] != 1 and
+                                               (other.shape[0] != self.shape[0]
+                                               or other._reg_shape[0] !=
+                                               self._reg_shape[0])):
+            raise NotImplementedError("Self addition not implemented for the "
+                                      "given objects")
+
+        if other.shape[0] == self.shape[0] and \
+                other._reg_shape[0] == self._reg_shape[0]:
+            if self._reg_shape[1] != other._reg_shape[1]:
+                raise ValueError("incorrect block sizes for the requested "
+                                 f"addition ("
+                                 f"{self._reg_shape[0]}"", "
+                                 f"{self._reg_shape[1]} !="
+                                 f"{other._reg_shape[0]}"", "
+                                 f"{other._reg_shape[1]})")
+
+            if self._top_left_shape != other._top_left_shape:
+                raise ValueError("Incompatible block sizes of the "
+                                 "top left block of the matrices"
+                                 "b._top_left_shape != b._top_left_shape")
+            # matrix + matrix
+            blocks = []
+
+            for hblock, others in zip(self._iterator("rows"),
+                                      other._iterator("rows")):
+                out_blocks = [object() for _ in range(hblock._n_blocks[1])]
+                _combine_blocks(hblock._blocks, others._blocks,
+                                Array._add, out_blocks)
+                blocks.append(out_blocks)
+            compss_delete_object(self._blocks)
+            self._blocks = blocks
+
+            return self
+        # matrix + vector
+        blocks = []
+
+        for hblock in self._iterator("rows"):
+            out_blocks = [object() for _ in range(hblock._n_blocks[1])]
+            _combine_blocks(hblock._blocks, other._blocks,
+                            Array._add, out_blocks)
+            blocks.append(out_blocks)
+        compss_delete_object(self._blocks)
+        self._blocks = blocks
+
+        return self
+
     def __truediv__(self, other):
         if not np.isscalar(other):
             raise NotImplementedError("Non scalar division not supported")
@@ -275,6 +392,22 @@ def _subtract(a, b):
         else:
             return a - b
 
+    @staticmethod
+    def _add(a, b):
+        sparse = issparse(a)
+
+        # needed because subtract with scipy.sparse does not support
+        # broadcasting
+        if sparse:
+            a = a.toarray()
+        if issparse(b):
+            b = b.toarray()
+
+        if sparse:
+            return csr_matrix(a + b)
+        else:
+            return a + b
+
     @staticmethod
     def _power(x_np, power):
         if issparse(x_np):

tests/test_array.py

@@ -1040,6 +1040,177 @@ def test_power(self):
         with self.assertRaises(NotImplementedError):
             x ** x
 
+    def test_add(self):
+        orig = np.array([[1, 2, 3], [4, 5, 6]])
+        x = ds.array(orig, block_size=(2, 1))
+        orig_vector = np.array([[1, 2, 3]])
+        vector = ds.array(orig_vector, block_size=(1, 1))
+
+        b = x + vector
+        self.assertTrue(_validate_array(b))
+        expected = np.array([[2, 4, 6], [5, 7, 9]])
+
+        self.assertTrue(_equal_arrays(expected, b.collect()))
+
+        orig = sp.csr_matrix([[1, 2, 3], [4, 5, 6]])
+        x = ds.array(orig, block_size=(2, 1))
+        orig_vector = sp.csr_matrix([[1, 2, 3]])
+        vector = ds.array(orig_vector, block_size=(1, 1))
+
+        b = x + vector
+        self.assertTrue(_validate_array(b))
+        expected = sp.csr_matrix([[2, 4, 6], [5, 7, 9]])
+
+        self.assertTrue(_equal_arrays(expected, b.collect()))
+
+        with self.assertRaises(NotImplementedError):
+            orig = np.array([[1, 2, 3], [4, 5, 6]])
+            x = ds.array(orig, block_size=(2, 1))
+            orig_vector = np.array([[1, 2]])
+            vector = ds.array(orig_vector, block_size=(1, 1))
+            b = x + vector
+
+    def test_sub(self):
+        orig = np.array([[1, 2, 3], [4, 5, 6]])
+        x = ds.array(orig, block_size=(2, 1))
+        orig_vector = np.array([[1, 2, 3]])
+        vector = ds.array(orig_vector, block_size=(1, 1))
+
+        b = x - vector
+        self.assertTrue(_validate_array(b))
+        expected = np.array([[0, 0, 0], [3, 3, 3]])
+
+        self.assertTrue(_equal_arrays(expected, b.collect()))
+
+        with self.assertRaises(NotImplementedError):
+            orig = np.array([[1, 2, 3], [4, 5, 6]])
+            x = ds.array(orig, block_size=(2, 1))
+            orig_vector = np.array([[1, 3]])
+            vector = ds.array(orig_vector, block_size=(1, 1))
+            b = x - vector
+
+    def test_iadd(self):
+        """ Tests ds-array magic method __iadd__ """
+        orig = np.array([[1, 2, 3], [4, 5, 6]])
+        x = ds.array(orig, block_size=(2, 1))
+        orig_vector = np.array([[1, 2, 3]])
+        vector = ds.array(orig_vector, block_size=(1, 1))
+
+        x += vector
+
+        self.assertTrue(_validate_array(x))
+        self.assertTrue(_validate_array(vector))
+
+        expected = np.array([[2, 4, 6], [5, 7, 9]])
+
+        self.assertTrue(_equal_arrays(expected, x.collect()))
+        self.assertTrue(_equal_arrays(orig_vector, vector.collect()))
+
+        orig = np.array([[1, 2, 3], [4, 5, 6]])
+        x = ds.array(orig, block_size=(2, 1))
+        orig_mat = np.array([[1, 2, 3], [4, 5, 6]])
+        matrix = ds.array(orig_mat, block_size=(2, 1))
+
+        x += matrix
+
+        self.assertTrue(_validate_array(x))
+        self.assertTrue(_validate_array(matrix))
+
+        expected = np.array([[2, 4, 6], [8, 10, 12]])
+
+        self.assertTrue(_equal_arrays(expected, x.collect()))
+        self.assertTrue(_equal_arrays(orig_mat, matrix.collect()))
+
+        with self.assertRaises(NotImplementedError):
+            orig = np.array([[1, 2, 3], [4, 5, 6]])
+            x = ds.array(orig, block_size=(2, 1))
+            orig_vector = np.array([[1, 2]])
+            vector = ds.array(orig_vector, block_size=(1, 1))
+            x += vector
+
+        with self.assertRaises(NotImplementedError):
+            orig = np.array([[1, 2, 3], [4, 5, 6]])
+            x = ds.array(orig, block_size=(2, 1))
+            orig_vector = np.array([[1, 2], [4, 5]])
+            vector = ds.array(orig_vector, block_size=(2, 1))
+            x += vector
+
+        with self.assertRaises(ValueError):
+            x1 = ds.array([[1, 2, 3], [4, 5, 6]], (2, 3))
+            x1.__init__([[1, 2, 3], [4, 5, 6], [7, 8, 9], [1, 2, 3],
+                         [4, 5, 6]], top_left_shape=(1, 3),
+                        reg_shape=(2, 3), shape=(5, 3), sparse=False)
+            x2 = ds.array([[1, 1, 1], [1, 1, 1],
+                           [1, 1, 1], [1, 1, 1],
+                           [1, 1, 1]], (2, 3))
+            x1 += x2
+        with self.assertRaises(ValueError):
+            orig = np.array([[1, 2, 3], [4, 5, 6]])
+            x = ds.array(orig, block_size=(2, 1))
+            orig_vector = np.array([[1, 2, 3], [4, 5, 6]])
+            vector = ds.array(orig_vector, block_size=(2, 2))
+            x += vector
+
+    def test_isub(self):
+        """ Tests ds-array magic method __isub__ """
+        orig = np.array([[1, 2, 3], [4, 5, 6]])
+        x = ds.array(orig, block_size=(2, 1))
+        orig_vector = np.array([[1, 2, 3]])
+        vector = ds.array(orig_vector, block_size=(1, 1))
+
+        x -= vector
+
+        self.assertTrue(_validate_array(x))
+        self.assertTrue(_validate_array(vector))
+
+        expected = np.array([[0, 0, 0], [3, 3, 3]])
+
+        self.assertTrue(_equal_arrays(expected, x.collect()))
+        self.assertTrue(_equal_arrays(orig_vector, vector.collect()))
+
+        orig = np.array([[1, 2, 3], [4, 5, 6]])
+        x = ds.array(orig, block_size=(2, 1))
+        orig_mat = np.array([[1, 2, 3], [4, 5, 6]])
+        matrix = ds.array(orig_mat, block_size=(2, 1))
+
+        x -= matrix
+
+        self.assertTrue(_validate_array(x))
+        self.assertTrue(_validate_array(matrix))
+
+        expected = np.array([[0, 0, 0], [0, 0, 0]])
+
+        self.assertTrue(_equal_arrays(expected, x.collect()))
+        self.assertTrue(_equal_arrays(orig_mat, matrix.collect()))
+        with self.assertRaises(NotImplementedError):
+            orig = np.array([[1, 2, 3], [4, 5, 6]])
+            x = ds.array(orig, block_size=(2, 1))
+            orig_vector = np.array([[1, 2]])
+            vector = ds.array(orig_vector, block_size=(1, 1))
+            x -= vector
+        with self.assertRaises(NotImplementedError):
+            orig = np.array([[1, 2, 3], [4, 5, 6]])
+            x = ds.array(orig, block_size=(2, 1))
+            orig_vector = np.array([[1, 2], [4, 5]])
+            vector = ds.array(orig_vector, block_size=(2, 1))
+            x -= vector
+
+        with self.assertRaises(ValueError):
+            x1 = ds.array([[1, 2, 3], [4, 5, 6]], (2, 3))
+            x1.__init__([[1, 2, 3], [4, 5, 6], [7, 8, 9], [1, 2, 3],
+                         [4, 5, 6]], top_left_shape=(1, 3),
+                        reg_shape=(2, 3), shape=(5, 3), sparse=False)
+            x2 = ds.array([[1, 1, 1], [1, 1, 1],
+                           [1, 1, 1], [1, 1, 1],
+                           [1, 1, 1]], (2, 3))
+            x1 -= x2
+        with self.assertRaises(ValueError):
+            orig = np.array([[1, 2, 3], [4, 5, 6]])
+            x = ds.array(orig, block_size=(2, 1))
+            orig_vector = np.array([[1, 2, 3], [4, 5, 6]])
+            vector = ds.array(orig_vector, block_size=(2, 2))
+            x -= vector
+
     def test_norm(self):
         """ Tests the norm """
         x_np = np.array([[1, 2, 3], [4, 5, 6]])