Skip to content

BUG: fix hard-coded 2x2x2 zone-count assumptions in octree neighbour-finding (issue 5402) #5508

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jurb33 wants to merge 5 commits into
base: main
Choose a base branch
from
Open

BUG: fix hard-coded 2x2x2 zone-count assumptions in octree neighbour-finding (issue 5402) #5508

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0f66b4a
add error case in tests for issue #5402
jurb33 Aug 26, 2026
2f2e4c6
fix issue #5402
jurb33 Aug 26, 2026
7e320fd
add test commenting and comment function references for tracing
jurb33 Aug 27, 2026
a7b5bbe
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 2, 2026
6fcff74
[STY] Cosmetic changes
cphyc Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 41 additions & 1 deletion yt/data_objects/tests/test_octree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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)):

@cphyc cphyc Sep 2, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would suggest also testing that (1, 1, 1) doesn't work (it doesn't on main).

Suggested change
for nz in ((2, 2, 2), (2, 3, 4)):
for nz in ((1, 1, 1), (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
Comment on lines +168 to +174

@cphyc cphyc Sep 2, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is fairly convoluted - do you think it'd be possible to use instead yt.load_octree(..., num_zones=(1, 2, 3))?

EDIT: I just tested, and the short answer is no... because load_octree expects num_zones to be a scalar.


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
18 changes: 15 additions & 3 deletions yt/geometry/oct_container.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
28 changes: 14 additions & 14 deletions yt/geometry/oct_visitors.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -412,7 +412,7 @@ cdef class BaseNeighbourVisitor(OctVisitor):

# Index of neighbouring cell within its oct
for i in range(3):
self.neigh_ind[i] = <np.uint32_t>(ishift[i]) % 2
self.neigh_ind[i] = <np.uint32_t>(ishift[i]) % self.nz[i]

self.other_oct = other_oct
if other_oct:
Expand Down Expand Up @@ -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):

@cphyc cphyc Sep 2, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How confident are you that this shouldn't be the following?

Suggested change
for i in range(i0, self.nz[0] + self.n_ghost_zones):
for i in range(i0, self.nz[2] + self.n_ghost_zones):

(and self.nz[0] for the innermost loop)

@jurb33 jurb33 Sep 2, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is correct - looking at these lines:
yt/geometry/oct_visitors.pyx:391
local_oct &= (0 <= ishift[i] < self.nz[i])
yt/geometry/oct_visitors.pyx:415
self.neigh_ind[i] = <np.uint32_t>(ishift[i]) % self.nz[i]
is using the same ishift[i] self.nz[i] pairing convention everywhere else in oct_visitors.pyx
Using self.nz[2] would break the test case which is what this was intended to fix. (ran this locally)

As far as the test implementation, yes, I can make these changes. May be a bit since school is starting up for me and I am pursuing different projects.

Thank you for reviewing my changes!

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)

Expand All @@ -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
Expand All @@ -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)

Expand All @@ -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
Expand Down
Loading