diff --git a/yt/data_objects/tests/test_octree.py b/yt/data_objects/tests/test_octree.py index 0c52ed337e..ae113562c9 100644 --- a/yt/data_objects/tests/test_octree.py +++ b/yt/data_objects/tests/test_octree.py @@ -1,8 +1,11 @@ import numpy as np from numpy.testing import assert_almost_equal, assert_equal -from yt.geometry.oct_container import OctreeContainer +from yt.geometry.oct_container import _ORDER_MAX, OctreeContainer +from yt.geometry.particle_oct_container import ParticleOctreeContainer +from yt.geometry.selection_routines import AlwaysSelector from yt.testing import fake_sph_grid_ds +from yt.utilities.lib.geometry_utils import get_morton_indices n_ref = 4 @@ -142,3 +145,40 @@ def test_num_zones_tuple(): assert oct_scalar is not None assert oct_tuple is not None assert oct_nonuniform is not None + + +def test_octcellindex_neighbours_num_zones(): + """ + Regression test for #5402: fill_octcellindex_neighbours hard-coded the + assumption that every oct holds 2x2x2 zones, so both the loop bounds and + the output buffer size were wrong whenever num_zones wasn't (2, 2, 2). + """ + DLE = np.array([0.0, 0.0, 0.0]) + DRE = np.array([8.0, 8.0, 8.0]) + dx = (DRE - DLE) / (2**_ORDER_MAX) + + # One particle per octant of the root oct: refines exactly one level deep. + centers = np.array( + [[x, y, z] for x in (2, 6) for y in (2, 6) for z in (2, 6)], dtype="float64" + ) + morton = get_morton_indices(np.floor((centers - DLE) / dx).astype("uint64")) + morton.sort() + + for nz in ((2, 2, 2), (2, 3, 4)): + octree = ParticleOctreeContainer((1, 1, 1), DLE, DRE, num_zones=nz) + octree.n_ref = 1 + octree.add(morton) + octree.finalize() + + nzones = nz[0] * nz[1] * nz[2] + n_per_oct = (nz[0] + 2) * (nz[1] + 2) * (nz[2] + 2) # +1 ghost zone each side + + selector = AlwaysSelector(None) + num_octs = selector.count_octs(octree, -1) + _, cell_inds = octree.fill_octcellindex_neighbours(selector) + + # old oct_visitors/containers hardcoded 4**3=64 cells/oct; for nz=(2,3,4) it's really + # 4*5*6=120. this assertion catches that + assert_equal(cell_inds.size, num_octs * n_per_oct) + assert cell_inds.min() >= 0 + assert cell_inds.max() <= nzones diff --git a/yt/geometry/oct_container.pyx b/yt/geometry/oct_container.pyx index 981085b15a..74063ab4f5 100644 --- a/yt/geometry/oct_container.pyx +++ b/yt/geometry/oct_container.pyx @@ -818,12 +818,19 @@ cdef class OctreeContainer: num_octs = selector.count_octs(self, domain_id) cdef NeighbourCellIndexVisitor visitor + cdef int n_per_oct cdef np.uint32_t[::1] cell_inds cdef np.int64_t[::1] oct_inds - cell_inds = np.full(num_octs*4**3, self.nz[0] * self.nz[1] * self.nz[2], dtype=np.uint32) - oct_inds = np.full(num_octs*4**3, -1, dtype=np.int64) + # Match per-oct cell count in NeighbourCellIndexVisitor.visit() + n_per_oct = ( + (self.nz[0] + 2 * n_ghost_zones) + * (self.nz[1] + 2 * n_ghost_zones) + * (self.nz[2] + 2 * n_ghost_zones) + ) + cell_inds = np.full(num_octs*n_per_oct, self.nz[0] * self.nz[1] * self.nz[2], dtype=np.uint32) + oct_inds = np.full(num_octs*n_per_oct, -1, dtype=np.int64) visitor = NeighbourCellIndexVisitor(self, -1, n_ghost_zones) visitor.cell_inds = cell_inds @@ -929,7 +936,12 @@ cdef class OctreeContainer: cdef int num_octs if num_cells < 0: num_octs = selector.count_octs(self, domain_id) - num_cells = num_octs * 4**3 + # Match per-oct cell count in NeighbourCellVisitor.visit() + num_cells = num_octs * ( + (self.nz[0] + 2 * n_ghost_zones) + * (self.nz[1] + 2 * n_ghost_zones) + * (self.nz[2] + 2 * n_ghost_zones) + ) cdef NeighbourCellVisitor visitor cdef np.ndarray[np.uint8_t, ndim=1] levels diff --git a/yt/geometry/oct_visitors.pyx b/yt/geometry/oct_visitors.pyx index e64cb6bff2..91b0f9e300 100644 --- a/yt/geometry/oct_visitors.pyx +++ b/yt/geometry/oct_visitors.pyx @@ -388,7 +388,7 @@ cdef class BaseNeighbourVisitor(OctVisitor): fcoords[i] += 1 elif fcoords[i] > 1: fcoords[i] -= 1 - local_oct &= (0 <= ishift[i] <= 1) + local_oct &= (0 <= ishift[i] < self.nz[i]) other_oct = not local_oct # Use octree to find neighbour @@ -412,7 +412,7 @@ cdef class BaseNeighbourVisitor(OctVisitor): # Index of neighbouring cell within its oct for i in range(3): - self.neigh_ind[i] = (ishift[i]) % 2 + self.neigh_ind[i] = (ishift[i]) % self.nz[i] self.other_oct = other_oct if other_oct: @@ -460,15 +460,15 @@ cdef class NeighbourCellIndexVisitor(BaseNeighbourVisitor): self.last = o.domain_ind - cdef int i0, i1 + cdef int i0 + cdef int out_of_bound_cell_ind = self.nz[0] * self.nz[1] * self.nz[2] i0 = -self.n_ghost_zones - i1 = 2 + self.n_ghost_zones # Loop over cells in and directly around oct - for i in range(i0, i1): + for i in range(i0, self.nz[0] + self.n_ghost_zones): ishift[0] = i - for j in range(i0, i1): + for j in range(i0, self.nz[1] + self.n_ghost_zones): ishift[1] = j - for k in range(i0, i1): + for k in range(i0, self.nz[2] + self.n_ghost_zones): ishift[2] = k self.set_neighbour_info(o, ishift) @@ -480,7 +480,7 @@ cdef class NeighbourCellIndexVisitor(BaseNeighbourVisitor): neigh_cell_ind = self.neighbour_rind() else: neigh_domain_ind = -1 - neigh_cell_ind = 8 + neigh_cell_ind = out_of_bound_cell_ind self.cell_inds[self.index] = neigh_cell_ind self.domain_inds[self.index] = neigh_domain_ind @@ -506,15 +506,15 @@ cdef class NeighbourCellVisitor(BaseNeighbourVisitor): self.last = o.domain_ind - cdef int i0, i1 + cdef int i0 + cdef int out_of_bound_cell_ind = self.nz[0] * self.nz[1] * self.nz[2] i0 = -self.n_ghost_zones - i1 = 2 + self.n_ghost_zones # Loop over cells in and directly around oct - for i in range(i0, i1): + for i in range(i0, self.nz[0] + self.n_ghost_zones): ishift[0] = i - for j in range(i0, i1): + for j in range(i0, self.nz[1] + self.n_ghost_zones): ishift[1] = j - for k in range(i0, i1): + for k in range(i0, self.nz[2] + self.n_ghost_zones): ishift[2] = k self.set_neighbour_info(o, ishift) @@ -532,7 +532,7 @@ cdef class NeighbourCellVisitor(BaseNeighbourVisitor): neigh_level = 255 neigh_domain = -1 neigh_file_ind = -1 - neigh_cell_ind = 8 + neigh_cell_ind = out_of_bound_cell_ind self.levels[self.index] = neigh_level self.file_inds[self.index] = neigh_file_ind