Skip to content

Commit

Permalink
Fix handling of ghost cells in fields
Browse files Browse the repository at this point in the history
  • Loading branch information
@dpgrote
dpgrote committed Jul 21, 2023
1 parent a999ad4 commit d1a143d
Showing 1 changed file with 135 additions and 60 deletions.
195 changes: 135 additions & 60 deletions Python/pywarpx/fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,9 @@ class _MultiFABWrapper(object):
The refinement level
include_ghosts: bool, default=False
Whether to include the ghost cells
Whether to include the ghost cells.
Note that when True, the first n-ghost negative indices will refer to the lower
ghost cells.
"""
def __init__(self, mf_name, level, include_ghosts=False):
self.mf_name = mf_name
Expand Down Expand Up @@ -103,6 +105,9 @@ def shape(self):
"""
min_box = self.mf.box_array().minimal_box()
shape = list(min_box.size - min_box.small_end)
if self.include_ghosts:
nghosts = self.mf.n_grow_vect()
shape = [shape[i] + 2*nghosts[i] for i in range(self.dim)]
shape.append(self.mf.nComp)
return tuple(shape)

Expand Down Expand Up @@ -138,6 +143,19 @@ def mesh(self, direction):
ilo = min_box.small_end[idir]
ihi = min_box.big_end[idir]

if self.include_ghosts:
# The ghost cells are added to the upper and lower end of the global domain.
nghosts = self.mf.n_grow_vect()
ilo = list(ilo)
ihi = list(ihi)
min_box = self.mf.box_array().minimal_box()
imax = min_box.big_end
for i in range(self.dim):
if ilo[i] == 0:
ilo[i] -= nghosts[i]
if ihi[i] == imax[i]:
ihi[i] += nghosts[i]

# Cell size in the direction
warpx = libwarpx.libwarpx_so.get_instance()
dd = warpx.Geom(self.level).data().CellSize(idir)
Expand All @@ -154,22 +172,83 @@ def mesh(self, direction):
lo = warpx.Geom(self.level).ProbLo(idir)
return lo + np.arange(ilo,ihi+1)*dd + shift

def _get_indices(self, index, missing):
"""Expand the index list to length three.
Parameters
----------
index: sequence of length dims
The indices for each dim
missing:
The value used to fill in the extra dimensions added
"""
if self.dim == 1:
return index[0], missing, missing
elif self.dim == 2:
return index[0], index[1], missing
elif self.dim == 3:
return index[0], index[1], index[2]

def _get_min_indices(self):
"""Returns the minimum indices, expanded to length 3"""
min_box = self.mf.box_array().minimal_box()
imin = self._get_indices(min_box.small_end, 0)
if self.include_ghosts:
nghosts = self._get_n_ghosts()
imin = [imin[i] - nghosts[i] for i in range(3)]
return imin

def _get_max_indices(self):
"""Returns the maximum indices, expanded to length 3.
"""
min_box = self.mf.box_array().minimal_box()
imax = self._get_indices(min_box.big_end, 0)
if self.include_ghosts:
nghosts = self._get_n_ghosts()
imax = [imax[i] + nghosts[i] for i in range(3)]
return imax

def _get_n_ghosts(self):
nghosts = list(self._get_indices(self.mf.n_grow_vect(), 0))
# The components always has nghosts = 0
nghosts.append(0)
return nghosts

def _fix_index(self, ii, imax, d):
"""Handle negative index, wrapping them as needed.
This depends on whether ghost cells are included. When true, the indices are
shifted by the number of ghost cells before being wrapped.
"""
nghosts = self._get_n_ghosts()
if self.include_ghosts:
ii += nghosts[d]
if ii < 0:
ii += imax
if self.include_ghosts:
ii -= nghosts[d]
return ii

def _find_start_stop(self, ii, imin, imax, d):
"""Given the input index, calculate the start and stop range of the indices.
Parameters
----------
ii: None, slice, integer
Input index, either None, a slice object, or an integer
Input index, either None, a slice object, or an integer.
Note that ii can be negative.
imin: integer
The global lowest index value in the specified direction
The global lowest lower bound in the specified direction.
This can include the ghost cells.
imax: integer
The global highest index value in the specified direction
The global highest upper bound in the specified direction.
This can include the ghost cells.
This should be the max index + 1.
d: integer
The dimension number, 1, 2, 3, or 4 (4 being the components)
The dimension number, 0, 1, 2, or 3 (3 being the components)
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.
Expand All @@ -182,70 +261,41 @@ def _find_start_stop(self, ii, imin, imax, d):
if ii.start is None:
iistart = imin
else:
iistart = ii.start
iistart = self._fix_index(ii.start, imax, d)
if ii.stop is None:
iistop = imax
else:
iistop = ii.stop
iistop = self._fix_index(ii.stop, imax, d)
else:
ii = self._fix_index(ii, imax, d)
iistart = ii
iistop = ii + 1
assert imin <= iistart <= imax, Exception(f'Dimension {d} lower index is out of bounds')
assert imin <= iistop <= imax, Exception(f'Dimension {d} upper index is out of bounds')
assert imin <= iistart <= imax, Exception(f'Dimension {d+1} lower index is out of bounds')
assert imin <= iistop <= imax, Exception(f'Dimension {d+1} upper index is out of bounds')
return iistart, iistop

def _get_indices(self, index, missing):
"""Expand the index list to length three.
Parameters
----------
index: sequence of length dims
The indices for each dim
missing:
The value used to fill in the extra dimensions added
"""
if self.dim == 1:
return index[0], missing, missing
elif self.dim == 2:
return index[0], index[1], missing
elif self.dim == 3:
return index[0], index[1], index[2]

def _get_min_indices(self):
"""Returns the minimum indices, expanded to length 3"""
min_box = self.mf.box_array().minimal_box()
return self._get_indices(min_box.small_end, 0)

def _get_max_indices(self):
"""Returns the maximum indices, expanded to length 3.
This is the index appropriate for the upper bound of a slice
and so includes the +1.
"""
min_box = self.mf.box_array().minimal_box()
indices = self._get_indices(min_box.big_end, 0)
indicesp1 = [i + 1 for i in indices]
return indicesp1

def _get_intersect_slice(self, box, starts, stops, ic):
"""Return the slices where the block intersects with the global slice.
If the block does not intersect, return None.
This also shifts the block slices by the number of guard cells in the
MultiFab arrays.
This also shifts the block slices by the number of ghost cells in the
MultiFab arrays since the arrays include the ghost cells.
Parameters
----------
box: Box instance
The box defining the block
starts: sequence
The minimum indices of the global slice
The minimum indices of the global slice.
These can be negative.
stops: sequence
The maximum indices of the global slice
The maximum indices of the global slice.
These can be negative.
ic: integer
The indices of the components
This can be negative.
Returns
-------
Expand All @@ -256,18 +306,33 @@ def _get_intersect_slice(self, box, starts, stops, ic):
The slice of the intersection relative to the global array where the data from individual block will go
"""
ilo = self._get_indices(box.small_end, 0)
ihi = self._get_indices(box.big_end, 0)
ihip1 = [i + 1 for i in ihi]
i1 = np.maximum(starts, ilo)
i2 = np.minimum(stops, ihip1)
nguard = self._get_indices(self.mf.n_grow_vect(), 0)
ilo_g = self._get_indices(box.small_end, 0)
ihi_g = self._get_indices(box.big_end, 0)
nghosts = self._get_indices(self.mf.n_grow_vect(), 0)
if self.include_ghosts:
# The starts and stops could include the ghost cells, so add the ghost cells
# to the lo and hi that they are compared against (to get i1 and i2 below).
# The ghost cells are only added to the lower and upper end of the global domain.
ilo_g = list(ilo_g)
ihi_g = list(ihi_g)
min_box = self.mf.box_array().minimal_box()
imax = self._get_indices(min_box.big_end, 0)
for i in range(3):
if ilo_g[i] == 0:
ilo_g[i] -= nghosts[i]
if ihi_g[i] == imax[i]:
ihi_g[i] += nghosts[i]
# Add 1 to the upper end to be consistent with the slicing notation
ihi_g_p1 = [i + 1 for i in ihi_g]
i1 = np.maximum(starts, ilo_g)
i2 = np.minimum(stops, ihi_g_p1)

if np.all(i1 < i2):

block_slices = []
global_slices = []
for i in range(3):
block_slices.append(slice(i1[i] - ilo[i] + nguard[i], i2[i] - ilo[i] + nguard[i]))
block_slices.append(slice(i1[i] - ilo[i] + nghosts[i], i2[i] - ilo[i] + nghosts[i]))
global_slices.append(slice(i1[i] - starts[i], i2[i] - starts[i]))

if ic is None:
Expand All @@ -286,12 +351,16 @@ def __getitem__(self, index):
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, and in fortran ordering, i.e. [ix,iy,iz].
This allows negative indexing, though with ghosts cells included, the first n-ghost negative
indices will refer to the lower guard cells.
Parameters
----------
index: integer, or sequence of integers or slices, or Ellipsis
Index of the slice to return
"""
# Note that the index can have negative values (which wrap around) and has 1 added to the upper
# limit using python style slicing
if index == Ellipsis:
index = self.dim*[slice(None)]
elif isinstance(index, slice):
Expand All @@ -311,14 +380,15 @@ def __getitem__(self, index):
ii = self._get_indices(index, None)
ic = index[-1]

# Global extent. These include the ghost cells when include_ghosts is True
ixmin, iymin, izmin = self._get_min_indices()
ixmax, iymax, izmax = self._get_max_indices()

# Setup the size of the array to be returned
ixstart, ixstop = self._find_start_stop(ii[0], ixmin, ixmax, 1)
iystart, iystop = self._find_start_stop(ii[1], iymin, iymax, 2)
izstart, izstop = self._find_start_stop(ii[2], izmin, izmax, 3)
icstart, icstop = self._find_start_stop(ic, 0, self.mf.n_comp(), 4)
ixstart, ixstop = self._find_start_stop(ii[0], ixmin, ixmax+1, 0)
iystart, iystop = self._find_start_stop(ii[1], iymin, iymax+1, 1)
izstart, izstop = self._find_start_stop(ii[2], izmin, izmax+1, 2)
icstart, icstop = self._find_start_stop(ic, 0, self.mf.n_comp(), 3)

# Gather the data to be included in a list to be sent to other processes
starts = [ixstart, iystart, izstart]
Expand Down Expand Up @@ -368,6 +438,8 @@ 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.
This allows negative indexing, though with ghosts cells included, the first n-ghost negative
indices will refer to the lower guard cells.
Parameters
----------
Expand All @@ -377,6 +449,8 @@ def __setitem__(self, index, value):
value: scalar or array
Input value to assign to the specified slice of the MultiFab
"""
# Note that the index can have negative values (which wrap around) and has 1 added to the upper
# limit using python style slicing
if index == Ellipsis:
index = tuple(self.dim*[slice(None)])
elif isinstance(index, slice):
Expand All @@ -396,14 +470,15 @@ def __setitem__(self, index, value):
ii = self._get_indices(index, None)
ic = index[-1]

# Global extent. These include the ghost cells when include_ghosts is True
ixmin, iymin, izmin = self._get_min_indices()
ixmax, iymax, izmax = self._get_max_indices()

# Setup the size of the global array to be set
ixstart, ixstop = self._find_start_stop(ii[0], ixmin, ixmax, 1)
iystart, iystop = self._find_start_stop(ii[1], iymin, iymax, 2)
izstart, izstop = self._find_start_stop(ii[2], izmin, izmax, 3)
icstart, icstop = self._find_start_stop(ic, 0, self.mf.n_comp(), 4)
ixstart, ixstop = self._find_start_stop(ii[0], ixmin, ixmax+1, 0)
iystart, iystop = self._find_start_stop(ii[1], iymin, iymax+1, 1)
izstart, izstop = self._find_start_stop(ii[2], izmin, izmax+1, 2)
icstart, icstop = self._find_start_stop(ic, 0, self.mf.n_comp(), 3)

if isinstance(value, np.ndarray):
# Expand the shape of the input array to match the shape of the global array
Expand Down

0 comments on commit d1a143d

Please sign in to comment.