ADD unary operator

Author: htwangtw

nibabel/arrayops.py

@@ -15,7 +15,7 @@
 
 
 class OperableImage:
-    def _op(self, other, op):
+    def _binop(self, val, *, op):
         """Apply operator to Nifti1Image.
 
         Arithmetic and logical operation on Nifti image.
@@ -28,49 +28,85 @@ def _op(self, other, op):
         op :
             Python operator.
         """
-        # Check orientations are the same
-        if aff2axcodes(self.affine) != aff2axcodes(other.affine):
-            raise ValueError("Two images should have the same orientation")
-        # Check affine
-        if (self.affine != other.affine).all():
-            raise ValueError("Two images should have the same affine.")
-        # Check shape. Handle identical stuff for now.
-        if self.shape != other.shape:
-            raise ValueError("Two images should have the same shape.")
-
-        # Check types? Problematic types will be caught by numpy,
-        # but might be cheaper to check before loading data.
-        # collect dtype
-        dtypes = [img.get_data_dtype().type for img in (self, other)]
-        # check allowed dtype based on the operator
-        if set(support_np_type).union(dtypes) == 0:
-            raise ValueError("Image contains illegal datatype for Nifti1Image.")
-
+        val = _input_validation(self, val)
+        # numerical operator should work work
         if op.__name__ in ["add", "sub", "mul", "truediv", "floordiv"]:
-            dataobj = op(np.asanyarray(self.dataobj), np.asanyarray(other.dataobj))
+            dataobj = op(np.asanyarray(self.dataobj), val)
         if op.__name__ in ["and_", "or_"]:
             self_ = self.dataobj.astype(bool)
-            other_ = other.dataobj.astype(bool)
+            other_ = val.astype(bool)
             dataobj = op(self_, other_).astype(int)
         return self.__class__(dataobj, self.affine, self.header)
 
+
+    def _unop(self, *, op):
+        """
+        Parameters
+        ----------
+        op :
+            Python operator.
+        """
+        _type_check(self)
+        if op.__name__ in ["pos", "neg", "abs"]:
+            dataobj = op(np.asanyarray(self.dataobj))
+        return self.__class__(dataobj, self.affine, self.header)
+
+
     def __add__(self, other):
-        return self._op(other, operator.__add__)
+        return self._binop(other, op=operator.__add__)
 
     def __sub__(self, other):
-        return self._op(other, operator.__sub__)
+        return self._binop(other, op=operator.__sub__)
 
     def __mul__(self, other):
-        return self._op(other, operator.__mul__)
+        return self._binop(other, op=operator.__mul__)
 
     def __truediv__(self, other):
-        return self._op(other, operator.__truediv__)
+        return self._binop(other, op=operator.__truediv__)
 
     def __floordiv__(self, other):
-        return self._op(other, operator.__floordiv__)
+        return self._binop(other, op=operator.__floordiv__)
 
     def __and__(self, other):
-        return self._op(other, operator.__and__)
+        return self._binop(other, op=operator.__and__)
 
     def __or__(self, other):
-        return self._op(other, operator.__or__)
+        return self._binop(other, op=operator.__or__)
+
+    def __pos__(self):
+        return self._unop(op=operator.__pos__)
+
+    def __neg__(self):
+        return self._unop(op=operator.__neg__)
+
+    def __abs__(self):
+        return self._unop(op=operator.__abs__)
+
+
+
+def _input_validation(self, val):
+    """Check images orientation, affine, and shape muti-images operation."""
+    _type_check(self)
+    if type(val) not in [float, int]:
+        # Check orientations are the same
+        if aff2axcodes(self.affine) != aff2axcodes(val.affine):
+            raise ValueError("Two images should have the same orientation")
+        # Check affine
+        if (self.affine != val.affine).all():
+            raise ValueError("Two images should have the same affine.")
+        # Check shape.
+        if self.shape[:3] != val.shape[:3]:
+            raise ValueError("Two images should have the same shape except"
+                                "the time dimension.")
+
+        _type_check(val)
+        val = np.asanyarray(val.dataobj)
+    return val
+
+def _type_check(*args):
+    """Ensure image contains correct nifti data type."""
+    # Check types
+    dtypes = [img.get_data_dtype().type for img in args]
+    # check allowed dtype based on the operator
+    if set(support_np_type).union(dtypes) == 0:
+        raise ValueError("Image contains illegal datatype for Nifti1Image.")

nibabel/tests/test_arrayops.py

@@ -3,12 +3,17 @@
 from numpy.testing import assert_array_equal
 import pytest
 
-def test_operations():
+
+def test_binary_operations():
     data1 = np.random.rand(5, 5, 2)
     data2 = np.random.rand(5, 5, 2)
     data1[0, 0, :] = 0
     img1 = Nifti1Image(data1, np.eye(4))
     img2 = Nifti1Image(data2, np.eye(4))
+
+    output = img1 + 2
+    assert_array_equal(output.dataobj, data1 + 2)
+
     output = img1 + img2
     assert_array_equal(output.dataobj, data1 + data2)
 
@@ -37,4 +42,15 @@ def test_operations():
     assert_array_equal(output.dataobj, (data1.astype(bool) & data2.astype(bool)).astype(int))
 
     output = img1 | img2
-    assert_array_equal(output.dataobj, (data1.astype(bool) | data2.astype(bool)).astype(int))
+    assert_array_equal(output.dataobj, (data1.astype(bool) | data2.astype(bool)).astype(int))
+
+
+def test_unary_operations():
+    data = np.random.rand(5, 5, 2)
+    img = Nifti1Image(data, np.eye(4))
+
+    output = -img
+    assert_array_equal(output.dataobj, -data)
+
+    output = abs(img)
+    assert_array_equal(output.dataobj, abs(data))