pyproject.toml

[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "jaccpot"
version = "0.0.1"
description = "JAX Fast Multipole Method solver package built on yggdrax tree artifacts"
readme = "README.md"
requires-python = ">=3.11"
license = "MIT"
authors = [
  { name = "AstroAI Lab", email = "astroai@iwr.uni-heidelberg.de" },
  { name = "Tobias Buck", email = "tobias.buck@iwr.uni-heidelberg.de" },
]
maintainers = [{ name = "AstroAI Lab", email = "astroai@iwr.uni-heidelberg.de" }]
keywords = ["fmm", "n-body", "jax", "gravity", "treecode"]
classifiers = [
  "Development Status :: 3 - Alpha",
  "Intended Audience :: Science/Research",
  "Programming Language :: Python :: 3",
  "Programming Language :: Python :: 3.11",
  "Programming Language :: Python :: 3.12",
  "Programming Language :: Python :: 3.13",
  "Topic :: Scientific/Engineering :: Physics",
]
dependencies = [
  # JAX FLOOR IS 0.10.2, and it is a hard floor: 0.9.1 and 0.9.2 CRASH on the CPU
  # backend. `tests/characterization/test_fmm_golden.py::test_fmm_golden[uni_real_n256_p6]`
  # dies with `Fatal Python error: Floating-point exception` -- a SIGFPE inside
  # XLA's CPU compiler while compiling the real-basis P2M `lax.scan`
  # (`upward/real_tree_expansions.py::_p2m_leaves_real`). Not ours to fix and not
  # catchable: it kills the worker. Measured on the CPU backend:
  #
  #     jax 0.9.1   core dump        <- was briefly the floor here; CI red
  #     jax 0.9.2   core dump        <- what a `<0.10` ceiling resolved to
  #     jax 0.10.2  passes (23.1 s)
  #     jax 0.11.0  passes (23.5 s)
  #
  # It is not Python-version specific (reproduced on 3.12 locally and 3.13 in CI).
  # This is why CI went red on this branch: bumping the floor to 0.9.1 made the
  # resolver pick 0.9.2, straight into the broken window. GPU-only validation does
  # not see it -- the CPU backend must be exercised for any JAX bump.
  #
  # Also 0.9.1+ is what the differentiable multi-GPU path wants: 0.9.1 is the first
  # release whose `jax.lax.ragged_all_to_all` survives having its reverse pass
  # executed (XLA fixed the stale peer-address rendezvous in `4e0cc7e356`). Below
  # it, one gradient silently breaks every later halo exchange and a forward loses
  # the whole cross-domain near field. `halo_exchange="auto"` still falls back
  # safely below the floor, so that part is policy, not requirement; the CPU crash
  # is the requirement. See docs/differentiable_fmm_distributed_audit.md.
  #
  # WHY 0.10.2 AND NOT 0.11.0. 0.11.0 also passes everywhere, but two things make
  # it the wrong floor:
  #   * it requires Python >= 3.12, while this package supports >= 3.11 -- so a
  #     0.11.0 floor silently drops 3.11 and fails at INSTALL on the CI matrix's
  #     3.11 legs (test-smoke, benchmark-guard);
  #   * it is ~2.6x slower than 0.10.2 on the CPU backend for this workload
  #     (characterization suite: 109.8 s on 0.10.2 and on 0.9.0.1, 289.5 s on
  #     0.11.0), which pushed CI's test-full from ~33 min past a 60-min cap.
  # 0.10.2 is as fast as the 0.9.0.x the project was developed against, keeps 3.11,
  # and has the ragged fix.
  #
  # Validated on BOTH backends. 0.10.2: GPU (2xA100) Pallas M2L + custom_vjp +
  # precision guards 54 passed, distributed gradients 10 passed; CPU golden
  # uni_real_n256_p6 passes. 0.11.0: GPU broad 506 / Pallas 51 / gradients 20 /
  # golden 13, CPU golden passes -- so 0.11.0 is *usable*, just not the floor.
  #
  # THE `<0.11` CEILING IS LOAD-BEARING, not caution. A floor alone does not stop
  # the resolver taking 0.11.0, and CI proved it in one run: with `>=0.10.2` and no
  # ceiling, the py3.11 leg got 0.10.2 and PASSED in 22.9 min, while the py3.12 and
  # py3.13 legs got 0.11.0 and blew 30- and 50-minute caps. Same code, same runner,
  # one variable -- which also confirms the 2.6x CPU slowdown in situ. py3.11 was
  # only green because 0.11.0 refuses to install there.
  #
  # So: lifting this ceiling means either accepting ~2.6x CPU CI time or first
  # finding out why 0.11.0 is that much slower for this workload. Note also that
  # `pip install yggdrax` pulls JAX unconstrained BEFORE this constraint applies, so
  # the ceiling is what downgrades it back.
  #
  # HISTORY: the ceiling used to be `<0.9.1` because `pallas_call`'s `backend=`
  # kwarg was removed in 0.9.1 and the fused M2L kernels passed it. They now go
  # through `jaccpot/pallas/_compat.pallas_backend_kwargs`, which selects Triton
  # via `backend=` on the old API and via `triton.CompilerParams` on the new one.
  "jax[cuda]>=0.10.2,<0.11",
  "jaxlib>=0.10.2,<0.11",
  "jaxtyping>=0.2.23",
  "beartype>=0.14.0",
  "yggdrax>=0.0.1,<0.1.0",
]

[project.urls]
Homepage = "https://github.com/AstroAI-Lab/jaccpot"
Repository = "https://github.com/AstroAI-Lab/jaccpot"
Issues = "https://github.com/AstroAI-Lab/jaccpot/issues"
Documentation = "https://github.com/AstroAI-Lab/jaccpot#readme"

[project.optional-dependencies]
dev = [
  # Pinned to match .pre-commit-config.yaml's black-pre-commit-mirror rev so a
  # local `black` and the hook/CI format identically (a floating ">=" let CI
  # resolve to a different black whose style disagreed with the hook).
  "black==26.5.1",
  # Pinned to match .pre-commit-config.yaml. isort 7.0.0/8.0.1 regressed the
  # black profile (force-wraps single-name imports); 9.0.0b1 fixes it. A
  # floating ">=" here lets CI's `isort --check-only .` resolve to the broken
  # 8.0.1 and disagree with pre-commit, so keep this equal to the hook rev.
  "isort==9.0.0b1",
  # Pinned to match .pre-commit-config.yaml's pydoclint rev, for the same reason as
  # black and isort above: the hook runs on every commit now, and a floating ">="
  # would let a local pydoclint disagree with it about what counts as a violation.
  "pydoclint==0.9.1",
  # Pinned, but NOT for the reason black/isort/pydoclint are: pyright is not a
  # pre-commit hook and not in CI, so there is nothing for it to disagree with.
  # It is pinned because the error counts the audit's E.2 records (705 on the
  # `self:`-annotated tree, 408 after the `_EngineBase` change) are specific to
  # this version, and with no gate to catch drift an unpinned checker would make
  # them quietly unreproducible. Bump it deliberately and re-record the numbers.
  "pyright==1.1.411",
  "pytest>=8.3.0",
  "pytest-cov>=5.0.0",
  "pytest-xdist>=3.6.0",
  "pre-commit>=3.8.0",
  "build>=1.2.2",
  "twine>=5.1.1",
]

[tool.setuptools.packages.find]
where = ["."]
include = ["jaccpot*"]

[tool.setuptools.package-data]
jaccpot = ["py.typed"]

[tool.black]
line-length = 88
target-version = ["py311", "py312", "py313"]
include = "\\.pyi?$"
exclude = '''
/(
  \.git
 |\.venv
 |build
 |dist
)/
'''

[tool.pydoclint]
style = "numpy"
# `.claude` holds git worktrees -- full transient copies of this repo. Without
# excluding them a manual run over `.` reports every violation four times over
# (measured against the old baseline: 409 of 521 file entries, ~13k of 14.7k
# violations) and the result churns as worktrees come and go.
exclude = '\.git|\.venv|\.claude|build|dist|jaccpot/experimental'

# The upstream default is `true`, and that made this hook VACUOUS for the whole
# Tier 2.5 docstring programme: it skips argument and return checking on any
# docstring that is a bare one-line summary, which is exactly the shape every
# undocumented function has. Measured 2026-08-17 on `main`: 0 violations under
# the default, 560 with it off. A function with a one-line summary and eight
# undocumented parameters passed CI.
#
# It also let finished work come undone. `runtime/_adaptive_policy.py` was
# driven to zero (PR #118), and `main` was never at zero: five functions from a
# concurrent branch merged into the same file without conflict, and the hook
# reported 0 the whole time (repaired in PR #149).
skip-checking-short-docstrings = false

# There is deliberately NO baseline any more. `.pydoclint-baseline.txt` carried
# the pre-existing tail while the Tier 2.5 docstring programme ran; that
# programme finished, `jaccpot/` measured 0 violations with the baseline out of
# play (2026-08-18), and a baseline with nothing real left to suppress is worse
# than none: it cannot hide a violation that exists, but it can silently absorb
# a NEW one in any file that still has an entry.
#
# Retiring it moved the enforcement scope into `.pre-commit-config.yaml`
# instead. The baseline's last 1215 entries were all outside `jaccpot/`
# (tests/ 1034, bench/ 121, examples/ 60 -- exactly the live count there), so
# deleting it repo-wide would have turned the hook red without catching
# anything it was not already suppressing. The hook is now scoped `^jaccpot/`,
# which leaves enforcement on the package unchanged and unconditional, and
# leaves tests/bench/examples exactly as unchecked as the baseline had them.

[tool.pyright]
# Scoped to the package, mirroring the pydoclint hook's `files: ^jaccpot/`: the
# same backlog argument applies -- `tests/`, `bench/` and `examples/` have never
# been checked and would swamp the signal.
include = ["jaccpot"]
# `experimental/` is opt-in prototype code, excluded here for the same reason
# `[tool.coverage.run].omit` excludes it and the `experimental` marker deselects
# its tests. Listed as the package, not a glob -- see the note on `omit` below.
exclude = ["jaccpot/experimental"]
# Resolve the sibling `yggdrax` checkout. Relative to this file, so it is the
# documented layout (ARCHITECTURE: "Sibling codebases ... yggdrax") rather
# than one machine's paths, and it is the same assumption
# `tests/conftest.py::_find_sibling_checkout` already makes.
extraPaths = ["../yggdrax"]
pythonVersion = "3.11"
# Set explicitly rather than left to the default, so a pyright release that moves
# its default mode cannot silently change what this reports.
typeCheckingMode = "standard"
#
# `reportMissingImports` is left ON deliberately, and `extraPaths` above is what
# makes that affordable. An earlier version of this comment said the rule had to
# stay noisy because `yggdrax` is an editable install pyright cannot follow --
# true, but the conclusion was wrong: those were not 100 cosmetic errors to
# subtract, they were 100 *suppressions*. An unresolved import makes every symbol
# from it `Unknown`, and `Unknown` is assignable to anything, so every call across
# the jaccpot/yggdrax boundary went unchecked. Resolving it removed 125 errors and
# switched 44 real checks back on. See audit E.5.
#
# There is no gate. `pyright jaccpot/` reports 210 errors and a hook at that count
# would need a baseline -- exactly what item 2.5 retired.

[tool.isort]
profile = "black"
line_length = 88
skip_gitignore = true

[tool.pytest.ini_options]
minversion = "7.4"
testpaths = ["tests"]
markers = [
    "slow: long-running / compile-heavy test (deselect with -m 'not slow')",
    "experimental: exercises experimental, non-production FMM code (octree/treecode prototypes); deselected by default, run explicitly with -m experimental",
]
# `-n auto` runs the (compilation-bound) JAX test suite across all logical cores
# via pytest-xdist. On a 72-core host this cuts wall-clock from ~25 min to ~2 min;
# a fixed `-n 32` measured marginally faster there. Override with `-n 0` to run
# serially (e.g. for pdb/`-x` debugging, which xdist does not support).
# `-m "not experimental"` keeps the everyday run on the production (radix real
# fast-lane) paths; the long-compile octree/treecode prototypes are opt-in via
# `pytest -m experimental`.
addopts = "-q --strict-markers --ignore=tests/perf -n auto -m \"not experimental\""

[tool.coverage.run]
source = ["jaccpot"]
# Kept ON. Turning it off was tried as a fix for the `test-full` OOM and is not
# worth it: measured on `tests/integration` with `-n 2 --cov`, peak RSS was
# 12,686 MB with branch coverage and 12,677 MB without -- a 9 MB (0.07%)
# difference, i.e. noise. The job's memory is JAX's retained compiled
# executables, not coverage's arc table, so this would have cost branch data for
# nothing. See `tests/conftest.py::_DIFF_FMM_TEST_FILES` for the fix that works.
#
# (Note for anyone retrying it: `--cov-branch` on the command line cannot turn
# this off -- setting it True here enables branch coverage regardless of the
# flag, so it has to change here.)
branch = true
omit = [
  # Opt-in experimental prototypes, listed INDIVIDUALLY rather than as
  # `jaccpot/experimental/*`, because that glob was hiding a module the default
  # suite fully covers. Measured:
  #
  #   octree_fmm_uvwx.py    0%    reached only under `pytest -m experimental`
  #   treecode_far_near.py  0%    reached only from the distributed treecode lane
  #                               (`local_walk="treecode"`), whose test skips
  #                               below 2 devices, so CPU CI never enters it
  #   treecode_walk.py    100%    NOT omitted -- see below
  #
  # `treecode_walk.py` is the pure-JAX reference that
  # `tests/unit/operators/test_pallas_treecode_walk.py` checks the Pallas walk
  # against. That test is neither `experimental`- nor `slow`-marked, so it runs on
  # every commit and covers the module 62/62 statements. Omitting it understated
  # coverage and, worse, the stated reason for omitting it was untrue.
  #
  # Keeping this an explicit list means a new file under `jaccpot/experimental/`
  # is measured until someone decides otherwise -- which is the right default,
  # since the failure this replaces was a glob quietly absorbing covered code.
  "jaccpot/experimental/octree_fmm_uvwx.py",
  "jaccpot/experimental/treecode_far_near.py",
  # `jaccpot/pallas/m2l_real_fused.py` used to be omitted here too, on the stated
  # grounds that it "has no interpret path, so CPU CI structurally cannot reach
  # it". That was untrue, and it is the same failure this list already documents
  # one entry above: an omit hiding a well-covered module behind a wrong reason.
  # Measured on CPU, from `tests/unit/operators/test_m2l_real_fused_pallas.py` and
  # `tests/unit/test_custom_vjp_parity.py` alone: 191 statements, 12 missed, 92%.
  # Both suites call `m2l_real_fused_pallas(..., interpret=True)` and say so in
  # their module docstrings. Removed, so the module is measured like any other.
]

[tool.coverage.report]
show_missing = true
skip_empty = true
exclude_lines = [
  "pragma: no cover",
  "if TYPE_CHECKING:",
  "if __name__ == .__main__.:",
]