aboutsummaryrefslogtreecommitdiff
path: root/Python/pywarpx/fields.py
blob: 8237ed867b9b79158dd3ee384322d1464d6beb43 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
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)