From ae10dc0f7be59de342ead10032d3cc2b3ff3e788 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Thu, 16 Jul 2026 15:33:58 +0200 Subject: [PATCH 1/2] feat(octree near): occupancy-bucketed compact W-tile near-field builder Adds `octree_nearfield_tiles.build_octree_nearfield_tiles`: re-packs an adaptive octree's occupancy-bounded near CSR into full uniform W-tiles (cell-respecting, so tile-pairs map 1:1 onto cell-pairs) with a per-tile source-tile list densely compacted to a capacity sized by actual near work (not the (1+u_cap)*mbpl worst case). Static-shape (compiles once, reused across steps); overflow surfaced as a single eager-checked boolean. Drops into jaccpot's octree near field as the `near_mode="compact_wtile"` structure. Tests (tests/unit/test_octree_nearfield_tiles.py, 100% module coverage): tile coverage / cell-respecting partition, source-tile completeness vs the reference leaf-near expansion, and overflow-flag correctness. Co-Authored-By: Claude Opus 4.8 --- tests/unit/test_octree_nearfield_tiles.py | 233 ++++++++++++++++++++ yggdrax/octree_nearfield_tiles.py | 249 ++++++++++++++++++++++ 2 files changed, 482 insertions(+) create mode 100644 tests/unit/test_octree_nearfield_tiles.py create mode 100644 yggdrax/octree_nearfield_tiles.py diff --git a/tests/unit/test_octree_nearfield_tiles.py b/tests/unit/test_octree_nearfield_tiles.py new file mode 100644 index 0000000..befe2f6 --- /dev/null +++ b/tests/unit/test_octree_nearfield_tiles.py @@ -0,0 +1,233 @@ +"""Compact near-field W-tiles (:mod:`yggdrax.octree_nearfield_tiles`). + +The builder re-packs an adaptive octree's occupancy-bounded near CSR into full uniform +W-tiles + a densely-compacted per-tile source-tile list. The gates here certify the three +correctness contracts the compact near-field relies on: + +* (COVERAGE / CELL-RESPECTING) every particle lands in exactly one tile, and every tile's + particles all belong to a single leaf cell (tiles never span cells -> tile-pairs map 1:1 + onto cell-pairs, so the near/far split is preserved with no double-counting). +* (SOURCE COMPLETENESS) for every target tile the compacted source-tile set equals the + reference expansion of the leaf-level near set (own cell's OTHER tiles + all tiles of the + cell's U-neighbour leaves), self-tile excluded. +* (OVERFLOW) the eager-only capacity gate trips exactly when a cap is under-sized. +""" + +from __future__ import annotations + +import jax +import numpy as np +import pytest + +jax.config.update("jax_enable_x64", True) + +from yggdrax.octree_nearfield_tiles import build_octree_nearfield_tiles +from yggdrax.octree_uvwx import build_adaptive_octree_execution_view_device + + +def _points(n, seed, dist): + rng = np.random.default_rng(seed) + if dist == "clustered": # concentrated Plummer-like cloud + r = rng.uniform(size=n) ** (1.0 / 3.0) + r = r / (1.0 - 0.85 * r) + d = rng.standard_normal((n, 3)) + d /= np.linalg.norm(d, axis=1, keepdims=True) + pos = r[:, None] * d + else: + pos = rng.uniform(-1.0, 1.0, size=(n, 3)) + return pos.astype(np.float64) + + +def _build_view(pos, max_depth, leaf_size, u_cap=256): + return build_adaptive_octree_execution_view_device( + pos, + max_depth, + leaf_size, + node_capacity=12000, + leaf_capacity=10000, + interactions=True, + u_capacity=u_cap, + ) + + +def _probe_caps(view, n, w): + """Eager reference: exact num_tiles, per-leaf tile counts, and max source tiles.""" + nr = np.asarray(view.node_ranges) + leaf_idx = np.asarray(view.leaf_indices) + u_nbr = np.asarray(view.u_neighbors) + node_cap = int(view.parent.shape[0]) + leaf_cap = int(leaf_idx.shape[0]) + u_width = u_nbr.shape[0] // leaf_cap + + occ = np.zeros(leaf_cap, np.int64) + for r in range(leaf_cap): + ln = int(leaf_idx[r]) + if ln < node_cap: + occ[r] = int(nr[ln, 1] - nr[ln, 0] + 1) + tiles_per_leaf = (occ + w - 1) // w + tile_base = np.concatenate([[0], np.cumsum(tiles_per_leaf)])[:-1] + total_tiles = int(tiles_per_leaf.sum()) + + # max source tiles over target tiles (self slot INCLUDED, matching the builder's counts) + max_src = 0 + for r in range(leaf_cap): + ln = int(leaf_idx[r]) + if ln >= node_cap or occ[r] <= 0: + continue + nbr_rows = [r] + for e in u_nbr[r * u_width : (r + 1) * u_width].tolist(): + if e != node_cap and e < leaf_cap and e != r: + nbr_rows.append(int(e)) + src_count = int(sum(int(tiles_per_leaf[x]) for x in nbr_rows)) + max_src = max(max_src, src_count) + return dict( + node_cap=node_cap, + leaf_cap=leaf_cap, + u_width=u_width, + occ=occ, + tiles_per_leaf=tiles_per_leaf, + tile_base=tile_base, + total_tiles=total_tiles, + max_src=max_src, + max_occ=int(occ.max()), + ) + + +@pytest.mark.parametrize( + ("n", "max_depth", "leaf_size", "dist", "w"), + [ + (2000, 6, 64, "uniform", 32), + (2000, 6, 64, "clustered", 32), + (4000, 7, 96, "clustered", 64), + (8000, 7, 128, "uniform", 64), + (8000, 8, 128, "clustered", 64), + ], +) +def test_tiles_coverage_and_cell_respecting(n, max_depth, leaf_size, dist, w): + """Every particle in exactly one tile; every tile confined to one leaf cell.""" + pos = _points(n, 7, dist) + view = _build_view(pos, max_depth, leaf_size) + ref = _probe_caps(view, n, w) + + tiles = build_octree_nearfield_tiles( + view, + tile_width=w, + num_tiles_cap=(n + w - 1) // w + ref["leaf_cap"], + max_source_tiles=ref["max_src"] + 8, + max_leaf_occupancy=ref["max_occ"], + ) + assert not bool(np.asarray(tiles.overflow)), "unexpected overflow with ample caps" + assert int(np.asarray(tiles.num_tiles_used)) == ref["total_tiles"] + + pidx = np.asarray(tiles.tile_particle_idx) + tmask = np.asarray(tiles.tile_mask) + tleaf = np.asarray(tiles.tile_leaf_row) + nr = np.asarray(view.node_ranges) + leaf_idx = np.asarray(view.leaf_indices) + + # (COVERAGE) union of all tiles' particles == [0, N), each exactly once + cov = np.zeros(n, np.int64) + n_tiles = int(np.asarray(tiles.num_tiles_used)) + for t in range(n_tiles): + parts = pidx[t][tmask[t]] + cov[parts] += 1 + # (CELL-RESPECTING) all this tile's particles lie in its owning leaf's Morton span + r = int(tleaf[t]) + assert r < ref["leaf_cap"], f"tile {t} has sentinel leaf row" + ln = int(leaf_idx[r]) + assert np.all(parts >= nr[ln, 0]) and np.all(parts <= nr[ln, 1]) + assert int(cov.min()) == 1 and int(cov.max()) == 1, ( + f"not an exact partition: missing={int((cov==0).sum())} " + f"double={int((cov>1).sum())}" + ) + + +@pytest.mark.parametrize( + ("n", "max_depth", "leaf_size", "dist", "w"), + [ + (2000, 6, 64, "clustered", 32), + (4000, 7, 96, "clustered", 64), + (8000, 8, 128, "clustered", 64), + ], +) +def test_tile_source_completeness(n, max_depth, leaf_size, dist, w): + """Each target tile's compacted source set == reference leaf-near expansion, minus self.""" + pos = _points(n, 11, dist) + view = _build_view(pos, max_depth, leaf_size) + ref = _probe_caps(view, n, w) + + tiles = build_octree_nearfield_tiles( + view, + tile_width=w, + num_tiles_cap=(n + w - 1) // w + ref["leaf_cap"], + max_source_tiles=ref["max_src"] + 8, + max_leaf_occupancy=ref["max_occ"], + ) + assert not bool(np.asarray(tiles.overflow)) + + tleaf = np.asarray(tiles.tile_leaf_row) + sids = np.asarray(tiles.tile_source_ids) + svalid = np.asarray(tiles.tile_source_valid) + leaf_idx = np.asarray(view.leaf_indices) + u_nbr = np.asarray(view.u_neighbors) + node_cap = ref["node_cap"] + leaf_cap = ref["leaf_cap"] + u_width = ref["u_width"] + tiles_per_leaf = ref["tiles_per_leaf"] + tile_base = ref["tile_base"] + n_tiles = int(np.asarray(tiles.num_tiles_used)) + + def tiles_of_leaf(x): + return set(range(int(tile_base[x]), int(tile_base[x]) + int(tiles_per_leaf[x]))) + + n_bad = 0 + for t in range(n_tiles): + r = int(tleaf[t]) + ln = int(leaf_idx[r]) + if ln >= node_cap: + continue + expected = set(tiles_of_leaf(r)) # own cell's tiles + for e in u_nbr[r * u_width : (r + 1) * u_width].tolist(): + if e != node_cap and e < leaf_cap and e != r: + expected |= tiles_of_leaf(int(e)) + expected.discard(t) # self tile handled by the self-block path + got = set(int(x) for x in sids[t][svalid[t]].tolist()) + assert len(got) == int(svalid[t].sum()), f"tile {t}: duplicate source ids" + if got != expected: + n_bad += 1 + if n_bad <= 3: + print( + f"tile {t} (leaf {r}): missing={sorted(expected-got)[:5]} " + f"extra={sorted(got-expected)[:5]}" + ) + assert n_bad == 0, f"{n_bad} tiles have an incorrect source set" + + +def test_overflow_flag_trips_on_undersized_caps(): + """The eager capacity gate trips on each under-sized cap and clears with ample caps.""" + pos = _points(4000, 3, "clustered") + view = _build_view(pos, 7, 96) + w = 64 + ref = _probe_caps(view, 4000, w) + ample = dict( + tile_width=w, + num_tiles_cap=(4000 + w - 1) // w + ref["leaf_cap"], + max_source_tiles=ref["max_src"] + 8, + max_leaf_occupancy=ref["max_occ"], + ) + assert not bool(np.asarray(build_octree_nearfield_tiles(view, **ample).overflow)) + + # under-size the source-tile capacity + bad_src = dict(ample) + bad_src["max_source_tiles"] = max(1, ref["max_src"] // 2) + assert bool(np.asarray(build_octree_nearfield_tiles(view, **bad_src).overflow)) + + # under-size the per-leaf occupancy bound (fewer tiles per leaf than needed) + bad_occ = dict(ample) + bad_occ["max_leaf_occupancy"] = w # forces mbpl=1 while dense leaves need more + assert bool(np.asarray(build_octree_nearfield_tiles(view, **bad_occ).overflow)) + + # under-size the global tile cap + bad_tiles = dict(ample) + bad_tiles["num_tiles_cap"] = max(1, ref["total_tiles"] // 2) + assert bool(np.asarray(build_octree_nearfield_tiles(view, **bad_tiles).overflow)) diff --git a/yggdrax/octree_nearfield_tiles.py b/yggdrax/octree_nearfield_tiles.py new file mode 100644 index 0000000..fc10369 --- /dev/null +++ b/yggdrax/octree_nearfield_tiles.py @@ -0,0 +1,249 @@ +"""Occupancy-bucketed, padding-free near-field W-tiles for the adaptive octree. + +The adaptive octree execution view (:func:`build_adaptive_octree_execution_view_device`) +exposes an occupancy-bounded near list as a fixed-width CSR (``u_offsets`` / ``u_neighbors`` +over leaf ROW ids, self-excluded). The straightforward near-field kernel groups each leaf +into fixed-``B`` blocks and gives every block a source list of width ``(1 + u_capacity) * +ceil(max_leaf_size / B)`` -- with ``u_capacity`` (256) hugely over-provisioned vs the ~27-56 +real neighbours, and the block factor sized by the *max* leaf occupancy. That ~48x padding is +the whole near-field wall on the concentrated-galaxy IC. + +This module builds the SAME near tile-pairs but **compacted**: it re-packs each adaptive cell +into full uniform **W-tiles** (small width ``W`` = 32/64, cell-respecting so tile-pairs map +1:1 onto cell-pairs -- no cross-cell masking, no double-counting), and it densely packs each +tile's source-tile ids into a fixed capacity sized by the *actual* near work (probed), not the +worst-case product. Overflow (a leaf needing more tiles than ``max_tiles_per_leaf``, a tile +needing more sources than ``max_source_tiles``, or more tiles than ``num_tiles_cap``) is +reported as a single boolean the caller checks EAGER-only -- never per-step under a scan. + +The output :class:`OctreeNearfieldTiles` drops straight into jaccpot's ``_octree_near_field``: +``tile_particle_idx``/``tile_mask`` become the ``(num_tiles, W, ...)`` gather table (tiles as +"leaves", ``W`` as ``leaf_width``) and ``tile_source_ids``/``tile_source_valid`` become the +per-tile source-id list for the topology-agnostic ``nearfield_leafpair`` Pallas kernel or the +pure-JAX chunked edge scan. Intra-tile self (``i != j``) is handled by the caller's existing +self-block path; this builder EXCLUDES the target tile from its own source list. + +Everything is static-shape (shapes depend only on ``(num_tiles_cap, W, max_source_tiles)``), +so the build compiles once and is reused across timesteps. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import jax.numpy as jnp +from jax import Array +from jax.typing import DTypeLike + +from .dtypes import INDEX_DTYPE +from .octree_uvwx import UniformOctreeExecutionView + +__all__ = [ + "OctreeNearfieldTiles", + "build_octree_nearfield_tiles", +] + + +class OctreeNearfieldTiles(NamedTuple): + """Compact, cell-respecting near-field W-tiles for the adaptive octree. + + All array shapes are static functions of ``(num_tiles, W, max_source_tiles)``; the two + ``num_*`` scalars are data-dependent (traced under jit) and ``overflow`` is the eager-only + capacity gate. ``sentinel`` marks unused padded rows/slots. + """ + + tile_particle_idx: Array # (num_tiles, W) Morton particle indices per tile + tile_mask: Array # (num_tiles, W) bool: real particle vs pad + tile_leaf_row: Array # (num_tiles,) owning leaf ROW id (sentinel for unused tiles) + tile_source_ids: Array # (num_tiles, max_source_tiles) source TILE ids (self-excluded) + tile_source_valid: Array # (num_tiles, max_source_tiles) bool + num_tiles_used: Array # () traced: total real tiles (<= num_tiles cap) + num_source_tiles_max: Array # () traced: max source tiles over target tiles (incl self slot) + overflow: Array # () bool: any capacity exceeded (check eager-only) + + +def _compact_rows_to_width( + values: Array, + valid: Array, + *, + out_width: int, + fill: int, + index_dtype: DTypeLike, +) -> tuple[Array, Array, Array]: + """Pack the valid entries of each row to the front of a fixed-width buffer. + + ``values`` / ``valid`` are ``(M, C)``. Returns ``(packed, packed_valid, counts)`` where + ``packed`` / ``packed_valid`` are ``(M, out_width)`` (valid entries of row ``m`` in + ``values[m]`` moved to columns ``0..counts[m]-1``, padded with ``fill`` / ``False``) and + ``counts`` is ``(M,)`` the number of valid entries per row (may exceed ``out_width``; the + overflow is dropped and reported via ``counts > out_width``). + """ + + m = int(values.shape[0]) + w = int(out_width) + valid_i = valid.astype(index_dtype) + # exclusive prefix sum along the row = destination column of each valid entry + prefix = jnp.cumsum(valid_i, axis=1) - valid_i # (M, C) + counts = jnp.sum(valid_i, axis=1) # (M,) + place = valid & (prefix < w) # entries that fit + row = jnp.arange(m, dtype=index_dtype)[:, None] + oor = jnp.asarray(m * w, dtype=index_dtype) # out-of-range slot -> dropped + dest = jnp.where(place, row * w + prefix, oor) + packed = ( + jnp.full((m * w,), fill, dtype=index_dtype) + .at[dest] + .set(jnp.where(place, values, fill).astype(index_dtype), mode="drop") + .reshape(m, w) + ) + packed_valid = ( + jnp.zeros((m * w,), dtype=bool) + .at[dest] + .set(place, mode="drop") + .reshape(m, w) + ) + return packed, packed_valid, counts + + +def build_octree_nearfield_tiles( + view: UniformOctreeExecutionView, + *, + tile_width: int, + num_tiles_cap: int, + max_source_tiles: int, + max_leaf_occupancy: int, + index_dtype: DTypeLike = INDEX_DTYPE, +) -> OctreeNearfieldTiles: + """Build compact, cell-respecting near-field W-tiles from an adaptive octree view. + + Args: + view: adaptive octree execution view built with ``interactions=True`` (must carry + ``leaf_indices``, ``node_ranges`` and the near CSR ``u_offsets`` / ``u_neighbors``). + tile_width: ``W``, the fixed tile width (e.g. 32 or 64). Small ``W`` collapses the + within-tile padding on under-full cells. + num_tiles_cap: static upper bound on the number of tiles; a safe value is + ``ceil(N / W) + leaf_capacity`` (each cell contributes at most one partial tail + tile beyond ``floor(occ / W)``). + max_source_tiles: static per-target-tile source-id capacity ``S``. Size from an eager + probe of the near CSR (max over tiles of the summed source-tile count, self slot + included) with a modest safety factor; under-sizing trips ``overflow`` (loud, not + silent) rather than corrupting forces. + max_leaf_occupancy: static bound on per-leaf particle count (the adaptive + ``leaf_size``); ``ceil(max_leaf_occupancy / W)`` is the per-leaf tile cap used to + expand source cells. A leaf exceeding it (a ``max_depth``-forced dense leaf) trips + ``overflow``. + index_dtype: integer dtype for id arrays. + + Returns: + :class:`OctreeNearfieldTiles`. + + A tile ``t`` owned by leaf row ``r`` has as its source tiles: all tiles of ``r``'s near + leaves (the CSR row ``u_neighbors[r]``, self-excluded) plus ``r``'s OWN other tiles + (intra-cell near), with ``t`` itself removed (intra-tile ``i != j`` self handled by the + caller). Because tiles never span cells, this reproduces exactly the leaf-pair near set the + non-compact path evaluates, up to the compaction. + """ + + idt = index_dtype + w = int(tile_width) + num_tiles = int(num_tiles_cap) + s_cap = int(max_source_tiles) + mbpl = (int(max_leaf_occupancy) + w - 1) // w # static max tiles per leaf + + node_ranges = jnp.asarray(view.node_ranges, dtype=idt) + leaf_indices = jnp.asarray(view.leaf_indices, dtype=idt) + num_nodes = int(node_ranges.shape[0]) + r_leaves = int(leaf_indices.shape[0]) # leaf_capacity (rows, sentinel-padded) + u_cap = int(jnp.asarray(view.u_neighbors).shape[0] // max(r_leaves, 1)) + + # ---- per-leaf geometry (Morton particle span) ---- + valid_leaf = leaf_indices < num_nodes # real leaf vs sentinel pad + leaf_safe = jnp.where(valid_leaf, leaf_indices, 0) + lo = node_ranges[leaf_safe, 0] # (R,) Morton start index + occ = jnp.where( + valid_leaf, + node_ranges[leaf_safe, 1] - node_ranges[leaf_safe, 0] + 1, + 0, + ) # (R,) particle count + + # ---- per-leaf tile counts + global (level-agnostic) tile numbering ---- + tiles_per_leaf = jnp.where(valid_leaf, (occ + (w - 1)) // w, 0) # (R,) + cum = jnp.cumsum(tiles_per_leaf) # (R,) inclusive + tile_base = cum - tiles_per_leaf # (R,) first global tile id of leaf r + total_tiles = jnp.sum(tiles_per_leaf) + + g = jnp.arange(num_tiles, dtype=idt) # (num_tiles,) + valid_tile = g < total_tiles + tile_leaf_row = jnp.clip( + jnp.searchsorted(cum, g, side="right"), 0, max(r_leaves - 1, 0) + ).astype(idt) + r_of_tile = tile_leaf_row + local = g - tile_base[r_of_tile] # (num_tiles,) local tile index within its leaf + + # ---- tile -> particle gather indices (Morton order), padded ---- + slot = jnp.arange(w, dtype=idt) # (W,) + within = local[:, None] * w + slot[None, :] # (num_tiles, W) local particle offset + tile_mask = valid_tile[:, None] & (within < occ[r_of_tile][:, None]) + tile_particle_idx = jnp.where(tile_mask, lo[r_of_tile][:, None] + within, 0) + + # ---- per-leaf compacted source-tile list ---- + # neighbour leaf rows = own leaf (self, first slot) + U-neighbours (CSR, self-excluded). + src_rows = jnp.asarray(view.u_neighbors, dtype=idt).reshape(r_leaves, u_cap) + own_row = jnp.arange(r_leaves, dtype=idt)[:, None] + nbr = jnp.concatenate([own_row, src_rows], axis=1) # (R, 1+u_cap) + nbr_valid = jnp.concatenate( + [ + valid_leaf[:, None], + (src_rows < r_leaves) & (src_rows != own_row), # guard against a self entry + ], + axis=1, + ) + nbr_safe = jnp.where(nbr_valid, nbr, 0) + + # expand each neighbour leaf into its (up to mbpl) tiles + k = jnp.arange(mbpl, dtype=idt) # (mbpl,) + cand = tile_base[nbr_safe][:, :, None] + k[None, None, :] # (R, 1+u_cap, mbpl) + cand_valid = nbr_valid[:, :, None] & ( + k[None, None, :] < tiles_per_leaf[nbr_safe][:, :, None] + ) + c_width = (1 + u_cap) * mbpl + cand = cand.reshape(r_leaves, c_width) + cand_valid = cand_valid.reshape(r_leaves, c_width) + + leaf_src_ids, leaf_src_valid, leaf_src_counts = _compact_rows_to_width( + cand, cand_valid, out_width=s_cap, fill=0, index_dtype=idt + ) + + # ---- broadcast per-leaf source list to tiles, then exclude the target tile itself ---- + tile_source_ids = leaf_src_ids[r_of_tile] # (num_tiles, S) + tile_source_valid = ( + leaf_src_valid[r_of_tile] + & valid_tile[:, None] + & (tile_source_ids != g[:, None]) # drop the self tile (intra-tile handled elsewhere) + ) + tile_source_ids = jnp.where(tile_source_ids < total_tiles, tile_source_ids, 0) + tile_source_ids = jnp.where(tile_source_valid, tile_source_ids, 0) + + # ---- unused-tile sentinel + capacity gate (traced; caller checks eager-only) ---- + sentinel = jnp.asarray(num_tiles, dtype=idt) + tile_leaf_row = jnp.where(valid_tile, r_of_tile, sentinel) + + tiles_per_leaf_max = jnp.max(tiles_per_leaf) if r_leaves > 0 else jnp.asarray(0, idt) + num_source_tiles_max = ( + jnp.max(leaf_src_counts) if r_leaves > 0 else jnp.asarray(0, idt) + ) + overflow = ( + (total_tiles > jnp.asarray(num_tiles, idt)) + | (tiles_per_leaf_max > jnp.asarray(mbpl, idt)) + | (num_source_tiles_max > jnp.asarray(s_cap, idt)) + ) + + return OctreeNearfieldTiles( + tile_particle_idx=tile_particle_idx.astype(idt), + tile_mask=tile_mask, + tile_leaf_row=tile_leaf_row.astype(idt), + tile_source_ids=tile_source_ids.astype(idt), + tile_source_valid=tile_source_valid, + num_tiles_used=total_tiles.astype(idt), + num_source_tiles_max=num_source_tiles_max.astype(idt), + overflow=overflow, + ) From 43c33a071ce39c3b0330d58eb4a35f281e2688a8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:35:02 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- yggdrax/octree_nearfield_tiles.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/yggdrax/octree_nearfield_tiles.py b/yggdrax/octree_nearfield_tiles.py index fc10369..61d8417 100644 --- a/yggdrax/octree_nearfield_tiles.py +++ b/yggdrax/octree_nearfield_tiles.py @@ -55,10 +55,14 @@ class OctreeNearfieldTiles(NamedTuple): tile_particle_idx: Array # (num_tiles, W) Morton particle indices per tile tile_mask: Array # (num_tiles, W) bool: real particle vs pad tile_leaf_row: Array # (num_tiles,) owning leaf ROW id (sentinel for unused tiles) - tile_source_ids: Array # (num_tiles, max_source_tiles) source TILE ids (self-excluded) + tile_source_ids: ( + Array # (num_tiles, max_source_tiles) source TILE ids (self-excluded) + ) tile_source_valid: Array # (num_tiles, max_source_tiles) bool num_tiles_used: Array # () traced: total real tiles (<= num_tiles cap) - num_source_tiles_max: Array # () traced: max source tiles over target tiles (incl self slot) + num_source_tiles_max: ( + Array # () traced: max source tiles over target tiles (incl self slot) + ) overflow: Array # () bool: any capacity exceeded (check eager-only) @@ -96,10 +100,7 @@ def _compact_rows_to_width( .reshape(m, w) ) packed_valid = ( - jnp.zeros((m * w,), dtype=bool) - .at[dest] - .set(place, mode="drop") - .reshape(m, w) + jnp.zeros((m * w,), dtype=bool).at[dest].set(place, mode="drop").reshape(m, w) ) return packed, packed_valid, counts @@ -218,7 +219,9 @@ def build_octree_nearfield_tiles( tile_source_valid = ( leaf_src_valid[r_of_tile] & valid_tile[:, None] - & (tile_source_ids != g[:, None]) # drop the self tile (intra-tile handled elsewhere) + & ( + tile_source_ids != g[:, None] + ) # drop the self tile (intra-tile handled elsewhere) ) tile_source_ids = jnp.where(tile_source_ids < total_tiles, tile_source_ids, 0) tile_source_ids = jnp.where(tile_source_valid, tile_source_ids, 0) @@ -227,7 +230,9 @@ def build_octree_nearfield_tiles( sentinel = jnp.asarray(num_tiles, dtype=idt) tile_leaf_row = jnp.where(valid_tile, r_of_tile, sentinel) - tiles_per_leaf_max = jnp.max(tiles_per_leaf) if r_leaves > 0 else jnp.asarray(0, idt) + tiles_per_leaf_max = ( + jnp.max(tiles_per_leaf) if r_leaves > 0 else jnp.asarray(0, idt) + ) num_source_tiles_max = ( jnp.max(leaf_src_counts) if r_leaves > 0 else jnp.asarray(0, idt) )