Merge pull request #280 from cosanlab/arith

added vector multiplication to Brain_Data

nltools/data/brain_data.py

@@ -235,26 +235,66 @@ def __add__(self, y):
             new.data = new.data + y.data
         return new
 
+    def __radd__(self, y):
+        new = deepcopy(self)
+        if isinstance(y, (int, float)):
+            new.data = y + new.data
+        elif isinstance(y, Brain_Data):
+            if self.shape() != y.shape():
+                raise ValueError("Both Brain_Data() instances need to be the "
+                                 "same shape.")
+            new.data = y.data + new.data
+        return new
+
     def __sub__(self, y):
         new = deepcopy(self)
         if isinstance(y, (int, float)):
             new.data = new.data - y
-        if isinstance(y, Brain_Data):
+        elif isinstance(y, Brain_Data):
             if self.shape() != y.shape():
                 raise ValueError('Both Brain_Data() instances need to be the '
                                  'same shape.')
             new.data = new.data - y.data
         return new
 
+    def __rsub__(self, y):
+        new = deepcopy(self)
+        if isinstance(y, (int, float)):
+            new.data = y - new.data
+        elif isinstance(y, Brain_Data):
+            if self.shape() != y.shape():
+                raise ValueError('Both Brain_Data() instances need to be the '
+                                 'same shape.')
+            new.data = y.data - new.data
+        return new
+
     def __mul__(self, y):
         new = deepcopy(self)
         if isinstance(y, (int, float)):
             new.data = new.data * y
-        if isinstance(y, Brain_Data):
+        elif isinstance(y, Brain_Data):
             if self.shape() != y.shape():
                 raise ValueError("Both Brain_Data() instances need to be the "
                                  "same shape.")
             new.data = np.multiply(new.data, y.data)
+        elif isinstance(y, (list, np.ndarray, np.array)):
+            if len(y) != len(self):
+                raise ValueError('Vector multiplication requires that the '
+                                 'length of the vector match the number of '
+                                 'images in Brain_Data instance.')
+            else:
+                new.data = np.dot(new.data.T, y).T
+        return new
+
+    def __rmul__(self, y):
+        new = deepcopy(self)
+        if isinstance(y, (int, float)):
+            new.data = y * new.data
+        elif isinstance(y, Brain_Data):
+            if self.shape() != y.shape():
+                raise ValueError("Both Brain_Data() instances need to be the "
+                                 "same shape.")
+            new.data = np.multiply(y.data, new.data)
         return new
 
     def __iter__(self):

nltools/tests/test_brain_data.py

@@ -68,17 +68,24 @@ def test_sum(sim_brain_data):
 def test_add(sim_brain_data):
     new = sim_brain_data + sim_brain_data
     assert new.shape() == shape_2d
-
+    value = 10
+    assert(value + sim_brain_data[0]).mean() == (sim_brain_data[0] + value).mean()
 
 def test_subtract(sim_brain_data):
     new = sim_brain_data - sim_brain_data
     assert new.shape() == shape_2d
-
+    value = 10
+    assert (-value-(-1)*sim_brain_data[0]).mean() == (sim_brain_data[0]-value).mean()
 
 def test_multiply(sim_brain_data):
     new = sim_brain_data * sim_brain_data
     assert new.shape() == shape_2d
-
+    value = 10
+    assert(value * sim_brain_data[0]).mean() == (sim_brain_data[0] * value).mean()
+    c1 = [.5, .5, -.5, -.5]
+    new = sim_brain_data[0:4]*c1
+    new2 = sim_brain_data[0]*.5 + sim_brain_data[1]*.5 - sim_brain_data[2]*.5 - sim_brain_data[3]*.5
+    np.testing.assert_almost_equal(0, (new-new2).sum(), decimal=5)
 
 def test_indexing(sim_brain_data):
     index = [0, 3, 1]