From 47f57581a8fc952c2b14da6484ba2defe980af09 Mon Sep 17 00:00:00 2001 From: David Grote Date: Wed, 23 Oct 2024 13:43:28 -0700 Subject: [PATCH] Fix getitem for cases with incomplete domains (#385) This is a fix for the global indexing mechanism in `MultiFab` that is needed in cases where the box array does not fully cover the domain. In those cases, the `__getitem__` is done twice, once to get all of the data including places where the ghost cells internal to the domain cover places that are not covered by valid cells, and a second time to ensure that in places where the ghost cells do overlap valid cells, that the data used is from the valid cells. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- src/Base/BoxArray.cpp | 10 +++++--- src/amrex/extensions/MultiFab.py | 42 ++++++++++++++++++++++++++------ 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/src/Base/BoxArray.cpp b/src/Base/BoxArray.cpp index f79e1b3a..a6b57d87 100644 --- a/src/Base/BoxArray.cpp +++ b/src/Base/BoxArray.cpp @@ -129,13 +129,15 @@ void init_BoxArray(py::module &m) { //! \brief Apply surroundingNodes(Box,int) to each Box in //! BoxArray. See the documentation of Box for details. BoxArray& surroundingNodes (int dir); +*/ //! Apply Box::enclosedCells() to each Box in the BoxArray. - BoxArray& enclosedCells (); - - //! Apply Box::enclosedCells(int) to each Box in the BoxArray. - BoxArray& enclosedCells (int dir); + .def("enclosed_cells", + py::overload_cast<>(&BoxArray::enclosedCells)) + .def("enclosed_cells", + py::overload_cast(&BoxArray::enclosedCells)) +/* //! Apply Box::convert(IndexType) to each Box in the BoxArray. BoxArray& convert (IndexType typ); diff --git a/src/amrex/extensions/MultiFab.py b/src/amrex/extensions/MultiFab.py index e91f9b02..3f6b4144 100644 --- a/src/amrex/extensions/MultiFab.py +++ b/src/amrex/extensions/MultiFab.py @@ -463,7 +463,7 @@ def _get_intersect_slice(self, mfi, index, with_internal_ghosts): return None, None -def __getitem__(self, index): +def __getitem__(self, index, with_internal_ghosts=False): """Returns slice of the MultiFab using global indexing, as a numpy array. This uses numpy array indexing, with the indexing relative to the global array. The slice ranges can cross multiple blocks and the result will be gathered into a single @@ -484,13 +484,18 @@ def __getitem__(self, index): ---------- index : the index using numpy style indexing Index of the slice to return. + with_internal_ghosts : bool, optional + Whether to include internal ghost cells. When true, data from ghost cells may be used that + overlaps valid cells. """ - index = _process_index(self, index) + index4 = _process_index(self, index) # Gather the data to be included in a list to be sent to other processes datalist = [] for mfi in self: - block_slices, global_slices = _get_intersect_slice(self, mfi, index, False) + block_slices, global_slices = _get_intersect_slice( + self, mfi, index4, with_internal_ghosts + ) if global_slices is not None: # Note that the array will always have 4 dimensions. device_arr = _get_field(self, mfi) @@ -522,11 +527,29 @@ def __getitem__(self, index): raise Exception("MultiFab.__getitem__ requires mpi4py") all_datalist = comm_world.allgather(datalist) - # Create the array to be returned - result_shape = tuple([max(0, ii.stop - ii.start) for ii in index]) + # The shape of the array to be returned + result_shape = tuple([max(0, ii.stop - ii.start) for ii in index4]) + + # If the boxes do not fill the domain, then include the internal ghost + # cells since they may cover internal regions not otherwise covered by valid cells. + # If the domain is not completely covered, __getitem__ is done twice, the first time + # including internal ghost cells, and the second time without. The second time is needed + # to ensure that in places where ghost cells overlap with valid cells, that the data + # from the valid cells is used. + # The domain is complete if the number of cells in the box array is the same as + # the number of cells in the minimal box. + cell_centered_box_array = self.box_array().enclosed_cells() + domain_complete = ( + cell_centered_box_array.numPts == cell_centered_box_array.minimal_box().numPts() + ) + + if domain_complete or with_internal_ghosts: + result_global = None + else: + # First get the data including the internal ghost cells + result_global = self.__getitem__(index, with_internal_ghosts=True) # Now, copy the data into the result array - result_global = None for datalist in all_datalist: for global_slices, f_arr in datalist: if result_global is None: @@ -540,7 +563,12 @@ def __getitem__(self, index): # Remove dimensions of length 1, and if all dimensions # are removed, return a scalar (that's what the [()] does) - return result_global.squeeze()[()] + if with_internal_ghosts: + # Return the data without the squeeze so that the result can be used in the loop + # above again. + return result_global + else: + return result_global.squeeze()[()] def __setitem__(self, index, value):