diff --git a/nose_ignores b/nose_ignores index 6db1089c71f..e98a303a897 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_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/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_pytest.py b/yt/geometry/tests/test_process_octree_locality_pytest.py new file mode 100644 index 00000000000..ab57372d2ef --- /dev/null +++ b/yt/geometry/tests/test_process_octree_locality_pytest.py @@ -0,0 +1,218 @@ +""" +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