feat: implement __sub__ feature for DiffractionObject

Author: bobleesj

src/diffpy/utils/diffraction_objects.py

@@ -217,42 +217,16 @@ def __add__(self, other):
 
     __radd__ = __add__
 
-    def _check_operation_compatibility(self, other):
-        if not isinstance(other, (DiffractionObject, int, float)):
-            raise TypeError(invalid_add_type_emsg)
-        if isinstance(other, DiffractionObject):
-            self_yarray = self.all_arrays[:, 0]
-            other_yarray = other.all_arrays[:, 0]
-            if len(self_yarray) != len(other_yarray):
-                raise ValueError(y_grid_length_mismatch_emsg)
-
     def __sub__(self, other):
-        subtracted = deepcopy(self)
-        if isinstance(other, int) or isinstance(other, float) or isinstance(other, np.ndarray):
-            subtracted.on_tth[1] = self.on_tth[1] - other
-            subtracted.on_q[1] = self.on_q[1] - other
-        elif not isinstance(other, DiffractionObject):
-            raise TypeError("I only know how to subtract two Scattering_object objects")
-        elif self.on_tth[0].all() != other.on_tth[0].all():
-            raise RuntimeError(y_grid_length_mismatch_emsg)
-        else:
-            subtracted.on_tth[1] = self.on_tth[1] - other.on_tth[1]
-            subtracted.on_q[1] = self.on_q[1] - other.on_q[1]
-        return subtracted
+        self._check_operation_compatibility(other)
+        subtracted_do = deepcopy(self)
+        if isinstance(other, (int, float)):
+            subtracted_do._all_arrays[:, 0] -= other
+        if isinstance(other, DiffractionObject):
+            subtracted_do._all_arrays[:, 0] -= other.all_arrays[:, 0]
+        return subtracted_do
 
-    def __rsub__(self, other):
-        subtracted = deepcopy(self)
-        if isinstance(other, int) or isinstance(other, float) or isinstance(other, np.ndarray):
-            subtracted.on_tth[1] = other - self.on_tth[1]
-            subtracted.on_q[1] = other - self.on_q[1]
-        elif not isinstance(other, DiffractionObject):
-            raise TypeError("I only know how to subtract two Scattering_object objects")
-        elif self.on_tth[0].all() != other.on_tth[0].all():
-            raise RuntimeError(y_grid_length_mismatch_emsg)
-        else:
-            subtracted.on_tth[1] = other.on_tth[1] - self.on_tth[1]
-            subtracted.on_q[1] = other.on_q[1] - self.on_q[1]
-        return subtracted
+    __rsub__ = __sub__
 
     def __mul__(self, other):
         multiplied = deepcopy(self)
@@ -268,19 +242,10 @@ def __mul__(self, other):
             multiplied.on_q[1] = self.on_q[1] * other.on_q[1]
         return multiplied
 
-    def __rmul__(self, other):
-        multiplied = deepcopy(self)
-        if isinstance(other, int) or isinstance(other, float) or isinstance(other, np.ndarray):
-            multiplied.on_tth[1] = other * self.on_tth[1]
-            multiplied.on_q[1] = other * self.on_q[1]
-        elif self.on_tth[0].all() != other.on_tth[0].all():
-            raise RuntimeError(y_grid_length_mismatch_emsg)
-        else:
-            multiplied.on_tth[1] = self.on_tth[1] * other.on_tth[1]
-            multiplied.on_q[1] = self.on_q[1] * other.on_q[1]
-        return multiplied
+    __rmul__ = __mul__
 
     def __truediv__(self, other):
+
         divided = deepcopy(self)
         if isinstance(other, int) or isinstance(other, float) or isinstance(other, np.ndarray):
             divided.on_tth[1] = other / self.on_tth[1]
@@ -294,17 +259,16 @@ def __truediv__(self, other):
             divided.on_q[1] = self.on_q[1] / other.on_q[1]
         return divided
 
-    def __rtruediv__(self, other):
-        divided = deepcopy(self)
-        if isinstance(other, int) or isinstance(other, float) or isinstance(other, np.ndarray):
-            divided.on_tth[1] = other / self.on_tth[1]
-            divided.on_q[1] = other / self.on_q[1]
-        elif self.on_tth[0].all() != other.on_tth[0].all():
-            raise RuntimeError(y_grid_length_mismatch_emsg)
-        else:
-            divided.on_tth[1] = other.on_tth[1] / self.on_tth[1]
-            divided.on_q[1] = other.on_q[1] / self.on_q[1]
-        return divided
+    __rmul__ = __mul__
+
+    def _check_operation_compatibility(self, other):
+        if not isinstance(other, (DiffractionObject, int, float)):
+            raise TypeError(invalid_add_type_emsg)
+        if isinstance(other, DiffractionObject):
+            self_yarray = self.all_arrays[:, 0]
+            other_yarray = other.all_arrays[:, 0]
+            if self_yarray.shape != other_yarray.shape:
+                raise ValueError(y_grid_length_mismatch_emsg)
 
     @property
     def all_arrays(self):

tests/test_diffraction_objects.py

@@ -713,59 +713,93 @@ def test_copy_object(do_minimal):
 
 
 @pytest.mark.parametrize(
-    "starting_all_arrays, scalar_to_add, expected_all_arrays",
+    "operation, starting_all_arrays, scalar_value, expected_all_arrays",
     [
-        # Test scalar addition to yarray values (intensity) and expect no change to xarrays (q, tth, d)
-        (  # C1: Add integer of 5, expect yarray to increase by by 5
+        # C1: Test scalar addition to yarray values (intensity), expect no change to xarrays (q, tth, d)
+        (  # 1. Add integer 5
+            "add",
             np.array([[1.0, 0.51763809, 30.0, 12.13818192], [2.0, 1.0, 60.0, 6.28318531]]),
             5,
             np.array([[6.0, 0.51763809, 30.0, 12.13818192], [7.0, 1.0, 60.0, 6.28318531]]),
         ),
-        (  # C2: Add float of 5.1, expect yarray to be added by 5.1
+        (  # 2. Add float 5.1
+            "add",
             np.array([[1.0, 0.51763809, 30.0, 12.13818192], [2.0, 1.0, 60.0, 6.28318531]]),
             5.1,
             np.array([[6.1, 0.51763809, 30.0, 12.13818192], [7.1, 1.0, 60.0, 6.28318531]]),
         ),
+        # C2. Test scalar subtraction to yarray values (intensity), expect no change to xarrays (q, tth, d)
+        (  # 1. Subtract integer 1
+            "sub",
+            np.array([[1.0, 0.51763809, 30.0, 12.13818192], [2.0, 1.0, 60.0, 6.28318531]]),
+            1,
+            np.array([[0.0, 0.51763809, 30.0, 12.13818192], [1.0, 1.0, 60.0, 6.28318531]]),
+        ),
+        (  # 2. Subtract float 0.5
+            "sub",
+            np.array([[1.0, 0.51763809, 30.0, 12.13818192], [2.0, 1.0, 60.0, 6.28318531]]),
+            0.5,
+            np.array([[0.5, 0.51763809, 30.0, 12.13818192], [1.5, 1.0, 60.0, 6.28318531]]),
+        ),
     ],
 )
-def test_addition_operator_by_scalar(starting_all_arrays, scalar_to_add, expected_all_arrays, do_minimal_tth):
+def test_scalar_operations(operation, starting_all_arrays, scalar_value, expected_all_arrays, do_minimal_tth):
     do = do_minimal_tth
     assert np.allclose(do.all_arrays, starting_all_arrays)
-    do_scalar_right_sum = do + scalar_to_add
-    assert np.allclose(do_scalar_right_sum.all_arrays, expected_all_arrays)
-    do_scalar_left_sum = scalar_to_add + do
-    assert np.allclose(do_scalar_left_sum.all_arrays, expected_all_arrays)
+
+    if operation == "add":
+        result_right = do + scalar_value
+        result_left = scalar_value + do
+    elif operation == "sub":
+        result_right = do - scalar_value
+        result_left = scalar_value - do
+
+    assert np.allclose(result_right.all_arrays, expected_all_arrays)
+    assert np.allclose(result_left.all_arrays, expected_all_arrays)
 
 
 @pytest.mark.parametrize(
-    "do_1_all_arrays, "
-    "do_2_all_arrays, "
-    "expected_do_1_all_arrays_with_y_summed, "
-    "expected_do_2_all_arrays_with_y_summed",
+    "operation, " "expected_do_1_all_arrays_with_y_modified, " "expected_do_2_all_arrays_with_y_modified",
     [
-        # Test addition of two DO objects, expect combined yarray values and no change to xarrays ((q, tth, d)
-        (  # C1: Add two DO objects, expect sum of yarray values
-            (np.array([[1.0, 0.51763809, 30.0, 12.13818192], [2.0, 1.0, 60.0, 6.28318531]]),),
-            (np.array([[1.0, 6.28318531, 100.70777771, 1], [2.0, 3.14159265, 45.28748053, 2.0]]),),
-            (np.array([[2.0, 0.51763809, 30.0, 12.13818192], [4.0, 1.0, 60.0, 6.28318531]]),),
-            (np.array([[2.0, 6.28318531, 100.70777771, 1], [4.0, 3.14159265, 45.28748053, 2.0]]),),
+        # Test addition of two DO objects, expect combined yarray values
+        (
+            "add",
+            np.array([[2.0, 0.51763809, 30.0, 12.13818192], [4.0, 1.0, 60.0, 6.28318531]]),
+            np.array([[2.0, 6.28318531, 100.70777771, 1], [4.0, 3.14159265, 45.28748053, 2.0]]),
+        ),
+        # Test subtraction of two DO objects, expect differences in yarray values
+        (
+            "sub",
+            np.array([[0.0, 0.51763809, 30.0, 12.13818192], [0.0, 1.0, 60.0, 6.28318531]]),
+            np.array([[0.0, 6.28318531, 100.70777771, 1], [0.0, 3.14159265, 45.28748053, 2.0]]),
         ),
     ],
 )
-def test_addition_operator_by_another_do(
-    do_1_all_arrays,
-    do_2_all_arrays,
-    expected_do_1_all_arrays_with_y_summed,
-    expected_do_2_all_arrays_with_y_summed,
+def test_binary_operator_on_do(
+    operation,
+    expected_do_1_all_arrays_with_y_modified,
+    expected_do_2_all_arrays_with_y_modified,
     do_minimal_tth,
     do_minimal_d,
 ):
     do_1 = do_minimal_tth
-    assert np.allclose(do_1.all_arrays, do_1_all_arrays)
     do_2 = do_minimal_d
-    assert np.allclose(do_2.all_arrays, do_2_all_arrays)
-    assert np.allclose((do_1 + do_2).all_arrays, expected_do_1_all_arrays_with_y_summed)
-    assert np.allclose((do_2 + do_1).all_arrays, expected_do_2_all_arrays_with_y_summed)
+    assert np.allclose(
+        do_1.all_arrays, np.array([[1.0, 0.51763809, 30.0, 12.13818192], [2.0, 1.0, 60.0, 6.28318531]])
+    )
+    assert np.allclose(
+        do_2.all_arrays, np.array([[1.0, 6.28318531, 100.70777771, 1], [2.0, 3.14159265, 45.28748053, 2.0]])
+    )
+
+    if operation == "add":
+        do_1_y_modified = do_1 + do_2
+        do_2_y_modified = do_2 + do_1
+    elif operation == "sub":
+        do_1_y_modified = do_1 - do_2
+        do_2_y_modified = do_2 - do_1
+
+    assert np.allclose(do_1_y_modified.all_arrays, expected_do_1_all_arrays_with_y_modified)
+    assert np.allclose(do_2_y_modified.all_arrays, expected_do_2_all_arrays_with_y_modified)
 
 
 def test_addition_operator_invalid_type(do_minimal_tth, invalid_add_type_error_msg):
@@ -775,6 +809,10 @@ def test_addition_operator_invalid_type(do_minimal_tth, invalid_add_type_error_m
         do + "string_value"
     with pytest.raises(TypeError, match=re.escape(invalid_add_type_error_msg)):
         "string_value" + do
+    with pytest.raises(TypeError, match=re.escape(invalid_add_type_error_msg)):
+        do - "string_value"
+    with pytest.raises(TypeError, match=re.escape(invalid_add_type_error_msg)):
+        "string_value" - do
 
 
 def test_addition_operator_invalid_yarray_length(do_minimal, do_minimal_tth, y_grid_size_mismatch_error_msg):
@@ -785,3 +823,5 @@ def test_addition_operator_invalid_yarray_length(do_minimal, do_minimal_tth, y_g
     assert len(do_2.all_arrays[:, 0]) == 2
     with pytest.raises(ValueError, match=re.escape(y_grid_size_mismatch_error_msg)):
         do_1 + do_2
+    with pytest.raises(ValueError, match=re.escape(y_grid_size_mismatch_error_msg)):
+        do_1 - do_2