add functin to subtract qnArrayTwo objects

xnx · Jul 4, 2017 · 30f8afc · 30f8afc
1 parent 33868f3
commit 30f8afc
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/pyqn/qn_array_two.py b/pyqn/qn_array_two.py
@@ -55,3 +55,23 @@ def __add__(self, other):
                 temp_q = Quantity(value=self[i], units=self.units) + other
                 v.append(temp_q.value)
             return qnArrayTwo(v, units = temp_q.units)
+        else:
+            raise qnArrayTwoError("Can only add two qnArray objects or a qnArray with Quantity")
+
+    def __sub__(self, other):
+        if type(other) is qnArrayTwo:
+            if len(other) != len(self):
+                raise qnArrayTwoError("Inconsistent array lengths")
+            v = []
+            for i in range(len(self)):
+                temp_q = Quantity(value = self[i], units = self.units) - Quantity(value = other[i], units = other.units)
+                v.append(temp_q.value)
+            return qnArrayTwo(v, units = temp_q.units)
+        if type(other) is Quantity:
+            v = []
+            for i in range(len(self)):
+                temp_q = Quantity(value=self[i], units=self.units) - other
+                v.append(temp_q.value)
+            return qnArrayTwo(v, units = temp_q.units)
+        else:
+            raise qnArrayTwoError("Can only add two qnArray objects or a qnArray with Quantity")