From 6df7654064cc7c6f114d8bdfe506d72772ed7a1d Mon Sep 17 00:00:00 2001
From: Michael Krayer <michael.krayer@posteo.de>
Date: Fri, 15 Oct 2021 15:04:22 +0200
Subject: [PATCH] arithmetic operations

---
 parallel.py | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/parallel.py b/parallel.py
index a213e15..4250677 100644
--- a/parallel.py
+++ b/parallel.py
@@ -295,17 +295,42 @@ class PPP:
 
     #def vorticity(self,axis=None,key_out=None,keep_derivatives=False,on_pressure_grid=True):
 
-    def multiply(self,key1,key2,key_out=None,on_pressure_grid=True):
+    def arithmetic_operation(self,key1,operation,key2,key_out=None,on_pressure_grid=True):
+        import operator
+        if operation in ('add','+'):
+            op = operator.iadd
+        elif operation in ('subtract','sub','-'):
+            op = operator.isub
+        elif operation in ('divide','div','/'):
+            op = operator.itruediv
+        elif operation in ('multiply','mul','*'):
+            op = operator.imul
         if key_out is None:
             key_out = key1
         if not self.field[key1].has_same_grid(self.field[key2]) or on_pressure_grid:
             self.shift_to_pressure_grid(key1,key_out)
             self.shift_to_pressure_grid(key2,'tmp')
-            self.field[key_out] *= self.field['tmp']
+            op(self.field[key_out],self.field['tmp'])
             self.delete('tmp')
         else:
             self.copy(key1,key_out)
-            self.field[key_out] *= self.field[key2]
+            op(self.field[key_out],self.field[key2])
+        return
+
+    def add(self,key1,key2,key_out=None,on_pressure_grid=True):
+        self.arithmetic_operation(key1,'+',key2,key_out=key_out,on_pressure_grid=on_pressure_grid)
+        return
+    
+    def subtract(self,key1,key2,key_out=None,on_pressure_grid=True):
+        self.arithmetic_operation(key1,'-',key2,key_out=key_out,on_pressure_grid=on_pressure_grid)
+        return
+    
+    def multiply(self,key1,key2,key_out=None,on_pressure_grid=True):
+        self.arithmetic_operation(key1,'*',key2,key_out=key_out,on_pressure_grid=on_pressure_grid)
+        return
+    
+    def divide(self,key1,key2,key_out=None,on_pressure_grid=True):
+        self.arithmetic_operation(key1,'/',key2,key_out=key_out,on_pressure_grid=on_pressure_grid)
         return
 
     def gaussian_filter(self,key,sigma,truncate=4.0,key_out=None,iterate=False,verbose=True):