Metadata-Version: 2.4
Name: hyperiax
Version: 3.0.0
Summary: Tree traversals using JAX
Author-email: Stefan Sommer <sommer@di.ku.dk>, Marcus Holte Teller <marcus@di.ku.dk>, Gefan Yang <gy@di.ku.dk>, Michael Baand Severinsen <michael.baand@sund.ku.dk>, Morten Akhøj Pedersen <morten.akhoej@di.ku.dk>, Nicklas Boserup <nicklas.boserup@sund.ku.dk>
Project-URL: Homepage, https://github.com/computationalevolutionarymorphometry/hyperiax
Project-URL: Issues, https://github.com/computationalevolutionarymorphometry/hyperiax/issues
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4.30
Requires-Dist: jaxlib>=0.4.30
Requires-Dist: numpy>=1.26
Provides-Extra: gpu
Requires-Dist: jax[cuda12]>=0.4.30; extra == "gpu"
Provides-Extra: notebook
Requires-Dist: matplotlib>=3.8; extra == "notebook"
Requires-Dist: ipykernel>=6.29; extra == "notebook"
Requires-Dist: numpyro>=0.15; extra == "notebook"
Provides-Extra: numpyro
Requires-Dist: numpyro>=0.15; extra == "numpyro"
Dynamic: license-file

# Hyperiax

<p align="center">
  <img width="300" height="250" src="./docs/figures/hyperiax_logo.png" alt="Hyperiax logo">
</p>

Hyperiax is a pure-JAX library for message passing on rooted trees. It provides
immutable tree data structures, typed per-node arrays, and decorator-based
sweeps that compose with `jax.jit`, `jax.vmap`, and `jax.lax.scan`.

Hyperiax is designed for phylogenetic and tree-structured scientific computing,
including Gaussian graphical models, phylogenetic means, and guided inference
for diffusion processes on trees.

## Installation

Hyperiax requires Python 3.11 or newer.

```bash
pip install hyperiax
```

The core package depends only on JAX, jaxlib, and NumPy. For accelerator-backed
JAX installations, follow the official JAX installation instructions for your
platform before or after installing Hyperiax.

## Quick Start

```python
import jax.numpy as jnp
import hyperiax as hx

topology = hx.symmetric_topology(depth=2, degree=2)
tree = hx.Tree.empty(topology, {"value": (2,)})

leaf_count = int(topology.is_leaf.sum())
tree = tree.at[topology.is_leaf].set(value=jnp.ones((leaf_count, 2)))


@hx.up(reads_children=("value",), writes=("value",))
def average_children(node, children, params):
    return {"value": children.value.mean(0)}


result = average_children(tree)
root_value = result.value[0]
```

Sweeps are ordinary `Tree -> Tree` functions. The `@hx.up` and `@hx.down`
decorators declare which fields are read and written, so dispatch remains
explicit and JAX-friendly.

## Architecture

```text
hyperiax/
├── core/      # Topology, Tree, Schema, views, sweep dispatch, builders
├── utils/     # Pure-JAX ODE and SDE solvers
└── prebuilt/  # Ready-to-use sweeps for BFFG and phylogenetic means
```

### Core

`hyperiax.core` contains the public primitives:

- `Topology` for rooted tree structure.
- `Tree` for immutable per-node JAX arrays.
- `Schema` for field shape and dtype validation.
- `@hx.up` and `@hx.down` for message-passing sweeps.
- Newick helpers for importing and exporting rooted trees.

### Utilities

`hyperiax.utils` contains pure-JAX numerical helpers, including ODE and SDE
solvers used by the prebuilt guided-inference routines.

### Prebuilt Sweeps

`hyperiax.prebuilt` contains focused implementations for common tree workflows:

- `phylo_mean` for weighted phylogenetic means.
- `bffg` for Backward Filtering Forward Guiding on discrete Gaussian and
  continuous SDE edges.

## Documentation

Tutorials and API reference are available in the project documentation:

- Quickstart and sweep-writing tutorials under `docs/source/notebooks/`.
- Public API reference under `docs/source/api/`.

## Reference

The BFFG implementation follows:

> van der Meulen, F. H. & Sommer, S. (2025). Backward Filtering Forward
> Guiding. JMLR 26(281), 1-51. https://arxiv.org/abs/2505.18239

## Contact

Hyperiax is maintained by CCEM, University of Copenhagen. For technical
questions, please open a GitHub issue or contact
[Stefan Sommer](mailto:sommer@di.ku.dk).
