Added wrappers for multifab fields, e.g. ExWrapper

1 files changed, 250 insertions, 0 deletions

diff --git a/Python/pywarpx/fields.py b/Python/pywarpx/fields.py
new file mode 100644
index 000000000..8237ed867
--- /dev/null
+++ b/Python/pywarpx/fields.py
@@ -0,0 +1,250 @@
+"""Provides wrappers around field and current density on multiFABs
+
+Available routines:
+
+ExWrapper, EyWrapper, EzWrapper
+BxWrapper, ByWrapper, BzWrapper
+JxWrapper, JyWrapper, JzWrapper
+
+"""
+import numpy as np
+from . import _libwarpx
+
+
+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.
+     - direction: component to access, one of the values (0, 1, 2)
+     - overlaps: is one along the axes where the grid boundaries overlap the neighboring grid
+     - get_lovects: routine that returns the list of lo vectors
+     - get_fabs: routine that returns the list of FABs
+     - level=0: ignored
+    """
+    def __init__(self, direction, overlaps, get_lovects, get_fabs, level=0):
+        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
+
+    def _getlovects(self):
+        return self.get_lovects(self.level, self.direction, self.include_ghosts)
+
+    def _gethivects(self):
+        lovects = self._getlovects()
+        fields = self._getfields()
+
+        hivects = np.zeros_like(lovects)
+        for i in range(len(fields)):
+            hivects[:,i] = lovects[:,i] + np.array(fields[i].shape) - self.overlaps
+
+        return hivects
+
+    def _getfields(self):
+        return self.get_fabs(self.level, self.direction, self.include_ghosts)
+
+    def __len__(self):
+        return lend(self._getlovects())
+
+    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.
+        """
+        if index == Ellipsis:
+            index = (slice(None), slice(None), slice(None))
+
+        ix = index[0]
+        iy = index[1]
+        iz = index[2]
+
+        lovects = self._getlovects()
+        hivects = self._gethivects()
+        fields = self._getfields()
+
+        nx = hivects[0,:].max()
+        ny = hivects[1,:].max()
+        nz = hivects[2,:].max()
+
+        if isinstance(ix, slice):
+            ixstart = max(ix.start, 0)
+            ixstop = min(ix.stop or nx + 1, nx + self.overlaps[0])
+        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])
+        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])
+        else:
+            izstart = iz
+            izstop = iz + 1
+
+        # --- Setup the size of the array to be returned and create it.
+        sss = (max(0, ixstop - ixstart),
+               max(0, iystop - iystart),
+               max(0, izstop - izstart))
+        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 or nx+1), 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])
+            iz1 = max(izstart, lovects[2,i])
+            iz2 = min((izstop or nz+1), lovects[2,i] + fields[i].shape[2])
+
+            if ix1 < ix2 and iy1 < iy2 and iz1 < iz2:
+
+                sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]),
+                       slice(iy1 - lovects[1,i], iy2 - lovects[1,i]),
+                       slice(iz1 - lovects[2,i], iz2 - lovects[2,i]))
+
+                vslice = (slice(ix1 - ixstart, ix2 - ixstart),
+                          slice(iy1 - iystart, iy2 - iystart),
+                          slice(iz1 - izstart, iz2 - izstart))
+
+                resultglobal[vslice] = fields[i][sss]
+
+        # --- Now remove any of the reduced dimensions.
+        sss = [slice(None), slice(None), slice(None)]
+        if not isinstance(ix, slice):
+            sss[0] = 0
+        if not isinstance(iy, slice):
+            sss[1] = 0
+        if not isinstance(iz, slice):
+            sss[2] = 0
+
+        return resultglobal[sss]
+
+    def __setitem__(self, index, value):
+        """Sets slices of a decomposed array. The shape of
+      the input object depends on the number of arguments specified, which can
+      be from none to all three.
+        - value: input array (must be supplied)
+        """
+        if index == Ellipsis:
+            index = (slice(None), slice(None), slice(None))
+
+        ix = index[0]
+        iy = index[1]
+        iz = index[2]
+
+        lovects = self._getlovects()
+        hivects = self._gethivects()
+        fields = self._getfields()
+
+        nx = hivects[0,:].max()
+        ny = hivects[1,:].max()
+        nz = hivects[2,:].max()
+
+        # --- 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(iy, slice): sss[1:1] = [1]
+            if not isinstance(iz, slice): sss[2:2] = [1]
+            value3d.shape = sss
+
+        if isinstance(ix, slice):
+            ixstart = max(ix.start, 0)
+            ixstop = min(ix.stop or nx + 1, nx + self.overlaps[0])
+        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])
+        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])
+        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 or nx+1), 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])
+            iz1 = max(izstart, lovects[2,i])
+            iz2 = min((izstop or nz+1), lovects[2,i] + fields[i].shape[2])
+
+            if ix1 < ix2 and iy1 < iy2 and iz1 < iz2:
+
+                sss = (slice(ix1 - lovects[0,i], ix2 - lovects[0,i]),
+                       slice(iy1 - lovects[1,i], iy2 - lovects[1,i]),
+                       slice(iz1 - lovects[2,i], iz2 - lovects[2,i]))
+
+                if isinstance(value, np.ndarray):
+                    vslice = (slice(ix1 - ixstart, ix2 - ixstart),
+                              slice(iy1 - iystart, iy2 - iystart),
+                              slice(iz1 - izstart, iz2 - izstart))
+                    fields[i][sss] = value3d[vslice]
+                else:
+                    fields[i][sss] = value
+
+
+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)
+
+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)
+
+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)
+
+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)
+
+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)
+
+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)
+
+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)
+
+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)
+
+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)
+