From 11ca851682a55e9fc6d5aed534ff92f3c589edf5 Mon Sep 17 00:00:00 2001 From: jurb33 Date: Mon, 10 Aug 2026 13:33:48 -0500 Subject: [PATCH 1/4] add local benchmark test for new particle deposit implementation. Review integration with CI or remove. Contributes to issue #2281. --- .../tests/test_process_octree_locality.py | 129 ++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 yt/geometry/tests/test_process_octree_locality.py diff --git a/yt/geometry/tests/test_process_octree_locality.py b/yt/geometry/tests/test_process_octree_locality.py new file mode 100644 index 00000000000..fef47e0b845 --- /dev/null +++ b/yt/geometry/tests/test_process_octree_locality.py @@ -0,0 +1,129 @@ +"""Regression tests for fake-octree mesh-sampling locality and perf behavior. + +These tests are intended to exercise the current mesh-sampling implementation in +`yt/geometry/particle_deposit.pyx`, specifically +`ParticleDepositOperation.process_octree()`. They validate correctness for +clustered/ordered/random/boundary particle access patterns and provide a local +benchmark for perf regressions. + +These tests are marked `local` / `perf` so they can be run manually with: + + pytest -m local yt/geometry/tests/test_process_octree_locality.py + pytest -m perf yt/geometry/tests/test_process_octree_locality.py + +They are not excluded from the regular suite via a global pytest marker filter, +and they do not require any special global CI flag to keep them separate. +""" + +import time + +import numpy as np +import pytest +from numpy.testing import assert_allclose + +from yt.testing import fake_octree_ds + + +def _create_fake_octree_dataset(): + """Return a small synthetic octree dataset for mesh-sampling tests.""" + return fake_octree_ds(num_zones=8, partial_coverage=0) + + +def _timeit(label, func, *args, **kwargs): + """Measure wall-clock time for a callable. + + This helper is only used by the perf-marked timing regression below. + """ + start = time.perf_counter() + result = func(*args, **kwargs) + elapsed = time.perf_counter() - start + print(f"{label}: {elapsed:.6f}s") + return result, elapsed + + +@pytest.mark.local +def test_fake_octree_mesh_sampling_clustered_points(): + """Clustered particle positions should still return correct mesh-sampled values.""" + ds = _create_fake_octree_dataset() + # This test is intended to exercise the bulk octree lookup path inside + # ParticleDepositOperation.process_octree() for points that are spatially + # clustered and therefore likely to traverse only a small set of nearby cells. + ds.add_mesh_sampling_particle_field(("gas", "density"), ptype="all") + ad = ds.all_data() + + positions = ds.r["all", "particle_position"].to_value("code_length") + center = positions[0] + d = np.linalg.norm(positions - center, axis=1) + cluster_idx = np.argsort(d)[:64] + + actual = ad["all", "cell_gas_density"][cluster_idx].to_value("code_density") + baseline = ds.find_field_values_at_points(("gas", "density"), positions[cluster_idx]).to_value( + "code_density" + ) + + assert_allclose(actual, baseline) + + +@pytest.mark.local +def test_fake_octree_mesh_sampling_ordered_and_random_positions(): + """Ordered and random position sets should match point-sampled reference values.""" + ds = _create_fake_octree_dataset() + # Compare ordered particle lookup against random access to ensure the updated + # octree implementation behaves correctly for both locality-friendly and + # locality-unfriendly access patterns. + ds.add_mesh_sampling_particle_field(("gas", "density"), ptype="all") + ad = ds.all_data() + + positions = ds.r["all", "particle_position"].to_value("code_length") + sorted_idx = np.lexsort((positions[:, 2], positions[:, 1], positions[:, 0]))[:128] + random_idx = np.random.RandomState(1).choice(len(positions), size=128, replace=False) + + actual_sorted = ad["all", "cell_gas_density"][sorted_idx].to_value("code_density") + baseline_sorted = ds.find_field_values_at_points(("gas", "density"), positions[sorted_idx]).to_value( + "code_density" + ) + actual_random = ad["all", "cell_gas_density"][random_idx].to_value("code_density") + baseline_random = ds.find_field_values_at_points(("gas", "density"), positions[random_idx]).to_value( + "code_density" + ) + + assert_allclose(actual_sorted, baseline_sorted) + assert_allclose(actual_random, baseline_random) + + +@pytest.mark.local +def test_fake_octree_mesh_sampling_boundary_positions(): + """Boundary particles should still sample the correct octree cells.""" + ds = _create_fake_octree_dataset() + # This test exercises the edge handling logic in process_octree() by using + # particles near domain boundaries where octree cell membership can be subtle. + ds.add_mesh_sampling_particle_field(("gas", "density"), ptype="all") + ad = ds.all_data() + + positions = ds.r["all", "particle_position"].to_value("code_length") + mask = np.any((positions < 0.1) | (positions > 0.9), axis=1) + boundary_idx = np.nonzero(mask)[0][:64] + if boundary_idx.size == 0: + boundary_idx = np.arange(min(64, len(positions))) + + actual = ad["all", "cell_gas_density"][boundary_idx].to_value("code_density") + baseline = ds.find_field_values_at_points(("gas", "density"), positions[boundary_idx]).to_value( + "code_density" + ) + + assert_allclose(actual, baseline) + + +@pytest.mark.perf +def test_fake_octree_mesh_sampling_timing(): + """Measure fake-octree mesh sampling performance for a small dataset.""" + ds = _create_fake_octree_dataset() + # This timing test is intended for local benchmarking of the updated octree + # lookup path, not for strict CI timing thresholds. + ds.add_mesh_sampling_particle_field(("gas", "density"), ptype="all") + ad = ds.all_data() + + positions = ds.r["all", "particle_position"] + actual, elapsed = _timeit("fake_octree_mesh_sampling", lambda: ad["all", "cell_gas_density"]) + assert actual.shape[0] == positions.shape[0] + assert elapsed >= 0.0 From 983128d340f795f7f8a46da9de3a4adf7a923dd0 Mon Sep 17 00:00:00 2001 From: jgurban2 Date: Wed, 12 Aug 2026 04:43:48 -0500 Subject: [PATCH 2/4] add cache for NNeighbors of particle_deposit's process_octree (get to get_near). Follow up with team. Closes #2281. --- yt/geometry/fake_octree.pyx | 10 +- yt/geometry/oct_container.pxd | 3 + yt/geometry/oct_container.pyx | 78 ++++ yt/geometry/oct_visitors.pxd | 1 + yt/geometry/oct_visitors.pyx | 1 + yt/geometry/particle_deposit.pyx | 33 +- .../tests/test_process_octree_locality.py | 347 +++++++++++------- 7 files changed, 336 insertions(+), 137 deletions(-) diff --git a/yt/geometry/fake_octree.pyx b/yt/geometry/fake_octree.pyx index fc4d52903d6..c84ed8c319c 100644 --- a/yt/geometry/fake_octree.pyx +++ b/yt/geometry/fake_octree.pyx @@ -36,13 +36,13 @@ def create_fake_octree(SparseOctreeContainer oct_handler, for i in range(3): ind[i] = 0 dd[i] = ndd[i] - oct_handler.allocate_domains([max_noct]) + oct_handler.allocate_domains([max_noct], 1) parent = oct_handler.next_root(1, ind) parent.domain = 1 cur_leaf = 8 #we've added one parent... mask = np.ones((max_noct,8),dtype='uint8') - while oct_handler.domains[0].n_assigned < max_noct: - print("root: nocts ", oct_handler.domains[0].n_assigned) + while oct_handler.nocts < max_noct: + print("root: nocts ", oct_handler.nocts) cur_leaf = subdivide(oct_handler, parent, ind, dd, cur_leaf, 0, max_noct, max_level, fsubdivide, mask) return cur_leaf @@ -61,7 +61,7 @@ cdef long subdivide(SparseOctreeContainer oct_handler, cdef float rf #random float from 0-1 if cur_level >= max_level: return cur_leaf - if oct_handler.domains[0].n_assigned >= max_noct: + if oct_handler.nocts >= max_noct: return cur_leaf for i in range(3): ind[i] = ((rand() * 1.0 / RAND_MAX) * dd[i]) @@ -69,7 +69,7 @@ cdef long subdivide(SparseOctreeContainer oct_handler, rf = rand() * 1.0 / RAND_MAX if rf > fsubdivide: ii = cind(ind[0], ind[1], ind[2]) - if parent.children[ii] == NULL: + if parent.children == NULL or parent.children[ii] == NULL: cur_leaf += 7 oct = oct_handler.next_child(1, ind, parent) oct.domain = 1 diff --git a/yt/geometry/oct_container.pxd b/yt/geometry/oct_container.pxd index 166981f98b7..1c44c763e50 100644 --- a/yt/geometry/oct_container.pxd +++ b/yt/geometry/oct_container.pxd @@ -64,6 +64,9 @@ cdef class OctreeContainer: cdef public int num_domains cdef Oct *get(self, np.float64_t ppos[3], OctInfo *oinfo = ?, int max_level = ?) noexcept nogil + cdef Oct *get_near(self, Oct *start_oct, OctInfo *start_oi, + np.float64_t ppos[3], OctInfo *oinfo = ?, + int max_level = ?) noexcept nogil cdef int get_root(self, int ind[3], Oct **o) noexcept nogil cdef Oct **neighbors(self, OctInfo *oinfo, np.int64_t *nneighbors, Oct *o, bint periodicity[3]) diff --git a/yt/geometry/oct_container.pyx b/yt/geometry/oct_container.pyx index 981085b15aa..422aca119b0 100644 --- a/yt/geometry/oct_container.pyx +++ b/yt/geometry/oct_container.pyx @@ -126,6 +126,7 @@ cdef class OctreeContainer: o = &cur.my_objs[cur.n_assigned] o.domain_ind = o.file_ind = 0 o.domain = 1 + o.parent = NULL obj.root_mesh[i][j][k] = o cur.n_assigned += 1 visitor.pos[0] = i @@ -265,6 +266,79 @@ cdef class OctreeContainer: oinfo.left_edge[i] = oinfo.ipos[i] * (oinfo.dds[i] * self.nz[i]) + self.DLE[i] oinfo.level = level return cur + + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.cdivision(True) + cdef Oct *get_near(self, Oct *start_oct, OctInfo *start_oi, + np.float64_t ppos[3], OctInfo *oinfo = NULL, + int max_level = 99) noexcept nogil: + cdef int i + cdef Oct *cur + cdef Oct *next + cdef np.float64_t dds[3] + cdef np.float64_t left_edge[3] + cdef np.float64_t cand_dds[3] + cdef np.float64_t cand_left_edge[3] + cdef np.float64_t mid + cdef np.int64_t ind[3] + cdef np.int64_t level + cdef bint contains + #instead of starting at root, we start at given oct + cur = start_oct + for i in range(3): + dds[i] = start_oi.dds[i] * self.nz[i] + left_edge[i] = start_oi.left_edge[i] + level = start_oi.level + + #walk up towards root, doubling each step until box contains ppos + while level > 0: + contains = True + for i in range(3): + if ppos[i] < left_edge[i] or ppos[i] >= left_edge[i] + dds[i]: + contains = False + break + if contains: + break + cur = cur.parent + for i in range(3): + dds[i] = dds[i] * 2.0 + left_edge[i] = floor((left_edge[i] - self.DLE[i]) / dds[i]) * dds[i] + self.DLE[i] + level -= 1 + #if still not in same root cell as before, start over + contains = True + for i in range(3): + if ppos[i] < left_edge[i] or ppos[i] >= left_edge[i] + dds[i]: + contains = False + break + if not contains: + return self.get(ppos, oinfo, max_level) + #walk back down toward ppos from wherever we ended up + while cur.children != NULL and level < max_level: + for i in range(3): + cand_dds[i] = dds[i] / 2.0 + mid = left_edge[i] + cand_dds[i] + if mid > ppos[i]: + ind[i] = 0 + cand_left_edge[i] = left_edge[i] + else: + ind[i] = 1 + cand_left_edge[i] = mid + next = cur.children[cind(ind[0], ind[1], ind[2])] + if next == NULL: + break + cur = next + level += 1 + for i in range(3): + dds[i] = cand_dds[i] + left_edge[i] = cand_left_edge[i] + if oinfo == NULL: return cur + for i in range(3): + oinfo.dds[i] = dds[i] / self.nz[i] + oinfo.left_edge[i] = left_edge[i] + oinfo.ipos[i] = 0 #not used by particle_deposit loop + oinfo.level = level + return cur def locate_positions(self, np.float64_t[:,:] positions): """ @@ -665,6 +739,7 @@ cdef class OctreeContainer: next = &cont.my_objs[cont.n_assigned] cont.n_assigned += 1 self.root_mesh[ind[0]][ind[1]][ind[2]] = next + next.parent = NULL self.nocts += 1 return next @@ -684,6 +759,7 @@ cdef class OctreeContainer: next = &cont.my_objs[cont.n_assigned] cont.n_assigned += 1 parent.children[cind(ind[0],ind[1],ind[2])] = next + next.parent = parent self.nocts += 1 return next @@ -1097,6 +1173,7 @@ cdef class SparseOctreeContainer(OctreeContainer): tsearch(ikey, &self.tree_root, root_node_compare) self.num_root += 1 self.nocts += 1 + next.parent = NULL return next def allocate_domains(self, domain_counts, int root_nodes): @@ -1217,6 +1294,7 @@ cdef class OctObjectPool(ObjectPool): octs[n].file_ind = octs[n].domain = - 1 octs[n].domain_ind = n + offset octs[n].children = NULL + octs[n].parent = NULL cdef void teardown_objs(self, void *obj, np.uint64_t n, np.uint64_t offset, np.int64_t con_id): diff --git a/yt/geometry/oct_visitors.pxd b/yt/geometry/oct_visitors.pxd index a46dafeeb2d..9b55caff2f3 100644 --- a/yt/geometry/oct_visitors.pxd +++ b/yt/geometry/oct_visitors.pxd @@ -17,6 +17,7 @@ cdef struct Oct: np.int64_t domain_ind # index within the global set of domains np.int64_t domain # (opt) addl int index Oct **children # Up to 8 long + Oct *parent # NULL for root octs cdef struct OctInfo: np.float64_t left_edge[3] diff --git a/yt/geometry/oct_visitors.pyx b/yt/geometry/oct_visitors.pyx index e64cb6bff2f..95d97de0617 100644 --- a/yt/geometry/oct_visitors.pyx +++ b/yt/geometry/oct_visitors.pyx @@ -320,6 +320,7 @@ cdef class LoadOctree(OctVisitor): o.children[ii + i].file_ind = -1 o.children[ii + i].domain = -1 o.children[ii + i].children = NULL + o.children[ii + i].parent = o self.nocts[0] += 1 else: print("SOMETHING IS AMISS", self.index) diff --git a/yt/geometry/particle_deposit.pyx b/yt/geometry/particle_deposit.pyx index 8df22123b73..ea68560dc6d 100644 --- a/yt/geometry/particle_deposit.pyx +++ b/yt/geometry/particle_deposit.pyx @@ -67,6 +67,10 @@ cdef class ParticleDepositOperation: cdef np.int64_t offset, moff cdef Oct *oct cdef np.int8_t use_lvlmax + cdef int this_max_level + cdef Oct *cached_oct = NULL + cdef OctInfo cached_oi + cdef bint in_cached_oct moff = octree.get_domain_offset(domain_id + domain_offset) if lvlmax is None: use_lvlmax = False @@ -87,10 +91,33 @@ cdef class ParticleDepositOperation: # previously generated. This way we can support not knowing the # full octree structure. All we *really* care about is some # arbitrary offset into a field value for deposition. - if not use_lvlmax: - oct = octree.get(pos, &oi) + + #only cache particles if max level is equal or looser with no children + #comparing the limit avoids wrong answers where refinement may be deeper. + this_max_level = lvlmaxval[i] if use_lvlmax else 99 + in_cached_oct = False + if cached_oct != NULL and this_max_level >= cached_oi.level and \ + not (this_max_level > cached_oi.level and cached_oct.children != NULL): + + in_cached_oct = True + for j in range(3): + if pos[j] < cached_oi.left_edge[j] or \ + pos[j] >= cached_oi.left_edge[j] + cached_oi.dds[j] * dims[j]: + in_cached_oct = False + break + if in_cached_oct: + oct = cached_oct + oi = cached_oi + elif cached_oct != NULL: + #different oct, find it but reuse last traversal as starting point + oct = octree.get_near(cached_oct, &cached_oi, pos, &oi, this_max_level) + cached_oct = oct + cached_oi = oi else: - oct = octree.get(pos, &oi, max_level=lvlmaxval[i]) + #should be first particle + oct = octree.get(pos, &oi, max_level=this_max_level) + cached_oct = oct + cached_oi = oi # This next line is unfortunate. Basically it says, sometimes we # might have particles that belong to octs outside our domain. # For the distributed-memory octrees, this will manifest as a NULL diff --git a/yt/geometry/tests/test_process_octree_locality.py b/yt/geometry/tests/test_process_octree_locality.py index fef47e0b845..ab57372d2ef 100644 --- a/yt/geometry/tests/test_process_octree_locality.py +++ b/yt/geometry/tests/test_process_octree_locality.py @@ -1,129 +1,218 @@ -"""Regression tests for fake-octree mesh-sampling locality and perf behavior. - -These tests are intended to exercise the current mesh-sampling implementation in -`yt/geometry/particle_deposit.pyx`, specifically -`ParticleDepositOperation.process_octree()`. They validate correctness for -clustered/ordered/random/boundary particle access patterns and provide a local -benchmark for perf regressions. - -These tests are marked `local` / `perf` so they can be run manually with: - - pytest -m local yt/geometry/tests/test_process_octree_locality.py - pytest -m perf yt/geometry/tests/test_process_octree_locality.py - -They are not excluded from the regular suite via a global pytest marker filter, -and they do not require any special global CI flag to keep them separate. -""" - -import time - -import numpy as np -import pytest -from numpy.testing import assert_allclose - -from yt.testing import fake_octree_ds - - -def _create_fake_octree_dataset(): - """Return a small synthetic octree dataset for mesh-sampling tests.""" - return fake_octree_ds(num_zones=8, partial_coverage=0) - - -def _timeit(label, func, *args, **kwargs): - """Measure wall-clock time for a callable. - - This helper is only used by the perf-marked timing regression below. - """ - start = time.perf_counter() - result = func(*args, **kwargs) - elapsed = time.perf_counter() - start - print(f"{label}: {elapsed:.6f}s") - return result, elapsed - - -@pytest.mark.local -def test_fake_octree_mesh_sampling_clustered_points(): - """Clustered particle positions should still return correct mesh-sampled values.""" - ds = _create_fake_octree_dataset() - # This test is intended to exercise the bulk octree lookup path inside - # ParticleDepositOperation.process_octree() for points that are spatially - # clustered and therefore likely to traverse only a small set of nearby cells. - ds.add_mesh_sampling_particle_field(("gas", "density"), ptype="all") - ad = ds.all_data() - - positions = ds.r["all", "particle_position"].to_value("code_length") - center = positions[0] - d = np.linalg.norm(positions - center, axis=1) - cluster_idx = np.argsort(d)[:64] - - actual = ad["all", "cell_gas_density"][cluster_idx].to_value("code_density") - baseline = ds.find_field_values_at_points(("gas", "density"), positions[cluster_idx]).to_value( - "code_density" - ) - - assert_allclose(actual, baseline) - - -@pytest.mark.local -def test_fake_octree_mesh_sampling_ordered_and_random_positions(): - """Ordered and random position sets should match point-sampled reference values.""" - ds = _create_fake_octree_dataset() - # Compare ordered particle lookup against random access to ensure the updated - # octree implementation behaves correctly for both locality-friendly and - # locality-unfriendly access patterns. - ds.add_mesh_sampling_particle_field(("gas", "density"), ptype="all") - ad = ds.all_data() - - positions = ds.r["all", "particle_position"].to_value("code_length") - sorted_idx = np.lexsort((positions[:, 2], positions[:, 1], positions[:, 0]))[:128] - random_idx = np.random.RandomState(1).choice(len(positions), size=128, replace=False) - - actual_sorted = ad["all", "cell_gas_density"][sorted_idx].to_value("code_density") - baseline_sorted = ds.find_field_values_at_points(("gas", "density"), positions[sorted_idx]).to_value( - "code_density" - ) - actual_random = ad["all", "cell_gas_density"][random_idx].to_value("code_density") - baseline_random = ds.find_field_values_at_points(("gas", "density"), positions[random_idx]).to_value( - "code_density" - ) - - assert_allclose(actual_sorted, baseline_sorted) - assert_allclose(actual_random, baseline_random) - - -@pytest.mark.local -def test_fake_octree_mesh_sampling_boundary_positions(): - """Boundary particles should still sample the correct octree cells.""" - ds = _create_fake_octree_dataset() - # This test exercises the edge handling logic in process_octree() by using - # particles near domain boundaries where octree cell membership can be subtle. - ds.add_mesh_sampling_particle_field(("gas", "density"), ptype="all") - ad = ds.all_data() - - positions = ds.r["all", "particle_position"].to_value("code_length") - mask = np.any((positions < 0.1) | (positions > 0.9), axis=1) - boundary_idx = np.nonzero(mask)[0][:64] - if boundary_idx.size == 0: - boundary_idx = np.arange(min(64, len(positions))) - - actual = ad["all", "cell_gas_density"][boundary_idx].to_value("code_density") - baseline = ds.find_field_values_at_points(("gas", "density"), positions[boundary_idx]).to_value( - "code_density" - ) - - assert_allclose(actual, baseline) - - -@pytest.mark.perf -def test_fake_octree_mesh_sampling_timing(): - """Measure fake-octree mesh sampling performance for a small dataset.""" - ds = _create_fake_octree_dataset() - # This timing test is intended for local benchmarking of the updated octree - # lookup path, not for strict CI timing thresholds. - ds.add_mesh_sampling_particle_field(("gas", "density"), ptype="all") - ad = ds.all_data() - - positions = ds.r["all", "particle_position"] - actual, elapsed = _timeit("fake_octree_mesh_sampling", lambda: ad["all", "cell_gas_density"]) - assert actual.shape[0] == positions.shape[0] - assert elapsed >= 0.0 +""" +Fake-octree mesh-sampling tests for +ParticleDepositOperation.process_octree() (yt/geometry/particle_deposit.pyx). + +Run with: + + pytest yt/geometry/tests/test_process_octree_locality.py -s + +To sweep different particle counts, set YT_PERF_N_PARTICLES (comma-separated) +before that same command -- it replaces the default sweep in +_DEFAULT_N_PARTICLES for both test_fake_octree_mesh_sampling and +test_deep_octree_locality: + + YT_PERF_N_PARTICLES=1000,100000,5000000 pytest yt/geometry/tests/test_process_octree_locality.py -s +""" + +import contextlib +import functools +import io +import os +import time + +import numpy as np +from numpy.random import RandomState +from numpy.testing import assert_allclose +from yt.testing import fake_octree_ds + +_CORRECTNESS_SAMPLE = 200 +_DEFAULT_N_PARTICLES = [128, 1_024, 8_192, 65_536, 262_144] + + +def _create_fake_octree_dataset(): + # mesh_sampling_particle_field() only handles the octree's native 2x2x2 + # branching, so num_zones must stay at 2. + #prng is shared, so initialize it randomly every time + return fake_octree_ds(num_zones=2, partial_coverage=0, prng=RandomState(0x4D3D3D3)) + + +def _clustered(ds, n, seed): + prng = np.random.RandomState(seed) + center = prng.random_sample(3) + pos = np.clip(center + prng.normal(scale=0.02, size=(n, 3)), 0.0, 1.0) + return ds.arr(pos, "code_length") + + +def _random(ds, n, seed): + prng = np.random.RandomState(seed) + return ds.arr(prng.random_sample((n, 3)), "code_length") + + +def _ordered(ds, n, seed): + positions = _random(ds, n, seed) + pos = positions.to_value("code_length") + idx = np.lexsort((pos[:, 2], pos[:, 1], pos[:, 0])) + return positions[idx] + + +def _boundary(ds, n, seed): + prng = np.random.RandomState(seed) + pos = prng.random_sample((n, 3)) + on_low_edge = prng.randint(0, 2, size=n).astype(bool) + pos[:, 0] = np.where(on_low_edge, pos[:, 0] * 0.01, 1.0 - pos[:, 0] * 0.01) + return ds.arr(pos, "code_length") + + +_PATTERNS = { + "clustered": _clustered, + "ordered": _ordered, + "random": _random, + "boundary": _boundary, +} + + +def _reference_density_at(ds, positions): + # Not using ds.find_field_values_at_points() -- it breaks under NumPy 2.x. + pos = positions.to_value("code_length") + values = [ds.point(p)["gas", "density"][0] for p in pos] + return ds.arr(values, "code_density") + + +def _timeit(label, func, *args, **kwargs): + start = time.perf_counter() + result = func(*args, **kwargs) + elapsed = time.perf_counter() - start + print(f"{label}: {elapsed:.6f}s") + return result, elapsed + + +def pytest_generate_tests(metafunc): + if "pattern" in metafunc.fixturenames: + metafunc.parametrize("pattern", list(_PATTERNS)) + if "n_particles" in metafunc.fixturenames: + raw = os.environ.get("YT_PERF_N_PARTICLES") + values = [int(v) for v in raw.split(",")] if raw else _DEFAULT_N_PARTICLES + metafunc.parametrize("n_particles", values) + + +def test_fake_octree_mesh_sampling(n_particles, pattern): + ds = _create_fake_octree_dataset() + positions = _PATTERNS[pattern](ds, n_particles, seed=1) + + ad = ds.all_data() + density = ad["gas", "density"] + obj = ad._current_chunk.objs[0] + mesh_field = np.asarray(density.T.reshape(-1)) + + # Only the actual octree lookup is inside the clock -- setup above and + # the correctness check below both happen outside it. + sampled, elapsed = _timeit( + f"fake_octree_mesh_sampling ({pattern}, n_particles={n_particles})", + obj.mesh_sampling_particle_field, + positions, + mesh_field, + ) + assert sampled.shape[0] == n_particles + assert elapsed >= 0.0 + + # The reference lookup is a slow per-point Python loop, so only check a + # capped sample instead of all n_particles. + check_n = min(n_particles, _CORRECTNESS_SAMPLE) + actual = ds.arr(sampled[:check_n], density.units).to_value("code_density") + baseline = _reference_density_at(ds, positions[:check_n]).to_value("code_density") + assert_allclose(actual, baseline) + + +# The fixture above is a shallow octree (depth ~2) -- there's nothing for a +# same-oct cache or a climb-instead-of-restart lookup to save, no matter how +# good the implementation is. This section builds a genuinely deep octree +# directly (bypassing the fake_octree_ds/load_octree path, which only goes a +# couple levels deep for this test's mask) so the perf numbers below actually +# have tree depth to exploit. + +_DEEP_MAX_LEVEL = 24 +_DEEP_MAX_NOCT = 4000 +_DEEP_FSUBDIVIDE = 0.05 + + +@functools.lru_cache(maxsize=1) +def _build_deep_octree(): + from yt.geometry.fake_octree import create_fake_octree + from yt.geometry.oct_container import RAMSESOctreeContainer + + # create_fake_octree() only ever populates a single root oct (at index + # [0, 0, 0]); domain_dimensions must be [1, 1, 1] so that root's cell + # spans the whole [0, 1)^3 domain. With [2, 2, 2] (8 root cells), + # particles in the other 7 root cells would map to a NULL oct and get + # silently dropped by process_octree. + dd = np.array([1, 1, 1], dtype="i4") + dle = np.array([0.0, 0.0, 0.0], dtype="f8") + dre = np.array([1.0, 1.0, 1.0], dtype="f8") + oct_handler = RAMSESOctreeContainer(dd, dle, dre) + # create_fake_octree prints a line per oct it creates -- silence it. + with contextlib.redirect_stdout(io.StringIO()): + create_fake_octree( + oct_handler, _DEEP_MAX_NOCT, _DEEP_MAX_LEVEL, dd, dle, dre, _DEEP_FSUBDIVIDE + ) + return oct_handler + + +def _random_raw(n, seed): + return RandomState(seed).random_sample((n, 3)) + + +def _ordered_raw(n, seed): + pos = _random_raw(n, seed) + idx = np.lexsort((pos[:, 2], pos[:, 1], pos[:, 0])) + return pos[idx] + + +def _clustered_raw(n, seed): + prng = RandomState(seed) + center = prng.random_sample(3) + return np.clip(center + prng.normal(scale=0.02, size=(n, 3)), 0.0, 1.0) + + +def _boundary_raw(n, seed): + prng = RandomState(seed) + pos = prng.random_sample((n, 3)) + on_low_edge = prng.randint(0, 2, size=n).astype(bool) + pos[:, 0] = np.where(on_low_edge, pos[:, 0] * 0.01, 1.0 - pos[:, 0] * 0.01) + return pos + + +_RAW_PATTERNS = { + "clustered": _clustered_raw, + "ordered": _ordered_raw, + "random": _random_raw, + "boundary": _boundary_raw, +} + + +def test_deep_octree_locality(n_particles, pattern): + from yt.geometry.particle_deposit import CountParticles + + oct_handler = _build_deep_octree() + positions = _RAW_PATTERNS[pattern](n_particles, seed=1) + + dom_ind = np.arange(oct_handler.nocts, dtype=np.int64) + # RAMSESOctreeContainer was built with the default num_zones=2 -- see + # _build_deep_octree, which never overrides it. + nz = (2, 2, 2, oct_handler.nocts) + op = CountParticles(nz, "cubic") + op.initialize() + + _, elapsed = _timeit( + f"deep_octree_process_octree ({pattern}, n_particles={n_particles}, " + f"max_level={_DEEP_MAX_LEVEL}, nocts={oct_handler.nocts})", + op.process_octree, + oct_handler, + dom_ind, + positions, + None, + 1, + 0, + ) + counted = op.finalize() + assert counted.sum() == n_particles + assert elapsed >= 0.0 From 9d05ba3dbd384c01418ca78bc835597905770bf8 Mon Sep 17 00:00:00 2001 From: jgurban2 Date: Wed, 12 Aug 2026 21:10:21 -0500 Subject: [PATCH 3/4] remove test_process_octree_locality from CI Job --- nose_ignores | 1 + 1 file changed, 1 insertion(+) diff --git a/nose_ignores b/nose_ignores index 6db1089c71f..e2583931b5c 100644 --- a/nose_ignores +++ b/nose_ignores @@ -47,6 +47,7 @@ --ignore-file=test_field_parsing\.py --ignore-file=test_disks\.py --ignore-file=test_offaxisprojection_pytestonly\.py +--ignore-file=test_process_octree_locality\.py --ignore-file=test_sph_pixelization_pytestonly\.py --ignore-file=test_time_series\.py --ignore-file=test_cf_radial_pytest\.py From 7b8001751b149c1cc73bcb53b8d2107135f1ca32 Mon Sep 17 00:00:00 2001 From: jgurban2 Date: Wed, 2 Sep 2026 12:43:56 -0500 Subject: [PATCH 4/4] update naming convention for ci_runner --- nose_ignores | 2 +- ...ctree_locality.py => test_process_octree_locality_pytest.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename yt/geometry/tests/{test_process_octree_locality.py => test_process_octree_locality_pytest.py} (100%) diff --git a/nose_ignores b/nose_ignores index e2583931b5c..e98a303a897 100644 --- a/nose_ignores +++ b/nose_ignores @@ -47,7 +47,7 @@ --ignore-file=test_field_parsing\.py --ignore-file=test_disks\.py --ignore-file=test_offaxisprojection_pytestonly\.py ---ignore-file=test_process_octree_locality\.py +--ignore-file=test_process_octree_locality_pytest\.py --ignore-file=test_sph_pixelization_pytestonly\.py --ignore-file=test_time_series\.py --ignore-file=test_cf_radial_pytest\.py diff --git a/yt/geometry/tests/test_process_octree_locality.py b/yt/geometry/tests/test_process_octree_locality_pytest.py similarity index 100% rename from yt/geometry/tests/test_process_octree_locality.py rename to yt/geometry/tests/test_process_octree_locality_pytest.py