yggdrax 0.0.1__tar.gz

yggdrax-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 AstroAI Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

yggdrax-0.0.1/PKG-INFO

@@ -0,0 +1,251 @@
+Metadata-Version: 2.4
+Name: yggdrax
+Version: 0.0.1
+Summary: Tree generation and traversal package for hierarchical solvers
+Author: AstroAI Lab
+License: MIT
+Project-URL: Homepage, https://github.com/AstroAI-Lab/yggdrax
+Project-URL: Repository, https://github.com/AstroAI-Lab/yggdrax
+Project-URL: Issues, https://github.com/AstroAI-Lab/yggdrax/issues
+Keywords: fmm,hierarchical,n-body,jax,treecode
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.8.0
+Requires-Dist: jaxlib>=0.4.0
+Requires-Dist: jaxtyping>=0.2.23
+Requires-Dist: beartype>=0.14.0
+Provides-Extra: dev
+Requires-Dist: black>=24.8.0; extra == "dev"
+Requires-Dist: isort>=5.13.2; extra == "dev"
+Requires-Dist: pydoclint>=0.6.0; extra == "dev"
+Requires-Dist: pytest>=8.3.0; extra == "dev"
+Requires-Dist: pytest-cov>=5.0.0; extra == "dev"
+Dynamic: license-file
+
+# Yggdrax
+
+![Black](https://img.shields.io/badge/code%20style-black-000000.svg)
+![isort](https://img.shields.io/badge/imports-isort-1674b1.svg)
+![pydoclint](https://img.shields.io/badge/docstrings-pydoclint-2ea44f.svg)
+
+![Yggdrax Logo](./yggdrax.png)
+
+Yggdrax is a JAX-first tree toolkit for hierarchical N-body solvers. It
+provides Morton ordering, radix tree builders, per-node geometry, and dual-tree
+interaction traversal primitives designed for downstream FMM and treecode
+pipelines. The public `octree` backend now layers explicit octree-cell metadata
+on top of the existing Morton/radix construction path so downstream FMM code
+can consume both the proven traversal buffers and octree-style child tables.
+
+## Features
+
+- Morton encode/decode and stable Morton sorting for 3D points
+- LBVH and fixed-depth radix tree construction
+- Explicit octree metadata derived from Morton/radix topology
+- Tree geometry extraction (bounds, centers, extents, radii)
+- Dual-tree far-field and near-field interaction builders
+- Dense and grouped interaction buffer transforms for batched kernels
+- Prepared artifact utilities for downstream solver integrations
+
+## Installation
+
+Install from source:
+
+```bash
+pip install -e .
+```
+
+Install with development tools:
+
+```bash
+pip install -e ".[dev]"
+```
+
+## Quick Start
+
+```python
+import jax
+import jax.numpy as jnp
+
+from yggdrax import (
+    DualTreeTraversalConfig,
+    build_interactions_and_neighbors,
+    build_octree,
+    compute_tree_geometry,
+)
+
+key = jax.random.PRNGKey(0)
+key_pos, key_mass = jax.random.split(key)
+positions = jax.random.uniform(key_pos, (512, 3), minval=-1.0, maxval=1.0)
+masses = jax.random.uniform(key_mass, (512,), minval=0.5, maxval=1.5)
+
+tree = build_octree(positions, masses, leaf_size=16)
+positions_sorted = positions[tree.particle_indices]
+geom = compute_tree_geometry(tree, positions_sorted)
+traversal_cfg = DualTreeTraversalConfig(
+    max_pair_queue=8192,
+    process_block=256,
+    max_interactions_per_node=2048,
+    max_neighbors_per_leaf=2048,
+)
+interactions, neighbors = build_interactions_and_neighbors(
+    tree,
+    geom,
+    theta=0.6,
+    mac_type="dehnen",
+    traversal_config=traversal_cfg,
+)
+```
+
+`build_tree(...)` continues to expose the radix/LBVH backend directly. The
+octree wrappers (`build_octree(...)`, `build_fixed_depth_octree(...)`) preserve
+the same compatibility fields while additionally exposing explicit octree
+buffers such as `oct_children`, `oct_node_depths`, and `radix_node_to_oct`.
+
+Advanced users can override the built-in MAC with a JAX-traceable pair policy:
+
+```python
+def pair_policy(policy_state, **pair_data):
+    action = ...
+    tag = ...
+    return action, tag
+
+interactions, neighbors, result = build_interactions_and_neighbors(
+    tree,
+    geom,
+    pair_policy=pair_policy,
+    policy_state=...,
+    return_result=True,
+)
+```
+
+The policy receives generic pair geometry/state and returns:
+- `action`: one of accept-far / accept-near / refine
+- `tag`: integer metadata stored for accepted far pairs
+
+When `return_result=True`, raw far-pair tags are available on
+`result.interaction_tags`. This is intended for downstream solvers that need
+solver-side scheduling or adaptive-order bucketing without moving solver logic
+into `yggdrax`.
+
+See `examples/getting_started.ipynb` for a runnable walkthrough.
+For the locked high-performance GPU benchmark configuration, see
+`docs/gpu_benchmark_recommended_setup.md` and
+`examples/tree_gpu_performance_scaling.ipynb`.
+
+## KD-Tree MAC Note
+
+When comparing Radix vs Octree vs KD-tree traversal outputs, use the same MAC settings as
+your downstream solver.
+
+- For FMM-style runs (e.g. jaccpot), `mac_type="dehnen"` is the recommended
+  path for apples-to-apples parity checks.
+- Octree builds currently share the radix traversal core, so interaction-count
+  parity between `radix` and `octree` should hold for the same build settings.
+- KD-tree traversal uses a calibrated default effective radius scale for
+  Dehnen MAC (`dehnen_radius_scale=1.2`) to match near-field/far-field split
+  behavior more closely with radix trees.
+- If you benchmark with `mac_type="bh"`, expect different KD/Radix split
+  behavior unless you tune parameters explicitly.
+
+## Backend Extensibility
+
+Yggdrax now supports backend-oriented tree dispatch and capability-based
+topology contracts:
+
+- Register builders via `register_tree_builder(...)`
+- Inspect available builders via `available_tree_types()`
+- Use `resolve_tree_topology(...)` for container/topology adapters
+- Use derivation helpers (`get_node_levels`, `get_level_offsets`,
+  `get_nodes_by_level`) when a backend does not precompute level metadata
+- Octree consumers can additionally use explicit buffers like `oct_children`
+  and `oct_level_offsets` when level-wise FMM scheduling is preferable to
+  binary traversal over `left_child` / `right_child`
+
+Contract details and required/optional fields are documented in
+`docs/backend_contract.md`.
+
+## Build And Traversal Configs
+
+Public config dataclasses provide a stable way to reuse tuned settings across
+repeated builds and traversals:
+
+- `TreeBuildConfig`: adaptive radix-tree settings (`leaf_size`,
+  `return_reordered`, reusable workspace handling)
+- `FixedDepthTreeBuildConfig`: fixed-depth tree settings, including local
+  Morton refinement controls
+- `DualTreeTraversalConfig`: traversal queue, block size, interaction capacity,
+  and neighbor capacity
+
+When a `config=...` object is passed to `build_tree(...)` or
+`build_fixed_depth_tree(...)`, or their octree counterparts, it takes precedence over the equivalent
+individual keyword arguments.
+
+Conformance tests:
+
+```bash
+pytest -q --no-cov tests/unit/test_backend_conformance.py
+```
+
+## Development
+
+Local quality gates:
+
+```bash
+pytest
+black --check .
+isort --check-only .
+pydoclint .
+```
+
+Or run the same checks via pre-commit:
+
+```bash
+pre-commit run --all-files
+```
+
+Coverage is enforced via `pytest-cov`:
+
+```bash
+pytest --cov=yggdrax --cov-report=term-missing
+```
+
+## Project Structure
+
+- `yggdrax/tree.py`, `yggdrax/_tree_impl.py`: tree building and radix internals
+- `yggdrax/octree.py`: explicit octree metadata derived from Morton/radix trees
+- `yggdrax/protocols.py`: backend capability protocols
+- `yggdrax/geometry.py`, `yggdrax/_geometry_impl.py`: geometry wrappers and implementations
+- `yggdrax/interactions.py`, `yggdrax/_interactions_impl.py`: traversal and interaction generation
+- `yggdrax/dense_interactions.py`, `yggdrax/grouped_interactions.py`: interaction layout utilities
+- `tests/unit`: unit test suite for API and implementation behavior
+- `examples`: runnable examples and notebooks
+
+## CI
+
+GitHub Actions runs:
+
+- unit tests with coverage threshold
+- `black --check`
+- `isort --check-only`
+- `pydoclint`
+
+Workflow file: `.github/workflows/ci.yml`.
+
+## Relationship to Rubix
+
+This repository follows the same engineering principles used in the Rubix codebase:
+
+- strict formatting and lint automation
+- tested public APIs
+- explicit artifact contracts
+- examples that reflect real usage paths

yggdrax-0.0.1/README.md

@@ -0,0 +1,218 @@
+# Yggdrax
+
+![Black](https://img.shields.io/badge/code%20style-black-000000.svg)
+![isort](https://img.shields.io/badge/imports-isort-1674b1.svg)
+![pydoclint](https://img.shields.io/badge/docstrings-pydoclint-2ea44f.svg)
+
+![Yggdrax Logo](./yggdrax.png)
+
+Yggdrax is a JAX-first tree toolkit for hierarchical N-body solvers. It
+provides Morton ordering, radix tree builders, per-node geometry, and dual-tree
+interaction traversal primitives designed for downstream FMM and treecode
+pipelines. The public `octree` backend now layers explicit octree-cell metadata
+on top of the existing Morton/radix construction path so downstream FMM code
+can consume both the proven traversal buffers and octree-style child tables.
+
+## Features
+
+- Morton encode/decode and stable Morton sorting for 3D points
+- LBVH and fixed-depth radix tree construction
+- Explicit octree metadata derived from Morton/radix topology
+- Tree geometry extraction (bounds, centers, extents, radii)
+- Dual-tree far-field and near-field interaction builders
+- Dense and grouped interaction buffer transforms for batched kernels
+- Prepared artifact utilities for downstream solver integrations
+
+## Installation
+
+Install from source:
+
+```bash
+pip install -e .
+```
+
+Install with development tools:
+
+```bash
+pip install -e ".[dev]"
+```
+
+## Quick Start
+
+```python
+import jax
+import jax.numpy as jnp
+
+from yggdrax import (
+    DualTreeTraversalConfig,
+    build_interactions_and_neighbors,
+    build_octree,
+    compute_tree_geometry,
+)
+
+key = jax.random.PRNGKey(0)
+key_pos, key_mass = jax.random.split(key)
+positions = jax.random.uniform(key_pos, (512, 3), minval=-1.0, maxval=1.0)
+masses = jax.random.uniform(key_mass, (512,), minval=0.5, maxval=1.5)
+
+tree = build_octree(positions, masses, leaf_size=16)
+positions_sorted = positions[tree.particle_indices]
+geom = compute_tree_geometry(tree, positions_sorted)
+traversal_cfg = DualTreeTraversalConfig(
+    max_pair_queue=8192,
+    process_block=256,
+    max_interactions_per_node=2048,
+    max_neighbors_per_leaf=2048,
+)
+interactions, neighbors = build_interactions_and_neighbors(
+    tree,
+    geom,
+    theta=0.6,
+    mac_type="dehnen",
+    traversal_config=traversal_cfg,
+)
+```
+
+`build_tree(...)` continues to expose the radix/LBVH backend directly. The
+octree wrappers (`build_octree(...)`, `build_fixed_depth_octree(...)`) preserve
+the same compatibility fields while additionally exposing explicit octree
+buffers such as `oct_children`, `oct_node_depths`, and `radix_node_to_oct`.
+
+Advanced users can override the built-in MAC with a JAX-traceable pair policy:
+
+```python
+def pair_policy(policy_state, **pair_data):
+    action = ...
+    tag = ...
+    return action, tag
+
+interactions, neighbors, result = build_interactions_and_neighbors(
+    tree,
+    geom,
+    pair_policy=pair_policy,
+    policy_state=...,
+    return_result=True,
+)
+```
+
+The policy receives generic pair geometry/state and returns:
+- `action`: one of accept-far / accept-near / refine
+- `tag`: integer metadata stored for accepted far pairs
+
+When `return_result=True`, raw far-pair tags are available on
+`result.interaction_tags`. This is intended for downstream solvers that need
+solver-side scheduling or adaptive-order bucketing without moving solver logic
+into `yggdrax`.
+
+See `examples/getting_started.ipynb` for a runnable walkthrough.
+For the locked high-performance GPU benchmark configuration, see
+`docs/gpu_benchmark_recommended_setup.md` and
+`examples/tree_gpu_performance_scaling.ipynb`.
+
+## KD-Tree MAC Note
+
+When comparing Radix vs Octree vs KD-tree traversal outputs, use the same MAC settings as
+your downstream solver.
+
+- For FMM-style runs (e.g. jaccpot), `mac_type="dehnen"` is the recommended
+  path for apples-to-apples parity checks.
+- Octree builds currently share the radix traversal core, so interaction-count
+  parity between `radix` and `octree` should hold for the same build settings.
+- KD-tree traversal uses a calibrated default effective radius scale for
+  Dehnen MAC (`dehnen_radius_scale=1.2`) to match near-field/far-field split
+  behavior more closely with radix trees.
+- If you benchmark with `mac_type="bh"`, expect different KD/Radix split
+  behavior unless you tune parameters explicitly.
+
+## Backend Extensibility
+
+Yggdrax now supports backend-oriented tree dispatch and capability-based
+topology contracts:
+
+- Register builders via `register_tree_builder(...)`
+- Inspect available builders via `available_tree_types()`
+- Use `resolve_tree_topology(...)` for container/topology adapters
+- Use derivation helpers (`get_node_levels`, `get_level_offsets`,
+  `get_nodes_by_level`) when a backend does not precompute level metadata
+- Octree consumers can additionally use explicit buffers like `oct_children`
+  and `oct_level_offsets` when level-wise FMM scheduling is preferable to
+  binary traversal over `left_child` / `right_child`
+
+Contract details and required/optional fields are documented in
+`docs/backend_contract.md`.
+
+## Build And Traversal Configs
+
+Public config dataclasses provide a stable way to reuse tuned settings across
+repeated builds and traversals:
+
+- `TreeBuildConfig`: adaptive radix-tree settings (`leaf_size`,
+  `return_reordered`, reusable workspace handling)
+- `FixedDepthTreeBuildConfig`: fixed-depth tree settings, including local
+  Morton refinement controls
+- `DualTreeTraversalConfig`: traversal queue, block size, interaction capacity,
+  and neighbor capacity
+
+When a `config=...` object is passed to `build_tree(...)` or
+`build_fixed_depth_tree(...)`, or their octree counterparts, it takes precedence over the equivalent
+individual keyword arguments.
+
+Conformance tests:
+
+```bash
+pytest -q --no-cov tests/unit/test_backend_conformance.py
+```
+
+## Development
+
+Local quality gates:
+
+```bash
+pytest
+black --check .
+isort --check-only .
+pydoclint .
+```
+
+Or run the same checks via pre-commit:
+
+```bash
+pre-commit run --all-files
+```
+
+Coverage is enforced via `pytest-cov`:
+
+```bash
+pytest --cov=yggdrax --cov-report=term-missing
+```
+
+## Project Structure
+
+- `yggdrax/tree.py`, `yggdrax/_tree_impl.py`: tree building and radix internals
+- `yggdrax/octree.py`: explicit octree metadata derived from Morton/radix trees
+- `yggdrax/protocols.py`: backend capability protocols
+- `yggdrax/geometry.py`, `yggdrax/_geometry_impl.py`: geometry wrappers and implementations
+- `yggdrax/interactions.py`, `yggdrax/_interactions_impl.py`: traversal and interaction generation
+- `yggdrax/dense_interactions.py`, `yggdrax/grouped_interactions.py`: interaction layout utilities
+- `tests/unit`: unit test suite for API and implementation behavior
+- `examples`: runnable examples and notebooks
+
+## CI
+
+GitHub Actions runs:
+
+- unit tests with coverage threshold
+- `black --check`
+- `isort --check-only`
+- `pydoclint`
+
+Workflow file: `.github/workflows/ci.yml`.
+
+## Relationship to Rubix
+
+This repository follows the same engineering principles used in the Rubix codebase:
+
+- strict formatting and lint automation
+- tested public APIs
+- explicit artifact contracts
+- examples that reflect real usage paths

yggdrax-0.0.1/pyproject.toml

@@ -0,0 +1,89 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "yggdrax"
+version = "0.0.1"
+description = "Tree generation and traversal package for hierarchical solvers"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+dependencies = [
+  "jax>=0.8.0",
+  "jaxlib>=0.4.0",
+  "jaxtyping>=0.2.23",
+  "beartype>=0.14.0",
+]
+authors = [{ name = "AstroAI Lab" }]
+keywords = ["fmm", "hierarchical", "n-body", "jax", "treecode"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Scientific/Engineering :: Physics",
+]
+
+[project.urls]
+Homepage = "https://github.com/AstroAI-Lab/yggdrax"
+Repository = "https://github.com/AstroAI-Lab/yggdrax"
+Issues = "https://github.com/AstroAI-Lab/yggdrax/issues"
+
+[project.optional-dependencies]
+dev = [
+  "black>=24.8.0",
+  "isort>=5.13.2",
+  "pydoclint>=0.6.0",
+  "pytest>=8.3.0",
+  "pytest-cov>=5.0.0",
+]
+
+[tool.setuptools]
+packages = ["yggdrax"]
+
+[tool.setuptools.package-dir]
+yggdrax = "yggdrax"
+
+[tool.setuptools.package-data]
+yggdrax = ["py.typed"]
+
+[tool.black]
+line-length = 88
+target-version = ["py310", "py311", "py312"]
+include = "\\.pyi?$"
+exclude = '''
+/(
+  \.git
+ |\.venv
+ |build
+ |dist
+)/
+'''
+
+[tool.isort]
+profile = "black"
+line_length = 88
+skip_gitignore = true
+
+[tool.pytest.ini_options]
+minversion = "7.4"
+testpaths = ["tests"]
+addopts = "-q --strict-markers --cov=yggdrax --cov-report=term-missing --cov-report=xml --cov-fail-under=85"
+
+[tool.coverage.run]
+source = ["yggdrax"]
+branch = true
+
+[tool.coverage.report]
+show_missing = true
+skip_empty = true
+fail_under = 85
+exclude_lines = [
+  "pragma: no cover",
+  "if TYPE_CHECKING:",
+  "if __name__ == .__main__.:",
+]

yggdrax-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+