More fixes for the Python fields module (#2502)

Fixed indexing for PML fields.
Add all of the fixes to the setitems routines.

1 files changed, 64 insertions, 60 deletions

diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py
index c4b0e29c6..41f2751bc 100644
--- a/Python/pywarpx/fields.py
+++ b/Python/pywarpx/fields.py
@@ -140,11 +140,11 @@ class _MultiFABWrapper(object):
         elif self.dim == 3:
             return self._getitem3d(index)
 
-    def _find_start_stop(self, ii, nn, ngrow, d):
+    def _find_start_stop(self, ii, imin, imax, d):
         """Given the input index, calculate the start and stop range of the indices.
-        - ii: index, either a slice object or an integer
-        - nn: the global number of cells along the specified direction
-        - ngrow: list of the number of extra cells, i.e. guard cells
+        - ii: input index, either a slice object or an integer
+        - imin: the global lowest lovect value in the specified diretion
+        - imax: the global highest hivect value in the specified diretion
         - d: the direction, an integer, 0, 1, or 2
         If ii is a slice, the start and stop values are used directly,
         unless they are None, then the lower or upper bound is used.
@@ -152,18 +152,18 @@ class _MultiFABWrapper(object):
         """
         if isinstance(ii, slice):
             if ii.start is None:
-                iistart = -ngrow[d]
+                iistart = imin
             else:
                 iistart = ii.start
             if ii.stop is None:
-                iistop = nn + self.overlaps[d] + ngrow[d]
+                iistop = imax + self.overlaps[d]
             else:
                 iistop = ii.stop
         else:
             iistart = ii
             iistop = ii + 1
-        assert -ngrow[d] <= iistart <= nn + self.overlaps[d] + ngrow[d], Exception(f'Dimension {d} lower index is out of bounds')
-        assert -ngrow[d] <= iistop <= nn + self.overlaps[d] + ngrow[d], Exception(f'Dimension {d} upper index is out of bounds')
+        assert imin <= iistart <= imax + self.overlaps[d], Exception(f'Dimension {d} lower index is out of bounds')
+        assert imin <= iistop <= imax + self.overlaps[d], Exception(f'Dimension {d} upper index is out of bounds')
         return iistart, iistop
 
     def _getitem3d(self, index):
@@ -191,18 +191,24 @@ class _MultiFABWrapper(object):
         else:
             ic = None
 
-        nx = hivects[0,:].max() - ngrow[0]
-        ny = hivects[1,:].max() - ngrow[1]
-        nz = hivects[2,:].max() - ngrow[2]
+        ixmin = lovects[0,:].min()
+        ixmax = hivects[0,:].max()
+        iymin = lovects[1,:].min()
+        iymax = hivects[1,:].max()
+        izmin = lovects[2,:].min()
+        izmax = hivects[2,:].max()
 
         if npes > 1:
-            nx = comm_world.allreduce(nx, op=mpi.MAX)
-            ny = comm_world.allreduce(ny, op=mpi.MAX)
-            nz = comm_world.allreduce(nz, op=mpi.MAX)
+            ixmin = comm_world.allreduce(ixmin, op=mpi.MIN)
+            ixmax = comm_world.allreduce(ixmax, op=mpi.MAX)
+            iymin = comm_world.allreduce(iymin, op=mpi.MIN)
+            iymax = comm_world.allreduce(iymax, op=mpi.MAX)
+            izmin = comm_world.allreduce(izmin, op=mpi.MIN)
+            izmax = comm_world.allreduce(izmax, op=mpi.MAX)
 
-        ixstart, ixstop = self._find_start_stop(ix, nx, ngrow, 0)
-        iystart, iystop = self._find_start_stop(iy, ny, ngrow, 1)
-        izstart, izstop = self._find_start_stop(iz, nz, ngrow, 2)
+        ixstart, ixstop = self._find_start_stop(ix, ixmin, ixmax, 0)
+        iystart, iystop = self._find_start_stop(iy, iymin, iymax, 1)
+        izstart, izstop = self._find_start_stop(iz, izmin, izmax, 2)
 
         # --- Setup the size of the array to be returned and create it.
         # --- Space is added for multiple components if needed.
@@ -282,15 +288,19 @@ class _MultiFABWrapper(object):
         else:
             ic = None
 
-        nx = hivects[0,:].max() - ngrow[0]
-        nz = hivects[1,:].max() - ngrow[1]
+        ixmin = lovects[0,:].min()
+        ixmax = hivects[0,:].max()
+        izmin = lovects[1,:].min()
+        izmax = hivects[1,:].max()
 
         if npes > 1:
-            nx = comm_world.allreduce(nx, op=mpi.MAX)
-            nz = comm_world.allreduce(nz, op=mpi.MAX)
+            ixmin = comm_world.allreduce(ixmin, op=mpi.MIN)
+            ixmax = comm_world.allreduce(ixmax, op=mpi.MAX)
+            izmin = comm_world.allreduce(izmin, op=mpi.MIN)
+            izmax = comm_world.allreduce(izmax, op=mpi.MAX)
 
-        ixstart, ixstop = self._find_start_stop(ix, nx, ngrow, 0)
-        izstart, izstop = self._find_start_stop(iz, nz, ngrow, 1)
+        ixstart, ixstop = self._find_start_stop(ix, ixmin, ixmax, 0)
+        izstart, izstop = self._find_start_stop(iz, izmin, izmax, 1)
 
         # --- Setup the size of the array to be returned and create it.
         # --- Space is added for multiple components if needed.
@@ -379,9 +389,24 @@ class _MultiFABWrapper(object):
         else:
             ic = None
 
-        nx = hivects[0,:].max() - ngrow[0]
-        ny = hivects[1,:].max() - ngrow[1]
-        nz = hivects[2,:].max() - ngrow[2]
+        ixmin = lovects[0,:].min()
+        ixmax = hivects[0,:].max()
+        iymin = lovects[1,:].min()
+        iymax = hivects[1,:].max()
+        izmin = lovects[2,:].min()
+        izmax = hivects[2,:].max()
+
+        if npes > 1:
+            ixmin = comm_world.allreduce(ixmin, op=mpi.MIN)
+            ixmax = comm_world.allreduce(ixmax, op=mpi.MAX)
+            iymin = comm_world.allreduce(iymin, op=mpi.MIN)
+            iymax = comm_world.allreduce(iymax, op=mpi.MAX)
+            izmin = comm_world.allreduce(izmin, op=mpi.MIN)
+            izmax = comm_world.allreduce(izmax, op=mpi.MAX)
+
+        ixstart, ixstop = self._find_start_stop(ix, ixmin, ixmax, 0)
+        iystart, iystop = self._find_start_stop(iy, iymin, iymax, 1)
+        izstart, izstop = self._find_start_stop(iz, izmin, izmax, 2)
 
         # --- Add extra dimensions so that the input has the same number of
         # --- dimensions as array.
@@ -393,25 +418,6 @@ class _MultiFABWrapper(object):
             if not isinstance(iz, slice): sss[2:2] = [1]
             value3d.shape = sss
 
-        if isinstance(ix, slice):
-            ixstart = max(ix.start or -ngrow[0], -ngrow[0])
-            ixstop = min(ix.stop or nx + 1 + ngrow[0], nx + self.overlaps[0] + ngrow[0])
-        else:
-            ixstart = ix
-            ixstop = ix + 1
-        if isinstance(iy, slice):
-            iystart = max(iy.start or -ngrow[1], -ngrow[1])
-            iystop = min(iy.stop or ny + 1 + ngrow[1], ny + self.overlaps[1] + ngrow[1])
-        else:
-            iystart = iy
-            iystop = iy + 1
-        if isinstance(iz, slice):
-            izstart = max(iz.start or -ngrow[2], -ngrow[2])
-            izstop = min(iz.stop or nz + 1 + ngrow[2], nz + self.overlaps[2] + ngrow[2])
-        else:
-            izstart = iz
-            izstop = iz + 1
-
         for i in range(len(fields)):
 
             # --- The ix1, 2 etc are relative to global indexing
@@ -461,8 +467,19 @@ class _MultiFABWrapper(object):
         else:
             ic = None
 
-        nx = hivects[0,:].max() - ngrow[0]
-        nz = hivects[1,:].max() - ngrow[1]
+        ixmin = lovects[0,:].min()
+        ixmax = hivects[0,:].max()
+        izmin = lovects[1,:].min()
+        izmax = hivects[1,:].max()
+
+        if npes > 1:
+            ixmin = comm_world.allreduce(ixmin, op=mpi.MIN)
+            ixmax = comm_world.allreduce(ixmax, op=mpi.MAX)
+            izmin = comm_world.allreduce(izmin, op=mpi.MIN)
+            izmax = comm_world.allreduce(izmax, op=mpi.MAX)
+
+        ixstart, ixstop = self._find_start_stop(ix, ixmin, ixmax, 0)
+        izstart, izstop = self._find_start_stop(iz, izmin, izmax, 1)
 
         # --- Add extra dimensions so that the input has the same number of
         # --- dimensions as array.
@@ -473,19 +490,6 @@ class _MultiFABWrapper(object):
             if not isinstance(iz, slice): sss[1:1] = [1]
             value3d.shape = sss
 
-        if isinstance(ix, slice):
-            ixstart = max(ix.start or -ngrow[0], -ngrow[0])
-            ixstop = min(ix.stop or nx + 1 + ngrow[0], nx + self.overlaps[0] + ngrow[0])
-        else:
-            ixstart = ix
-            ixstop = ix + 1
-        if isinstance(iz, slice):
-            izstart = max(iz.start or -ngrow[1], -ngrow[1])
-            izstop = min(iz.stop or nz + 1 + ngrow[1], nz + self.overlaps[1] + ngrow[1])
-        else:
-            izstart = iz
-            izstop = iz + 1
-
         for i in range(len(fields)):
 
             # --- The ix1, 2 etc are relative to global indexing