From fcca37ad93f9d99ca5d285727f10124e987d183a Mon Sep 17 00:00:00 2001
From: Michael Krayer <michael.krayer@posteo.de>
Date: Sun, 19 Apr 2020 18:15:58 +0200
Subject: [PATCH] simplified maskParticles which lead to 5x speedup

---
 python/ibmppp/ibmppp.py | 43 +++++++++++++++++++----------------------
 1 file changed, 20 insertions(+), 23 deletions(-)

diff --git a/python/ibmppp/ibmppp.py b/python/ibmppp/ibmppp.py
index daa185c..104a8e9 100644
--- a/python/ibmppp/ibmppp.py
+++ b/python/ibmppp/ibmppp.py
@@ -1092,6 +1092,10 @@ class ibmppp:
 
     def maskParticles(self,key,reconstruct=False,cval=np.nan,padding=0.0):
       '''Fills grid points which lie inside of solid phase with values.'''
+      # Shortcut to local grid
+      xg = self.localGrid[key][0]
+      yg = self.localGrid[key][1]
+      zg = self.localGrid[key][2]
       # Loop through local particles
       for ipart in range(0,self.__npartl):
         # Slice a box from the field around the particle
@@ -1120,29 +1124,22 @@ class ibmppp:
             coeff_roty =  0.0
             coeff_rotz =  0.0
         # Get bounding box of particle
-        idx_x = np.nonzero((self.localGrid[key][0]>=xp-rp) & (self.localGrid[key][0]<=xp+rp))[0]
-        idx_y = np.nonzero((self.localGrid[key][1]>=yp-rp) & (self.localGrid[key][1]<=yp+rp))[0]
-        idx_z = np.nonzero((self.localGrid[key][2]>=zp-rp) & (self.localGrid[key][2]<=zp+rp))[0]
-        # Slice the bounding box
-        sli_x = slice(idx_x[0],idx_x[-1]+1)
-        sli_y = slice(idx_y[0],idx_y[-1]+1)
-        sli_z = slice(idx_z[0],idx_z[-1]+1)
-        # Construct a grid of the subarray
-        xg,yg,zg = np.meshgrid(self.localGrid[key][0][sli_x],self.localGrid[key][1][sli_y],self.localGrid[key][2][sli_z],indexing='ij',copy=True)
-        # Iterate through subarray and mask field
-        it = np.nditer((self.field[key][sli_x,sli_y,sli_z],xg,yg,zg),
-                       op_flags=[['writeonly'],['readonly'],['readonly'],['readonly']]
-                       )
-        while not it.finished:
-          xx,yy,zz = it[1:4]
-          Dx,Dy,Dz = xx-xp, yy-yp, zz-zp
-          isInside = Dx**2+Dy**2+Dz**2 <= rp**2
-          if isInside:
-            if reconstruct:
-              it[0] = coeff_lin + coeff_rotx*Dx + coeff_roty*Dy + coeff_rotz*Dz
-            else:
-              it[0] = cval
-          it.iternext()
+        idx_x = np.nonzero((xg>=xp-rp) & (xg<=xp+rp))[0]
+        idx_y = np.nonzero((yg>=yp-rp) & (yg<=yp+rp))[0]
+        idx_z = np.nonzero((zg>=zp-rp) & (zg<=zp+rp))[0]
+        # Triple for loop
+        for ii in range(idx_x[0],idx_x[-1]+1):
+          Dx = xg[ii]-xp
+          for jj in range(idx_y[0],idx_y[-1]+1):
+            Dy = yg[jj]-yp
+            for kk in range(idx_z[0],idx_z[-1]+1):
+              Dz = zg[kk]-zp
+              isInside = Dx*Dx+Dy*Dy+Dz*Dz <= rp*rp
+              if isInside:
+                if reconstruct:
+                  self.field[key][ii,jj,kk] = coeff_lin + coeff_rotx*Dx + coeff_roty*Dy + coeff_rotz*Dz
+                else:
+                  self.field[key][ii,jj,kk] = cval
 
     def shiftField(self,key,shift,subsequent=True):
       '''Shifts a field by dx/2 in specified direction as in-place operation. The new values are determined by linear interpolation.'''