Updated Python interface, adding WarpInterface

1 files changed, 266 insertions, 74 deletions

diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py
index 8237ed867..9068a22ea 100644
--- a/Python/pywarpx/fields.py
+++ b/Python/pywarpx/fields.py
@@ -11,7 +11,7 @@ import numpy as np
 from . import _libwarpx
 
 
-class MultiFABWrapper(object):
+class _MultiFABWrapper(object):
     """Wrapper around field arrays at level 0
     This provides a convenient way to query and set fields that are broken up into FABs.
     The indexing is based on global indices.
@@ -21,16 +21,23 @@ class MultiFABWrapper(object):
      - get_fabs: routine that returns the list of FABs
      - level=0: ignored
     """
-    def __init__(self, direction, overlaps, get_lovects, get_fabs, level=0):
+    def __init__(self, direction, overlaps, get_lovects, get_fabs, level=0, include_ghosts=False):
         self.direction = direction
         self.overlaps = np.array(overlaps)
         self.get_lovects = get_lovects
         self.get_fabs = get_fabs
         self.level = 0 #level
-        self.include_ghosts = False
+        self.include_ghosts = include_ghosts
+
+        self.dim = _libwarpx.dim
+        if self.dim == 2:
+            # --- Grab the first and last overlaps (for x and z)
+            self.overlaps = self.overlaps[::2]
 
     def _getlovects(self):
-        return self.get_lovects(self.level, self.direction, self.include_ghosts)
+        lovects = self.get_lovects(self.level, self.direction, self.include_ghosts)
+        self.nghosts = -lovects.min()
+        return lovects
 
     def _gethivects(self):
         lovects = self._getlovects()
@@ -38,7 +45,7 @@ class MultiFABWrapper(object):
 
         hivects = np.zeros_like(lovects)
         for i in range(len(fields)):
-            hivects[:,i] = lovects[:,i] + np.array(fields[i].shape) - self.overlaps
+            hivects[:,i] = lovects[:,i] + np.array(fields[i].shape[:self.dim]) - self.overlaps
 
         return hivects
 
@@ -50,14 +57,29 @@ class MultiFABWrapper(object):
 
     def __getitem__(self, index):
         """Returns slices of a decomposed array, The shape of
-      the object returned depends on the number of ix, iy and iz specified, which
-      can be from none to all three. Note that the values of ix, iy and iz are
-      relative to the fortran indexing, meaning that 0 is the lower boundary
-      of the domain.
+        the object returned depends on the number of ix, iy and iz specified, which
+        can be from none to all three. Note that the values of ix, iy and iz are
+        relative to the fortran indexing, meaning that 0 is the lower boundary
+        of the whole domain.
         """
         if index == Ellipsis:
-            index = (slice(None), slice(None), slice(None))
-
+            index = tuple(self.dim*[slice(None)])
+
+        if len(index) < self.dim:
+            # --- Add extra dims to index if needed
+            index = list(index)
+            for i in range(len(index), self.dim):
+                index.append(slice(None))
+            index = tuple(index)
+
+        if self.dim == 2:
+            return self._getitem2d(index)
+        elif self.dim == 3:
+            return self._getitem3d(index)
+
+    def _getitem3d(self, index):
+        """Returns slices of a 3D decomposed array,
+        """
         ix = index[0]
         iy = index[1]
         iz = index[2]
@@ -66,25 +88,25 @@ class MultiFABWrapper(object):
         hivects = self._gethivects()
         fields = self._getfields()
 
-        nx = hivects[0,:].max()
-        ny = hivects[1,:].max()
-        nz = hivects[2,:].max()
+        nx = hivects[0,:].max() - self.nghosts
+        ny = hivects[1,:].max() - self.nghosts
+        nz = hivects[2,:].max() - self.nghosts
 
         if isinstance(ix, slice):
-            ixstart = max(ix.start, 0)
-            ixstop = min(ix.stop or nx + 1, nx + self.overlaps[0])
+            ixstart = max(ix.start or -self.nghosts, -self.nghosts)
+            ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts)
         else:
             ixstart = ix
             ixstop = ix + 1
         if isinstance(iy, slice):
-            iystart = max(iy.start, 0)
-            iystop = min(iy.stop or ny + 1, ny + self.overlaps[1])
+            iystart = max(iy.start or -self.nghosts, -self.nghosts)
+            iystop = min(iy.stop or ny + 1 + self.nghosts, ny + self.overlaps[1] + self.nghosts)
         else:
             iystart = iy
             iystop = iy + 1
         if isinstance(iz, slice):
-            izstart = max(iz.start, 0)
-            izstop = min(iz.stop or nz + 1, nz + self.overlaps[2])
+            izstart = max(iz.start or -self.nghosts, -self.nghosts)
+            izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts)
         else:
             izstart = iz
             izstop = iz + 1
@@ -99,11 +121,11 @@ class MultiFABWrapper(object):
 
             # --- The ix1, 2 etc are relative to global indexing
             ix1 = max(ixstart, lovects[0,i])
-            ix2 = min((ixstop or nx+1), lovects[0,i] + fields[i].shape[0])
+            ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0])
             iy1 = max(iystart, lovects[1,i])
-            iy2 = min((iystop or ny+1), lovects[1,i] + fields[i].shape[1])
+            iy2 = min(iystop, lovects[1,i] + fields[i].shape[1])
             iz1 = max(izstart, lovects[2,i])
-            iz2 = min((izstop or nz+1), lovects[2,i] + fields[i].shape[2])
+            iz2 = min(izstop, lovects[2,i] + fields[i].shape[2])
 
             if ix1 < ix2 and iy1 < iy2 and iz1 < iz2:
 
@@ -126,7 +148,84 @@ class MultiFABWrapper(object):
         if not isinstance(iz, slice):
             sss[2] = 0
 
-        return resultglobal[sss]
+        return resultglobal[tuple(sss)]
+
+    def _getitem2d(self, index):
+        """Returns slices of a 2D decomposed array,
+        """
+
+        lovects = self._getlovects()
+        hivects = self._gethivects()
+        fields = self._getfields()
+
+        ix = index[0]
+        iz = index[1]
+
+        if len(fields[0].shape) > self.dim:
+            ncomps = fields[0].shape[-1]
+        else:
+            ncomps = 1
+
+        if len(index) > self.dim:
+            if ncomps > 1:
+                ic = index[2]
+            else:
+                raise Exception('Too many indices given')
+        else:
+            ic = None
+
+        nx = hivects[0,:].max() - self.nghosts
+        nz = hivects[1,:].max() - self.nghosts
+
+        if isinstance(ix, slice):
+            ixstart = max(ix.start or -self.nghosts, -self.nghosts)
+            ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts)
+        else:
+            ixstart = ix
+            ixstop = ix + 1
+        if isinstance(iz, slice):
+            izstart = max(iz.start or -self.nghosts, -self.nghosts)
+            izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[1] + self.nghosts)
+        else:
+            izstart = iz
+            izstop = iz + 1
+
+        # --- Setup the size of the array to be returned and create it.
+        # --- Space is added for multiple components if needed.
+        sss = (max(0, ixstop - ixstart),
+               max(0, izstop - izstart))
+        if ncomps > 1 and ic is None:
+            sss = tuple(list(sss) + [ncomps])
+        resultglobal = np.zeros(sss)
+
+        for i in range(len(fields)):
+
+            # --- The ix1, 2 etc are relative to global indexing
+            ix1 = max(ixstart, lovects[0,i])
+            ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0])
+            iz1 = max(izstart, lovects[1,i])
+            iz2 = min(izstop, lovects[1,i] + fields[i].shape[1])
+
+            if ix1 < ix2 and iz1 < iz2:
+
+                sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]),
+                       slice(iz1 - lovects[1,i], iz2 - lovects[1,i]))
+                if ic is not None:
+                    sss = tuple(list(sss) + [ic])
+
+                vslice = (slice(ix1 - ixstart, ix2 - ixstart),
+                          slice(iz1 - izstart, iz2 - izstart))
+
+                resultglobal[vslice] = fields[i][sss]
+
+        # --- Now remove any of the reduced dimensions.
+        sss = [slice(None), slice(None)]
+        if not isinstance(ix, slice):
+            sss[0] = 0
+        if not isinstance(iz, slice):
+            sss[1] = 0
+
+        return resultglobal[tuple(sss)]
 
     def __setitem__(self, index, value):
         """Sets slices of a decomposed array. The shape of
@@ -135,8 +234,23 @@ class MultiFABWrapper(object):
         - value: input array (must be supplied)
         """
         if index == Ellipsis:
-            index = (slice(None), slice(None), slice(None))
-
+            index = tuple(self.dim*[slice(None)])
+
+        if len(index) < self.dim:
+            # --- Add extra dims to index if needed
+            index = list(index)
+            for i in range(len(index), self.dim):
+                index.append(slice(None))
+            index = tuple(index)
+
+        if self.dim == 2:
+            return self._setitem2d(index, value)
+        elif self.dim == 3:
+            return self._setitem3d(index, value)
+
+    def _setitem3d(self, index, value):
+        """Sets slices of a decomposed 3D array.
+        """
         ix = index[0]
         iy = index[1]
         iz = index[2]
@@ -145,9 +259,9 @@ class MultiFABWrapper(object):
         hivects = self._gethivects()
         fields = self._getfields()
 
-        nx = hivects[0,:].max()
-        ny = hivects[1,:].max()
-        nz = hivects[2,:].max()
+        nx = hivects[0,:].max() - self.nghosts
+        ny = hivects[1,:].max() - self.nghosts
+        nz = hivects[2,:].max() - self.nghosts
 
         # --- Add extra dimensions so that the input has the same number of
         # --- dimensions as array.
@@ -160,20 +274,20 @@ class MultiFABWrapper(object):
             value3d.shape = sss
 
         if isinstance(ix, slice):
-            ixstart = max(ix.start, 0)
-            ixstop = min(ix.stop or nx + 1, nx + self.overlaps[0])
+            ixstart = max(ix.start or -self.nghosts, -self.nghosts)
+            ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts)
         else:
             ixstart = ix
             ixstop = ix + 1
         if isinstance(iy, slice):
-            iystart = max(iy.start, 0)
-            iystop = min(iy.stop or ny + 1, ny + self.overlaps[1])
+            iystart = max(iy.start or -self.nghosts, -self.nghosts)
+            iystop = min(iy.stop or ny + 1 + self.nghosts, ny + self.overlaps[1] + self.nghosts)
         else:
             iystart = iy
             iystop = iy + 1
         if isinstance(iz, slice):
-            izstart = max(iz.start, 0)
-            izstop = min(iz.stop or nz + 1, nz + self.overlaps[2])
+            izstart = max(iz.start or -self.nghosts, -self.nghosts)
+            izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts)
         else:
             izstart = iz
             izstop = iz + 1
@@ -182,11 +296,11 @@ class MultiFABWrapper(object):
 
             # --- The ix1, 2 etc are relative to global indexing
             ix1 = max(ixstart, lovects[0,i])
-            ix2 = min((ixstop or nx+1), lovects[0,i] + fields[i].shape[0])
+            ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0])
             iy1 = max(iystart, lovects[1,i])
-            iy2 = min((iystop or ny+1), lovects[1,i] + fields[i].shape[1])
+            iy2 = min(iystop, lovects[1,i] + fields[i].shape[1])
             iz1 = max(izstart, lovects[2,i])
-            iz2 = min((izstop or nz+1), lovects[2,i] + fields[i].shape[2])
+            iz2 = min(izstop, lovects[2,i] + fields[i].shape[2])
 
             if ix1 < ix2 and iy1 < iy2 and iz1 < iz2:
 
@@ -202,49 +316,127 @@ class MultiFABWrapper(object):
                 else:
                     fields[i][sss] = value
 
+    def _setitem2d(self, index, value):
+        """Sets slices of a decomposed 2D array.
+        """
+        ix = index[0]
+        iz = index[2]
 
-def ExWrapper(level=0):
-    return MultiFABWrapper(direction=0, overlaps=[0,1,1],
-                           get_lovects=_libwarpx.get_mesh_electric_field_lovects,
-                           get_fabs=_libwarpx.get_mesh_electric_field, level=level)
+        lovects = self._getlovects()
+        hivects = self._gethivects()
+        fields = self._getfields()
 
-def EyWrapper(level=0):
-    return MultiFABWrapper(direction=1, overlaps=[1,0,1],
-                           get_lovects=_libwarpx.get_mesh_electric_field_lovects,
-                           get_fabs=_libwarpx.get_mesh_electric_field, level=level)
+        if len(fields[0].shape) > self.dim:
+            ncomps = fields[0].shape[-1]
+        else:
+            ncomps = 1
 
-def EzWrapper(level=0):
-    return MultiFABWrapper(direction=2, overlaps=[1,1,0],
-                           get_lovects=_libwarpx.get_mesh_electric_field_lovects,
-                           get_fabs=_libwarpx.get_mesh_electric_field, level=level)
+        if len(index) > self.dim:
+            if ncomps > 1:
+                ic = index[2]
+            else:
+                raise Exception('Too many indices given')
+        else:
+            ic = None
 
-def BxWrapper(level=0):
-    return MultiFABWrapper(direction=0, overlaps=[1,0,0],
-                           get_lovects=_libwarpx.get_mesh_magnetic_field_lovects,
-                           get_fabs=_libwarpx.get_mesh_magnetic_field, level=level)
+        nx = hivects[0,:].max() - self.nghosts
+        nz = hivects[2,:].max() - self.nghosts
 
-def ByWrapper(level=0):
-    return MultiFABWrapper(direction=1, overlaps=[0,1,0],
-                           get_lovects=_libwarpx.get_mesh_magnetic_field_lovects,
-                           get_fabs=_libwarpx.get_mesh_magnetic_field, level=level)
+        # --- Add extra dimensions so that the input has the same number of
+        # --- dimensions as array.
+        if isinstance(value, np.ndarray):
+            value3d = np.array(value, copy=False)
+            sss = list(value3d.shape)
+            if not isinstance(ix, slice): sss[0:0] = [1]
+            if not isinstance(iz, slice): sss[1:1] = [1]
+            value3d.shape = sss
 
-def BzWrapper(level=0):
-    return MultiFABWrapper(direction=2, overlaps=[0,0,1],
-                           get_lovects=_libwarpx.get_mesh_magnetic_field_lovects,
-                           get_fabs=_libwarpx.get_mesh_magnetic_field, level=level)
+        if isinstance(ix, slice):
+            ixstart = max(ix.start or -self.nghosts, -self.nghosts)
+            ixstop = min(ix.stop or nx + 1 + self.nghosts, nx + self.overlaps[0] + self.nghosts)
+        else:
+            ixstart = ix
+            ixstop = ix + 1
+        if isinstance(iz, slice):
+            izstart = max(iz.start or -self.nghosts, -self.nghosts)
+            izstop = min(iz.stop or nz + 1 + self.nghosts, nz + self.overlaps[2] + self.nghosts)
+        else:
+            izstart = iz
+            izstop = iz + 1
+
+        for i in range(len(fields)):
+
+            # --- The ix1, 2 etc are relative to global indexing
+            ix1 = max(ixstart, lovects[0,i])
+            ix2 = min(ixstop, lovects[0,i] + fields[i].shape[0])
+            iz1 = max(izstart, lovects[2,i])
+            iz2 = min(izstop, lovects[2,i] + fields[i].shape[2])
+
+            if ix1 < ix2 and iz1 < iz2:
 
-def JxWrapper(level=0):
-    return MultiFABWrapper(direction=0, overlaps=[0,1,1],
-                           get_lovects=_libwarpx.get_mesh_current_density_lovects,
-                           get_fabs=_libwarpx.get_mesh_current_density, level=level)
+                sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]),
+                       slice(iz1 - lovects[2,i], iz2 - lovects[2,i]))
+                if ic is not None:
+                    sss = tuple(list(sss) + [ic])
 
-def JyWrapper(level=0):
-    return MultiFABWrapper(direction=1, overlaps=[1,0,1],
-                           get_lovects=_libwarpx.get_mesh_current_density_lovects,
-                           get_fabs=_libwarpx.get_mesh_current_density, level=level)
+                if isinstance(value, np.ndarray):
+                    vslice = (slice(ix1 - ixstart, ix2 - ixstart),
+                              slice(iz1 - izstart, iz2 - izstart))
+                    fields[i][sss] = value3d[vslice]
+                else:
+                    fields[i][sss] = value
 
-def JzWrapper(level=0):
-    return MultiFABWrapper(direction=2, overlaps=[1,1,0],
-                           get_lovects=_libwarpx.get_mesh_current_density_lovects,
-                           get_fabs=_libwarpx.get_mesh_current_density, level=level)
 
+def ExWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=0, overlaps=[0,1,1],
+                            get_lovects=_libwarpx.get_mesh_electric_field_lovects,
+                            get_fabs=_libwarpx.get_mesh_electric_field,
+                            level=level, include_ghosts=include_ghosts)
+
+def EyWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=1, overlaps=[1,0,1],
+                            get_lovects=_libwarpx.get_mesh_electric_field_lovects,
+                            get_fabs=_libwarpx.get_mesh_electric_field,
+                            level=level, include_ghosts=include_ghosts)
+
+def EzWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=2, overlaps=[1,1,0],
+                            get_lovects=_libwarpx.get_mesh_electric_field_lovects,
+                            get_fabs=_libwarpx.get_mesh_electric_field,
+                            level=level, include_ghosts=include_ghosts)
+
+def BxWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=0, overlaps=[1,0,0],
+                            get_lovects=_libwarpx.get_mesh_magnetic_field_lovects,
+                            get_fabs=_libwarpx.get_mesh_magnetic_field,
+                            level=level, include_ghosts=include_ghosts)
+
+def ByWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=1, overlaps=[0,1,0],
+                            get_lovects=_libwarpx.get_mesh_magnetic_field_lovects,
+                            get_fabs=_libwarpx.get_mesh_magnetic_field,
+                            level=level, include_ghosts=include_ghosts)
+
+def BzWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=2, overlaps=[0,0,1],
+                            get_lovects=_libwarpx.get_mesh_magnetic_field_lovects,
+                            get_fabs=_libwarpx.get_mesh_magnetic_field,
+                            level=level, include_ghosts=include_ghosts)
+
+def JxWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=0, overlaps=[0,1,1],
+                            get_lovects=_libwarpx.get_mesh_current_density_lovects,
+                            get_fabs=_libwarpx.get_mesh_current_density,
+                            level=level, include_ghosts=include_ghosts)
+
+def JyWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=1, overlaps=[1,0,1],
+                            get_lovects=_libwarpx.get_mesh_current_density_lovects,
+                            get_fabs=_libwarpx.get_mesh_current_density,
+                            level=level, include_ghosts=include_ghosts)
+
+def JzWrapper(level=0, include_ghosts=False):
+    return _MultiFABWrapper(direction=2, overlaps=[1,1,0],
+                            get_lovects=_libwarpx.get_mesh_current_density_lovects,
+                            get_fabs=_libwarpx.get_mesh_current_density,
+                            level=level, include_ghosts=include_ghosts)