feat(corrfunc): leaf-only KD-tree backend (exact pair coverage)#38
Merged
Merged
Conversation
…traversal Multi-GPU tree build and cross-walk now support all three backends. Distributed backend selector (radix + octree): - local_tree: DISTRIBUTED_TREE_TYPES + _validate_distributed_tree_type + shard_map-safe _build_local_tree dispatch; distributed_tree_moments / build_local_moments gain tree_type (root via argmin(parent)). - let: tree_type threaded to the local build in build_distributed_coarse_tree / classify_against_remote / distributed_let_import. Coarse tree stays radix (internal remote-COM representation; robust at coarse leaf_size=1). Leaf-only (bucket) median-split KD-tree (build_leaf_kdtree, LeafKDTree): - All points in leaves; internal nodes are pure split planes; compact contiguous node_ranges tiling [0,N) + nodes_by_level -> satisfies the same topology contract as radix/octree, giving complete pair coverage. The heap KD-tree (build_kdtree) used by the KNN kernels is untouched; the generic pipeline / tree_type="kdtree" / distributed now use the leaf-only build. Validated: 46 kd unit tests (heap KNN preserved + leaf-kd interactions parity vs radix); distributed CPU 9 passed incl. kd cross-walk complete disjoint source partition (the invariant the heap KD-tree violated). GPU validation in progress. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> (cherry picked from commit 3b1138f)
build_pair_topology no longer rejects backend="kdtree": the leaf-only bucket KD-tree (yggdrax.kdtree.build_leaf_kdtree, via tree_type="kdtree") stores every point in a leaf, so the near-list tiles all pairs. Parametrize the vs-baseline tests over kdtree and replace the rejection test with a near-field-exact check (kdtree soft counts == brute force to <1e-10 at theta=0). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
…pydoclint) Same pydoclint fix as on feat/distributed-tree-backends; pre-commit.ci already auto-fixed black. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…dep) The core type gate installs `.[dev]`, not `.[paper]`, so basedpyright cannot resolve optax in applications/corrfunc/inference_demo.py (a paper illustration). Exclude just that demo file; estimator/baselines/binning and svgd stay type-checked. No test imports optax, so pytest is unaffected. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…d/tested The differentiable-applications code (yggdrax/applications/corrfunc/inference_demo.py) and the corrfunc bench smoke tests import optax, which lives in the `paper` optional-deps extra. The CI job installed only `.[dev]`, so basedpyright could not resolve optax and `test_corrfunc_bench_smoke[recovery]` failed (its subprocess imports inference_demo -> optax). Install `.[dev,paper]` instead; this fixes both and lets basedpyright type-check inference_demo again (so the earlier targeted exclude is reverted). Verified locally: basedpyright 0 errors; tests/bench corrfunc (4) + differentiability (3) smoke tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Enables the
kdtreebackend for the corrfunc pair-counting estimator, targetingpaper/differentiable-applications(where the corrfunc application lives).Because the paper branch does not yet contain the leaf-only KD-tree work, this branch stacks that change onto paper (cherry-pick of the distributed/leaf-kd commit) and then adds the corrfunc-specific enablement:
applications/corrfunc/estimator.py:build_pair_topologyno longer rejectsbackend="kdtree". The leaf-only bucket KD-tree stores every point in a leaf, so the near-list tiles all pairs — the heap KD-tree could not (its internal-node pivots were absent from the near-list).tests/applications/test_corrfunc_vs_baseline.py: parametrize the vs-baseline tests overkdtree, and replace the rejection test with a near-field-exact check.Validation
corrfunc vs brute-force baseline: 8 passed — including
kdtreenear-field exact vs brute force (rel < 1e-10attheta=0) and the far-field-error-shrinks-with-theta test.Relationship to the core change
The leaf-only KD-tree + distributed backend selector is proposed standalone against
mainin the companion PR #37. This PR carries the same change (re-applied onto the paper line, which lacks it) plus the corrfunc enablement, so the paper branch can usebackend="kdtree"directly.🤖 Generated with Claude Code