jaxht 0.1.0__tar.gz

jaxht-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,52 @@
+name: release
+
+# Build and publish to PyPI on a version tag (e.g. `git tag v0.1.0 && git push --tags`).
+# Uses PyPI Trusted Publishing (OIDC) -- no API token / secret.
+#
+# ONE-TIME PyPI SETUP before the first tag (the project need not exist yet -- use a
+# "pending publisher"): on https://pypi.org -> your account -> Publishing, add a
+# trusted publisher with:
+#     PyPI project name: jaxht        (distribution name; the import name stays `jht`)
+#     Owner:             jrcheshire
+#     Repository name:   jht
+#     Workflow name:     release.yml
+#     Environment name:  pypi
+# Keep __version__ in src/jht/__init__.py in sync with the tag (the build reads the
+# version from there via hatch).
+
+on:
+  push:
+    tags: ["v*"]
+
+jobs:
+  build:
+    name: Build sdist + wheel
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: python -m pip install --upgrade build twine
+      - run: python -m build
+      - run: python -m twine check dist/*
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    name: Publish to PyPI (Trusted Publishing)
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/project/jaxht/
+    permissions:
+      id-token: write # OIDC for Trusted Publishing -- no token needed
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

jaxht-0.1.0/.github/workflows/test.yml

@@ -0,0 +1,36 @@
+name: test
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+jobs:
+  lint-and-test:
+    name: Lint + typecheck + tests (linux)
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: prefix-dev/setup-pixi@v0.8.10
+        with:
+          # `locked: true` runs `pixi install --locked`, so any pyproject.toml
+          # drift from pixi.lock surfaces as a CI failure (the co-commit invariant).
+          locked: true
+          cache: true
+
+      - name: Lint
+        run: pixi run lint
+
+      - name: Typecheck
+        run: pixi run typecheck
+
+      # The default pixi env pulls the ducc0/healpy oracles, so this is the full
+      # gated parity suite (CPU, float64). The GPU parity gate auto-skips (no device).
+      - name: Test
+        run: pixi run test

jaxht-0.1.0/.gitignore

@@ -0,0 +1,25 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+build/
+dist/
+
+# Pixi / envs
+.pixi/
+.venv/
+
+# Tool caches
+.mypy_cache/
+.ruff_cache/
+.pytest_cache/
+
+# Scratch / run outputs (not fixtures)
+runs/
+scratch/
+gpu_diag_results.jsonl
+
+# OS / editor
+.DS_Store
+*.swp

jaxht-0.1.0/CHANGELOG.md

@@ -0,0 +1,34 @@
+# Changelog
+
+All notable changes to jht are documented here. The format is based on
+[Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project
+adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+## [0.1.0] - 2026-06-07
+
+Inaugural release.
+
+### Added
+- **On-grid transforms** — spin-0 and spin-2 `synthesis` (a_lm → map) and the
+  exact adjoint `adjoint_synthesis`, validated to machine precision against
+  healpy **and** ducc0 (spin-2 inverse at the HEALPix floor, no s2fft-style
+  structural defect).
+- **Approximate inverse** — `bare_analysis` (Sᵀ W) and `map2alm` (jht's own ring
+  quadrature weights + Jacobi iteration; ~1e-13 on band-limited maps).
+- **Quadrature weights** — `ring_weights`, `pixel_weights`.
+- **Partial-sky / masked** — `pseudo_alm`, `deconvolve` (cut-sky CG), `wiener`
+  (masked Wiener filter / MUSE inner solve), `constrained_realization`.
+- **Off-grid (NUFFT)** — `synthesis_general` / `adjoint_synthesis_general` for
+  spin 0–3 at arbitrary pointings; alm- **and** pointing-differentiable under
+  native autodiff. JAX-native replacement for ducc0's `synthesis_general`.
+- **Differentiable real-DOF interface** — `synthesis_real`, `analysis_real`,
+  `bandpower`, with the `alm_to_real` / `real_to_alm` isometry and the
+  `alm_metric_weight` (2 − δ_m0) bridge.
+- **GPU (CUDA)** — pure JAX, runs on GPU with no code change; a parity harness
+  (`scripts/gpu_check.py`) and a single-slot performance diagnostic
+  (`scripts/gpu_diagnostic.py`).
+
+[Unreleased]: https://github.com/jrcheshire/jht/compare/v0.1.0...HEAD
+[0.1.0]: https://github.com/jrcheshire/jht/releases/tag/v0.1.0

jaxht-0.1.0/CITATION.cff

@@ -0,0 +1,27 @@
+cff-version: 1.2.0
+message: "If you use jht in your research, please cite it."
+title: "jht: JAX-native spherical harmonic transforms"
+abstract: >-
+  A clean-room, pure-JAX implementation of the forward and inverse spherical
+  harmonic transform (map <-> a_lm) for spin-0 and spin-2 fields on the HEALPix
+  ring pixelization, plus an off-grid (NUFFT) path for arbitrary pointings.
+  GPU-capable, fully differentiable under JAX autodiff, and dependency-controlled
+  (JAX + numpy at runtime).
+type: software
+authors:
+  - given-names: James
+    family-names: Cheshire
+    email: cheshire@caltech.edu
+    affiliation: California Institute of Technology
+version: 0.1.0
+date-released: 2026-06-07
+license: MIT
+repository-code: "https://github.com/jrcheshire/jht"
+keywords:
+  - spherical harmonic transform
+  - HEALPix
+  - JAX
+  - GPU
+  - automatic differentiation
+  - NUFFT
+  - cosmic microwave background

jaxht-0.1.0/CLAUDE.md

@@ -0,0 +1,195 @@
+# jht — JAX Harmonic Transforms
+
+JAX-native spherical harmonic transforms. GPU-capable, fully differentiable,
+and **dependency-controlled** — pure JAX + numpy at runtime, no compiled C++
+extension and no heavyweight third-party SHT library. Born from the BICEP/Keck
+`bk-jax` pipeline's need for a transform it can run on GPU, differentiate
+through, and *own* the numerics of.
+
+**Status (2026-06-06): scaffolding only.** This file, `ROADMAP.md`, and the
+`docs/` notes were written to carry the design into a future development
+session. No transforms are implemented yet. **Read `ROADMAP.md` next; the first
+action is the Phase-0 feasibility spike.**
+
+---
+
+## What jht is, in one paragraph
+
+A focused, clean-room implementation of the forward and inverse spherical
+harmonic transform (map ↔ aₗₘ) for **spin-0 and spin-2** fields on the
+**HEALPix ring** pixelization, in the **BICEP/Keck angular regime**
+(ℓ_max ≲ 1000, nside ≤ ~2048). It exists to serve the GPU / differentiable
+tier of analysis that a CPU-only C++ transform structurally cannot. It is
+*not* an attempt to reimplement a general all-spin / all-sampling library.
+
+## Scope
+
+**In scope (the target):**
+- Spin-0 and spin-2 (T and Q/U / E/B) transforms, forward and adjoint.
+- HEALPix ring pixelization.
+- BK regime: ℓ_max ≲ 1000, nside ≤ ~2048; partial-sky (masked) analysis.
+- GPU execution (the primary motivation) and JAX autodiff (alm-linear).
+- Validation to the relevant accuracy tier against ducc0 / healpy.
+
+**Out of scope (for now — keep the surface small):**
+- Arbitrary spin-s, non-HEALPix samplings, extreme ℓ_max. We are not chasing
+  s2fft-style generality.
+- **Off-grid synthesis (the NUFFT)** — synthesis at arbitrary detector
+  pointings, i.e. the `bk-jax` *sim-forward* path. That is a distinct
+  capability (a JAX NUFFT) and a separate, later phase. "JAX-native SHT" ≠
+  "escape ducc everywhere": jht replaces the **on-grid** SHT only.
+- Being gratuitously BK-coupled. Write the core so BK assumptions live at the
+  edges; the transform itself should be reusable.
+
+## Why jht exists (decision record — see `docs/motivation.md` for the full story)
+
+The recurring constraint is a single C++ black box (ducc0) that does the
+harmonic transforms. It is excellent but:
+- **CPU-only.** The forecast / field-level-inference (MUSE) ambitions are
+  GPU-hungry; ducc can't go there.
+- **A black box.** ducc is vendored in `bk-jax` as a submodule wrapped in a
+  custom XLA FFI — that is *build* control, not *code* control. Its internals
+  are unmodifiable C++, and off-grid pointing tangents raise
+  `NotImplementedError`.
+- **A build tax.** The C++/scikit-build/CUDA toolchain has repeatedly broken
+  CI and local installs.
+
+The leading JAX-native alternative, **s2fft**, was evaluated and rejected:
+its **spin-2 HEALPix inverse has a structural precision defect** (measured
+3–28% error at production ℓ, vs its own documented ~1e-3), and as of
+2026-06 it is **still at v1.4.0 (12 Feb 2026)** with the relevant issues triaged
+and parked ("no time") and a stagnant PR queue. (The 2026-05 evaluation spike
+ran on v1.3.0; the v1.3.0→v1.4.0 diff is CUDA/build-only, so the spin-2 accuracy
+is byte-identical — the date in the original record was wrong, the conclusion
+stands.) **Do not depend on s2fft, and do not fork it** (owning a large codebase
+you find broken is the opposite of dependency control). Clean-room, *informed
+by* the published methods (libsharp / Kostelec–Rockmore for the spin-2-safe
+ℓ-recursion; Price & McEwen for the JAX branch-free renorm *technique* only —
+explicitly NOT their m-recursion, which is the suspected defect home), is the
+control-aligned path.
+
+Crucially, the s2fft evaluation **validated the architecture**: a JAX-native
+SHT integrates cleanly (JIT-clean, spin-0 machine-precision, finite/well-scaled
+VJP). The concept works; one library's spin-2 corner does not.
+
+## Relationship to bk-jax (the consumer)
+
+- jht is a **future dependency** of `bk-jax` (`~/bicepkeck/bk-jax`), added as a
+  path/pip dependency once Phase 0–1 land. It is developed and versioned here,
+  independently — that's the point.
+- **ducc stays in bk-jax for the production path.** The ~1e-4 purity /
+  E→B-leakage-critical bandpower pipeline keeps using ducc; that accuracy is
+  person-decades of hard-won tuning and is not the thing to risk first.
+- **jht earns its keep in the tier ducc can't serve:** GPU forecast sweeps,
+  the MUSE inner Wiener solve, and (eventually) fully-differentiable geometry —
+  where the HEALPix ~1e-3 accuracy floor is acceptable.
+- `bk-jax` already dispatches its transform via a `BK_JAX_SHT_BACKEND` env var;
+  jht becomes a new backend option there (Phase 3).
+
+## Accuracy philosophy (the validation contract)
+
+This is the make-or-break of the project; treat it with the same rigor `bk-jax`
+applies — **a-priori tolerances, gate-driven, and tolerances are NOT relaxed
+without an explicit discussion and sign-off.**
+
+- HEALPix has **no sampling theorem**, so any HEALPix SHT is approximate. The
+  fundamental floor without ring weights is ~1e-3; ducc reaches ~1e-4 and
+  better via **ring quadrature weights + Jacobi (map2alm-style) iteration**.
+- jht's first target is **reach the HEALPix floor with no structural defect**
+  (i.e. spin-2 must NOT show s2fft's 3–28% — it must sit at ~1e-3 like spin-0),
+  then close toward ducc by adding ring weights + iteration.
+- **Spin-2 analysis is the gate.** It is exactly what s2fft got wrong — and the
+  2026-06 dive pinned the defect to the *analysis quadrature* (missing ring
+  weights + polar folding + a stalling unweighted iteration), NOT the recursion
+  (see the crux section). Gate the weighted spin-2 round-trip **per-(ℓ,m), below
+  the ℓ≤1.5·nside band-limit ceiling** so recursion / quadrature / ceiling errors
+  are never conflated; if it reaches the floor flat across all m (no s2fft-style
+  m<ℓ leakage), the project is viable.
+- Every transform lands with a parity test vs **both** healpy and ducc0
+  (cross-checking conventions), at a documented tolerance. Log any residual
+  mismatch (a `DISCREPANCIES.md`, mirroring bk-jax discipline) rather than
+  burying it.
+
+## The crux technical risk
+
+Two distinct risks — and the 2026-06 literature dive **corrected which one is
+make-or-break** (full reasoning in `docs/design.md` and the session plan
+`~/.claude/plans/yeah-let-s-get-to-lovely-feather.md`):
+
+1. **Recursion stability to ℓ_max~1000 — necessary, but NOT s2fft's bug.**
+   Associated-Legendre / Wigner-d recursions underflow at high ℓ; the solved
+   technique is libsharp-style enhanced-exponent ("X-number") scaling. jht uses
+   the libsharp 3-term ℓ-recursion (Kostelec–Rockmore), in the increasing-ℓ
+   direction, with a **branch-free per-step log-renorm** — a *known, solved*
+   problem with a published spin-2≡spin-0 accuracy proof (libsharp §5.1.2).
+   This is **table-stakes, not the differentiator**: s2fft already rescales its
+   recursion and still fails spin-2, so getting the recursion right is necessary
+   but not sufficient.
+
+2. **Spin-2 HEALPix *analysis quadrature* — the actual make-or-break.** s2fft's
+   defect lives here, not in the recursion: its spin-2 single-mode errors appear
+   at ℓ=8/16/32 and *shrink* with ℓ (35→28→22%), spin-0 (identical recursion) is
+   machine-precision, and the failures sit far below the ℓ≤1.5·nside ceiling —
+   all incompatible with recursion underflow, all consistent with an
+   aliasing/folding + missing-weights defect in the spin-2 map2alm path. jht's
+   gate is the **weighted spin-2 round-trip, validated per-(ℓ,m) below the
+   band-limit ceiling**, so recursion / quadrature / ceiling errors are never
+   conflated. Test both risks explicitly and early (Phase 0).
+
+## Conventions — pin and document these from day one
+
+These have repeatedly bitten `bk-jax`; do not leave them implicit.
+
+- **Polarization sign:** HEALPix-internal vs IAU U-sign. Decide the internal
+  convention, document it, and cross-validate against healpy (`pol=True`) and
+  ducc. (bk-jax stores HEALPix-internal and flips U at I/O boundaries.)
+- **Spin sign convention** (spin +2 vs −2 / E,B sign).
+- **aₗₘ layout** (healpy m-major triangular packing; real vs complex; which ℓ,m
+  ordering) and **normalization** (4π vs orthonormal). Mirror healpy/ducc and
+  document the exact mapping.
+- **Differentiability convention:** the on-grid SHT is alm-linear, so JVP/VJP
+  are straightforward — but note the **JAX-VJP-convention vs strict-math-adjoint**
+  subtlety (the `2·conj(·)` on m>0 modes) that bit bk-jax. Register transposes
+  deliberately and test `jacfwd ≡ jacrev` and the inner-product adjoint
+  identity.
+
+## Tooling & conventions
+
+- **Pixi** is the env tool (`pixi run test|lint|format|typecheck`). Standard.
+- **Pure Python / JAX — no compiled extension.** This is a deliberate feature:
+  zero build-toolchain pain, unlike bk-jax's vendored-ducc FFI. Keep it that
+  way; if a kernel ever truly needs C++/CUDA, that's a major decision, not a
+  drive-by.
+- **Runtime deps = jax + numpy only.** ducc0 / healpy are dev/test oracles in
+  the pixi env, never `[project.dependencies]`.
+- **No Jupyter notebooks for working code** — scripts + `ipython %run` for
+  exploration (user-level standard).
+- **JAX float64 is opt-in per entry point** (`jax.config.update("jax_enable_x64",
+  True)` in tests/scripts; library code does not touch global config) — same
+  rule as bk-jax; SHT accuracy work needs float64.
+
+## Attribution (release blocker)
+
+Published / shippable author fields must read **"James Cheshire"**
+(`cheshire@caltech.edu`) — already set in `pyproject.toml` `authors`/
+`maintainers`. "jamie" is fine only in the local git author config. This repo
+may become a published dependency; verify the name before any tag / build /
+publish. ("Jamie" must never appear in a shipped artifact.)
+
+## Dev workflow (per transform)
+
+1. Implement the transform over typed JAX arrays (parametric where a systematic
+   knob could later live).
+2. Gate it against healpy **and** ducc0 at an a-priori documented tolerance.
+3. Pin & document the conventions it assumes (`docs/design.md`).
+4. `pixi run lint` + `typecheck` before committing; terse single-line commit
+   messages (user preference).
+5. Log any numeric mismatch in `DISCREPANCIES.md`, don't bury it as a TODO.
+
+## Pointers
+
+- `ROADMAP.md` — phased plan + hard gates + open questions.
+- `docs/motivation.md` — why jht exists (full decision record).
+- `docs/design.md` — technical design, accuracy tiers, conventions, algorithm.
+- Consumer / sibling context: `~/bicepkeck/bk-jax` (the pipeline that motivates
+  jht; do not make jht depend on reading bk-jax to be understood).

jaxht-0.1.0/DISCREPANCIES.md

@@ -0,0 +1,84 @@
+# DISCREPANCIES
+
+Residual numeric mismatches and known-limitation notes, with their cause — logged
+rather than buried as TODOs (mirroring bk-jax discipline). Each entry says
+whether it is a *defect* (to fix) or an *expected property* (to document).
+
+## Spin-2 cut-sky recovery — E/B ambiguous modes (expected, not a defect)
+
+`jht.masked.deconvolve` recovers the true a_lm from a cut sky only where the cut
+leaves the modes constrained. For **spin-2** there is an additional fundamental
+floor: on a cut sky, E and B are not cleanly separable — a subspace of modes
+(the *ambiguous modes*, Bunn et al. 2003) is supported almost entirely under the
+mask and is therefore unconstrained by the data. These near-null directions of
+`Sᵀ(MW)S` are not recoverable; CG from `x0=0` returns the minimum-norm solution
+within that subspace, and `reg>0` (Tikhonov) bounds it.
+
+- **Status:** expected physics (information loss from the cut), **not** a
+  transform defect. The *operator* is correct — gated to ~1e-13 against the
+  explicit dense `Sᵀ diag(MW) S` (`tests/test_masked.py::test_operator_matches_dense`)
+  and symmetric/PSD — and the pseudo-a_lm reproduces healpy/ducc to machine
+  precision. Only the *recoverability* of cut-localized E/B modes is limited.
+- **Manifestation:** in `scripts/masked_sweep.py`, deconvolution error grows as
+  fsky shrinks and blows up under heavy apodization (which down-weights data).
+  spin-2 reaches the same recovery floor as spin-0 for mild cuts (no s2fft-style
+  m<ℓ defect — that was ruled out in Phase 0); the spin-2-specific limitation is
+  only the ambiguous-mode subspace at aggressive cuts.
+- **Mitigation / scope:** use `reg>0` to stabilise, restrict the recovered band,
+  or — now implemented — bound the ambiguous subspace with a signal prior via
+  `jht.masked.wiener` (the Cℓ-informed generalization of `reg`; the MUSE
+  inner-solve tier), with `constrained_realization` for posterior draws. The
+  operator and recovery are unchanged; the prior simply gives the unconstrained
+  subspace a principled, finite answer instead of the min-norm CG one. See
+  `docs/masked.md`.
+
+## Differentiability: `jax.linear_transpose` / forward-mode of the raw kernels (limitation, not a defect)
+
+The supported autodiff path is JAX's **native reverse/forward AD** through the
+public transforms — `jax.grad` / `vjp` / `jacrev` (complex aₗₘ), and `jacfwd ≡
+jacrev` / FD on the real-DOF layer (`jht.diff`). All Phase-2 gates pass on these
+(`tests/test_grad.py`).
+
+Two lower-level entry points are **not** supported, by deliberate choice:
+
+- `jax.linear_transpose(synthesis, …)` — the forward `lax.scan` in the recursion
+  engine is not directly transposable as a standalone linear function ("Error
+  interpreting argument to scan as a JAX value"). `jax.vjp`/`grad` are unaffected:
+  they linearize first (baking the static recursion plan as constants) and
+  transpose the linearized jaxpr, which is clean. Forward scatters carry
+  `unique_indices=True` so the *scatter*-transpose blocker is gone; the residual is
+  the scan, which native AD sidesteps.
+- `jax.jacfwd` / `jax.jvp` on the **complex** `synthesis` — `jacfwd` requires
+  real-valued inputs (a JAX constraint, not jht-specific). Forward-mode lives on
+  the **real-DOF layer** (`synthesis_real`/`analysis_real`), where `jacfwd ≡
+  jacrev` to ~1e-16.
+
+- **Status:** expected JAX-mechanics limitations with first-class supported
+  alternatives, **not** transform defects. A `custom_vjp` would *block* forward
+  mode entirely; the only kernel-routing fix needs JAX's internal `jax.core.
+  Primitive` API (removed in jax 0.9.2). Native AD is the chosen, version-stable
+  path. See `docs/design.md` §Differentiability.
+
+## Off-grid (`synthesis_general`): spin>0 m=0 phantom vs ducc (expected, not a defect)
+
+For spin > 0 the off-grid synthesis keeps both real components at every m (Re→first
+plane, Im→second), so the field depends on the **unphysical `Im(a_{l,0})`** at m=0 —
+the same phantom documented on-grid in `docs/design.md`. jht's *strict adjoint*
+`adjoint_synthesis_general` therefore carries a nonzero m=0 imaginary part, while
+ducc0's `adjoint_synthesis_general` zeros it.
+
+- **Status:** expected; jht's adjoint is the exact transpose (the inner-product
+  identity `<S a, v> == <a, Sᵀ v>_G` holds to ~1e-16, `tests/test_offgrid.py::
+  test_adjoint_identity`), ducc applies a different m=0 convention. **Harmless** —
+  physical skies have real `a_{l,0}`, so the phantom is never excited. The ducc
+  cross-check (`test_adjoint_vs_ducc`) compares the physical m≥1 modes (matched to
+  ~1e-13); the m≥1 / real-m=0 content agrees with ducc bit-for-bit.
+
+## Off-grid: the −spin channel carries `(−1)^s` (convention pin, now matched to ducc)
+
+The spin-weighted E/B → ±spin mapping is `+s a = −(aE + i aB)`, `−s a =
+−(−1)^s (aE − i aB)`. The `(−1)^s` on the −spin channel is **invisible for even
+spin** (the on-grid 0/2 path never exposed it) but flips the result for **odd spin
+(1, 3)**. Pinned by matching ducc0's `synthesis_general` for spin 1/3 (a sign probe;
+`tests/test_offgrid.py::test_forward_vs_ducc` covers all of spin 0–3). Documented so
+the even-only on-grid convention is not assumed to carry over to odd spin.

jaxht-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 James Cheshire
+
+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.

jaxht-0.1.0/PKG-INFO

@@ -0,0 +1,165 @@
+Metadata-Version: 2.4
+Name: jaxht
+Version: 0.1.0
+Summary: JAX-native spherical harmonic transforms (BK-regime first): GPU-capable, differentiable, dependency-controlled.
+Project-URL: Homepage, https://github.com/jrcheshire/jht
+Project-URL: Repository, https://github.com/jrcheshire/jht
+Project-URL: Issues, https://github.com/jrcheshire/jht/issues
+Project-URL: Changelog, https://github.com/jrcheshire/jht/blob/main/CHANGELOG.md
+Author-email: James Cheshire <cheshire@caltech.edu>
+Maintainer-email: James Cheshire <cheshire@caltech.edu>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: automatic-differentiation,cosmic-microwave-background,cosmology,gpu,healpix,jax,nufft,spherical-harmonic-transform
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.4.30
+Requires-Dist: numpy>=1.24
+Provides-Extra: dev
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Provides-Extra: validation
+Requires-Dist: ducc0; extra == 'validation'
+Requires-Dist: healpy; extra == 'validation'
+Description-Content-Type: text/markdown
+
+# jht — JAX Harmonic Transforms
+
+JAX-native spherical harmonic transforms (map ↔ aₗₘ): **GPU-capable**, **fully
+differentiable**, and **dependency-controlled** (pure JAX + numpy at runtime — no
+compiled C++ extension, no heavyweight third-party SHT library). Scoped to the
+BICEP/Keck regime — **spin-0 and spin-2** on the **HEALPix RING** pixelization,
+ℓ_max ≲ 1000, nside ≤ ~2048 — but written cleanly so it can serve as a general
+transform dependency.
+
+It exists to serve the GPU / differentiable tier of analysis that a CPU-only C++
+transform (ducc0) structurally cannot, while *owning the numerics*. See
+[`docs/motivation.md`](docs/motivation.md) for the full decision record.
+
+## Status (2026-06-10)
+
+Phases 0–4 **complete and validated** (190 tests pass + 8 GPU-gated skips,
+CPU/float64):
+
+- **On-grid transforms** — spin-0 & spin-2 synthesis (`aₗₘ→map`) and the exact
+  adjoint `Sᵀ`, validated to machine precision vs healpy **and** ducc0; spin-2
+  inverse at the HEALPix floor with **no s2fft-style structural defect**.
+- **Accuracy** — jht's own ring quadrature weights + Jacobi iteration reach
+  ~1e-13 on band-limited maps (matches `healpy.map2alm(use_weights=True)`); see
+  [`docs/accuracy.md`](docs/accuracy.md).
+- **Partial-sky** — masked pseudo-aₗₘ, a cut-sky CG deconvolution, and a masked
+  Wiener filter / constrained realization (the MUSE inner solve); see
+  [`docs/masked.md`](docs/masked.md).
+- **Off-grid (NUFFT)** — `synthesis_general` / `adjoint_synthesis_general`
+  evaluate a band-limited field at **arbitrary pointings** (spin 0–3), alm- **and**
+  pointing-differentiable. The JAX-native replacement for ducc0's
+  `sht.experimental.synthesis_general` (on-grid SHT + this NUFFT = the full ducc0
+  surface bk-jax depends on); see [`docs/offgrid.md`](docs/offgrid.md).
+- **Differentiability** — native JAX autodiff (`jacfwd ≡ jacrev`, tight adjoint
+  identity), plus a convention-clean real-DOF layer `jht.diff`; see
+  [`docs/design.md`](docs/design.md) §Differentiability.
+- **GPU** — pure JAX, so the transforms run on CUDA with no code change. Measured on
+  Cannon A100/V100 (fp64): GPU==CPU parity ~1e-13 across the BK regime **including
+  nside=2048**, forward synthesis 14–60× the CPU. See **Performance** below and
+  [`docs/gpu.md`](docs/gpu.md).
+
+## Install
+
+Standard env is [pixi](https://pixi.sh):
+
+```bash
+pixi install          # CPU env (osx-arm64 / linux-64)
+pixi run test         # the gated suite
+```
+
+GPU (CUDA, linux-64 — see [`docs/gpu.md`](docs/gpu.md)):
+
+```bash
+pixi run -e gpu python scripts/gpu_check.py   # on an NVIDIA box
+```
+
+As a dependency in another project (runtime deps are just `jax` + `numpy`):
+
+```bash
+pip install jaxht        # once released on PyPI — then `import jht`
+# or track the repo directly:
+pip install "jaxht @ git+https://github.com/jrcheshire/jht.git"
+```
+
+## Quick start
+
+float64 is **opt-in per entry point** — enable it *before creating any array*
+(library code never touches the global config):
+
+```python
+import jax
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+import jht
+
+nside, lmax, spin = 256, 512, 0
+m = jht.synthesis(alm, nside, lmax, spin=spin)          # aₗₘ -> map
+a = jht.map2alm(m, nside, lmax, spin=spin, niter=3)     # map -> aₗₘ (weighted + iterated)
+cl = jht.bandpower(a, lmax, spin=spin)                  # angular auto-power C_ℓ
+```
+
+`spin=2` takes/returns `(E, B)` aₗₘ of shape `(2, …)` and `(Q, U)` maps of shape
+`(2, npix)`. `jht.adjoint_synthesis` is the **exact unweighted transpose** `Sᵀ`
+(the operator seam / VJP), distinct from `map2alm` (the approximate inverse). For
+gradient-based work use the real-DOF layer `jht.synthesis_real` /
+`jht.analysis_real` (plain ℝⁿ→ℝᵐ, no complex-conjugate convention subtlety).
+
+## Conventions
+
+healpy m-major triangular aₗₘ packing, orthonormal Yₗₘ with the Condon–Shortley
+phase, HEALPix-internal (COSMO) polarization — verified against healpy 1.19.0 and
+ducc0 0.41.0. Pinned in [`docs/design.md`](docs/design.md).
+
+## Accuracy tiers (the contract)
+
+jht targets the **GPU/differentiable tier** where the HEALPix ~1e-3 sampling
+floor is acceptable; weights + iteration close it to ~1e-13 on band-limited
+inputs. It is **not** a drop-in for ducc's purity-critical (~1e-4 E→B-leakage)
+production path. Tolerances are a-priori and gate-driven, never relaxed without
+sign-off. Residual mismatches are logged in
+[`DISCREPANCIES.md`](DISCREPANCIES.md).
+
+## Performance
+
+Pure JAX runs unchanged on CUDA. Measured on Cannon A100 (incl. a 20 GB MIG) / V100, fp64:
+
+- **GPU==CPU parity ~1e-13** across the BK regime, **including nside=2048** (synthesis and `map2alm`).
+- **Forward synthesis 14–60×** the 8-core CPU; fp64/fp32 ≈ 2.2×.
+- **Off-grid forward** ~0.5–0.9 s at ℓ_max=1000 (independent of the number of points; recursion-bound), with the pointing gradient ~1× a forward.
+- **nside=2048** compiles and runs on GPU — a ~20 GB slice holds synthesis + `map2alm`; the one-time compile is multi-minute (jit-cached).
+
+The recurring GPU lesson: fp64/complex scatters are catastrophic on GPU, so jht packs and assembles via **gathers**. CPU perf model + memory in [`docs/performance.md`](docs/performance.md); GPU detail in [`docs/gpu.md`](docs/gpu.md).
+
+## Using jht as a dependency
+
+jht is standalone and consumer-agnostic. The operator/grad seam a downstream
+needs (e.g. to use jht *in place of ducc0*) — and the accuracy boundary — are
+documented in [`docs/consumers.md`](docs/consumers.md). Any backend-selection
+wiring lives in the consumer, not here.
+
+## Docs
+
+- [`docs/design.md`](docs/design.md) — technical design, conventions, the crux, differentiability.
+- [`docs/accuracy.md`](docs/accuracy.md) — the accuracy contract + ring-weight algorithm.
+- [`docs/masked.md`](docs/masked.md) — partial-sky estimators.
+- [`docs/performance.md`](docs/performance.md) — CPU perf model + memory.
+- [`docs/gpu.md`](docs/gpu.md) — the GPU env, the x64 requirement, the harness.
+- [`docs/consumers.md`](docs/consumers.md) — the downstream-dependency seam.
+- [`docs/motivation.md`](docs/motivation.md) — why jht exists.
+- [`ROADMAP.md`](ROADMAP.md) — phased plan + gates.