diff --git a/.github/workflows/physics-validation.yml b/.github/workflows/physics-validation.yml new file mode 100644 index 00000000..9ed2b83f --- /dev/null +++ b/.github/workflows/physics-validation.yml @@ -0,0 +1,83 @@ +name: physics-validation + +# Nightly + release-tag lane for the headline trust anchors that the PR gate skips: +# the @pytest.mark.slow validations (LIMEPY reference parity, multimass equilibrium, +# Engine A-vs-B anchor, Engine B AD-vs-FD). PROGENAX_STRICT_REFS=1 makes a missing +# reference cache a hard failure here (audit R8/T4), so these gates cannot silently +# no-op. Only the PUBLIC jaxstro sibling is needed (the lock is decoupled from the +# private jaxstroviz/fluxax — audit R1/D2); jaxstroviz-dependent experimental SBC +# tests import-skip. + +on: + schedule: + - cron: "17 9 * * *" # nightly + workflow_dispatch: + push: + tags: ["v*"] + +env: + JAX_ENABLE_X64: "1" + XLA_PYTHON_CLIENT_PREALLOCATE: "false" + XLA_PYTHON_CLIENT_ALLOCATOR: platform + OMP_NUM_THREADS: "1" + PROGENAX_STRICT_REFS: "1" + +jobs: + slow-physics: + runs-on: ubuntu-latest + timeout-minutes: 90 + steps: + - name: Checkout progenax + uses: actions/checkout@v4 + with: + path: progenax + - name: Checkout jaxstro (public path dependency, sibling layout) + uses: actions/checkout@v4 + with: + repository: jaxstro/jaxstro + path: jaxstro + - name: Setup uv + uses: astral-sh/setup-uv@v6 + with: + python-version: "3.13" + - name: Sync from lockfile (dev + experimental) + working-directory: progenax + run: uv sync --locked --extra dev --extra experimental + - name: Strict reference guard (cache must be present) + working-directory: progenax + run: uv run --no-sync pytest tests/validation/test_strict_refs.py -q + - name: Slow physics-validation suite (the advertised trust anchors) + working-directory: progenax + run: uv run --no-sync pytest tests -m slow -q + + # Version matrix (audit T3): requires-python is >=3.10 but only 3.13 is exercised + # on the PR lane. Nightly, run a representative memory-safe core slice (imf + + # profiles — the IMF maths + the diffrax ODE profiles, the most version-sensitive + # numerics) on the floor and ceiling Python. Not the PR lane (keeps PR latency flat). + version-matrix: + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.13"] + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - name: Checkout progenax + uses: actions/checkout@v4 + with: + path: progenax + - name: Checkout jaxstro (public path dependency, sibling layout) + uses: actions/checkout@v4 + with: + repository: jaxstro/jaxstro + path: jaxstro + - name: Setup uv + uses: astral-sh/setup-uv@v6 + with: + python-version: ${{ matrix.python-version }} + - name: Sync (dev) on Python ${{ matrix.python-version }} + working-directory: progenax + run: uv sync --extra dev --python ${{ matrix.python-version }} + - name: Core version-compat slice (imf + profiles) + working-directory: progenax + run: uv run --no-sync pytest tests/unit/imf tests/unit/profiles -m "not slow" -q diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index b2eac1f5..1183942e 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -40,14 +40,27 @@ jobs: working-directory: progenax run: uv lock --check - # Released-core invariant, SHARDED across the three test tiers. Each shard runs in a - # fresh process on its own runner -> bounded memory (no OOM), parallel (fast wall-clock), - # and isolates which tier fails. fail-fast: false so all shards report. + # Released-core invariant, SHARDED. Each shard runs in a FRESH process on its own + # runner -> bounded memory, parallel wall-clock, isolates which shard fails. + # The ~900-test unit tier is split into 3 sub-shards (~300 each): one process + # accumulating all of it OOM-killed the 2-core/7 GB runner (exit 143, audit R1 + # follow-up). unit-c is a CATCH-ALL (tests/unit minus the two heavy groups), so + # a NEW tests/unit subdir is still run without editing this matrix. + # fail-fast: false so all shards report. released-core: strategy: fail-fast: false matrix: - suite: [unit, integration, validation] + include: + # imf and profiles are the two heaviest modules (profiles compiles the + # diffrax-ODE King/Michie/LIMEPY + Eddington families) -> each gets its + # own shard. unit-c is the catch-all (everything else) so a new subdir is + # still covered. The conftest CI cache-clear bounds per-module memory. + - { name: unit-a, paths: "tests/unit/imf" } + - { name: unit-b, paths: "tests/unit/profiles" } + - { name: unit-c, paths: "tests/unit --ignore=tests/unit/imf --ignore=tests/unit/profiles" } + - { name: integration, paths: "tests/integration" } + - { name: validation, paths: "tests/validation" } runs-on: ubuntu-latest timeout-minutes: 30 steps: @@ -67,9 +80,9 @@ jobs: - name: Sync from lockfile (dev extra) working-directory: progenax run: uv sync --locked --extra dev - - name: Run tests/${{ matrix.suite }} + - name: Run ${{ matrix.name }} working-directory: progenax - run: uv run --no-sync pytest tests/${{ matrix.suite }} -m "not slow" -q + run: uv run --no-sync pytest ${{ matrix.paths }} -m "not slow" -q # Experimental gravoturb_fdf subsystem (repo-only, NOT in the wheel). Needs the # [experimental] extra (blackjax, optax); runs the non-slow experimental suite. @@ -97,22 +110,55 @@ jobs: working-directory: progenax run: uv run --no-sync pytest tests/experimental -m "not slow" -q + # Wheel-smoke (audit M6): build the wheel, install it in a CLEAN venv with only + # the public jaxstro sibling, and import the package + the diagnostics subpackage. + # Catches packaging regressions an editable dev install hides — missing runtime + # deps, experimental/gravoturb leakage, and the R9 eager numpy/scipy import (the + # clean venv has no scipy, so progenax.diagnostics must import lazily). + wheel-smoke: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout progenax + uses: actions/checkout@v4 + with: + path: progenax + - name: Checkout jaxstro (path dependency, sibling layout) + uses: actions/checkout@v4 + with: + repository: jaxstro/jaxstro + path: jaxstro + - name: Setup uv + uses: astral-sh/setup-uv@v6 + with: + python-version: "3.13" + - name: Build wheel + working-directory: progenax + run: uv build --wheel -o dist/ + - name: Install wheel in a clean venv and import + working-directory: progenax + run: | + uv venv /tmp/clean + uv pip install --python /tmp/clean/bin/python dist/*.whl ../jaxstro + /tmp/clean/bin/python -c "import progenax; import progenax.diagnostics; print(progenax.__name__, 'imports clean (no scipy in venv)')" + # Aggregator gate: a single `tests` check that is green iff every required job above # succeeded — preserves the `tests` branch-protection check name. `if: always()` so a # failed/cancelled dependency reds it (a skipped dependent job would otherwise let # failures slip through branch protection). tests: if: always() - needs: [lock-check, released-core, experimental] + needs: [lock-check, released-core, experimental, wheel-smoke] runs-on: ubuntu-latest timeout-minutes: 5 steps: - name: All required jobs succeeded run: | - echo "lock-check=${{ needs.lock-check.result }} released-core=${{ needs.released-core.result }} experimental=${{ needs.experimental.result }}" + echo "lock-check=${{ needs.lock-check.result }} released-core=${{ needs.released-core.result }} experimental=${{ needs.experimental.result }} wheel-smoke=${{ needs.wheel-smoke.result }}" if [ "${{ needs.lock-check.result }}" = "success" ] \ && [ "${{ needs.released-core.result }}" = "success" ] \ - && [ "${{ needs.experimental.result }}" = "success" ]; then + && [ "${{ needs.experimental.result }}" = "success" ] \ + && [ "${{ needs.wheel-smoke.result }}" = "success" ]; then echo "all required jobs green" else echo "::error::a required job did not succeed" diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 00000000..8e89cf81 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,76 @@ +# Changelog + +All notable changes to progenax are documented here. This project adheres to +[Semantic Versioning](https://semver.org/). + +## 0.1.0 (unreleased) + +First public-ready release candidate. This cycle resolves the release blockers from +the 2026-06-11 pre-release adversarial audit +(`docs/reviews/2026-06-11-prerelease-adversarial-audit.md`), moving progenax from Beta +toward Release Candidate. Finding IDs (R*, J*, S*, A*, T*, D*, M*) refer to that report. + +### Science fixes + +- **Moe & Di Stefano F_twin normalization** (R3): the twin excess was mixed against the + whole q ∈ [0.1, 1] population instead of the paper's q > 0.3 convention (realized + F_twin 0.367 vs Table-13 0.300 at solar logP=1). `MoeDiStefano2017Full` rebuilt as a + jointly-normalized PL + ft/(1−ft)·I_B twin block; pinned at four Table-13 nodes (`0adc74e`). +- **King/Michie/EFF core resolution at high concentration** (R4): the linear 1000-pt + position-CDF grid under-resolved the core (measured +18% at W0=9, +270% at W0=12, 0.3 r_c). + Replaced with a sqrt-stretched `r = r_t·u²` grid + non-uniform trapezoid; EFF's + `compute_profile_potential` made consistent (`d35164e`, `f57e34d`). +- **`sample_fixed_n` silent mass shortfall** (R5): asking 500 M☉ for n=1000 silently + returned 349 M☉. Now raises on an unreachable target; first tests for all four + mass-target sampling modes (`9085e88`). +- **`compute_stellar_radii` mass–radius relation** (R6): exponents were inverted vs MS + homology (10 M☉ → 6.3 R☉). Adopted Demircan & Kahraman (1991) (`20c62ec`). +- **PowerLawIMF α=1 NaN gradient** (R10): ported the `exp_safe` double-where so gradients + are finite at exactly α=1 (`b98ff0f`). +- **EFF γ=3 sub-virial offset** pinned (S2); honest equilibrium caveats for rotation + overlays, tidal truncation, mass segregation, the binary energy budget, and subgroup + virial (S3/S4/S6/S10/S16) (`240ac81`, `7946fdf`). + +### JAX / API hardening + +- King `r_t` boundary-pinning is now flagged (traced) or refused (eager) instead of + silently wrong (J4) (`43d5481`). +- `virial_scale` refuses T=0 input and is deduped with `rescale_velocities_to_virial` (J5) + (`6d3d7b6`). +- `q_approx` uses Python-`if` dispatch on static N (no double-branch compile) (J6) (`cb1b94b`). +- Inconsistent `KingProfile` r_t now warns; the dead `KingVelocityDF.r_t` field removed + (S1/A3) (`3a86878`). Zero-rotation-axis now raises (S15) (`7946fdf`). +- Single source of truth for the `VelocityDF` protocol (A1) (`83948ba`). + +### Documentation + +- Exported the 15 documented IMF/binary symbols from `progenax.__init__`; removed the + phantom `IGIMF`/`EnvironmentIMF` class claims (the environment-dependent IMF is the + functional `BirthEnvironment` + `env_to_imf_params` API) (R7) (`a09e5a8`). +- README rewritten against the current API, with an executable-examples smoke test that + runs every code block; installation.md fixed (R7/D3/D4/H2) (`b825b2a`). +- Removed the fabricated `pip install progenax-legacy` claim (D5/H4) (`7ac2547`). +- Reconciled the three conflicting test-count claims (T5/L3/D6) (`e0943be`). + +### Packaging & CI + +- **CI un-broken** (R1/D2): the lock was decoupled from the private `jaxstroviz` sibling + (CI had been red since 2026-06-10); the unit tier is sharded and the JAX compilation + cache is bounded in CI to fix OOM kills (`7fe4cec`, `2bb8f44`, `bc98083`). +- Nightly + release-tag lane running the slow-marked headline validations, with + `PROGENAX_STRICT_REFS=1` turning a missing reference cache into a failure (R8/T4) + (`02689fb`, `095e3b2`). +- **LICENSE** added (Apache-2.0, © 2026 Anna Rosen) + author/classifiers/URL in the wheel + metadata (D1) (`a688e53`). +- `[diagnostics]` extra + lazy numpy/scipy imports with actionable errors — `import + progenax.diagnostics` no longer crashes in a clean install (R9) (`b140bac`). +- wheel-smoke (clean-venv import) + nightly Python version matrix (3.10/3.13) (M6/T3) + (`bc98083`). + +### Known limitations / deferred (require a separate decision or arc) + +- **Not yet pip-installable for outsiders** (R2): the runtime depends on the unpublished + sibling `jaxstro`. Resolution (publish `jaxstro` to PyPI, vendor its `units`+`jaxconfig`, + or a git-URL dependency) is a strategy decision tracked separately. +- Units-policy A2 (DF `G=None` defaults vs protocol-wide explicit G) — a breaking sweep, + deferred. A two-sided quantile stretch for `sample_fixed_n` (R5 follow-up) is deferred. diff --git a/CLAUDE.md b/CLAUDE.md index e7eaeefa..f1a4c099 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,8 +6,9 @@ Differentiable initial conditions for N-body simulations in JAX. Part of the **j **Status**: Phase 1 + 2026-06 audit hardening + binaries SoTA arc (Batches 4f–4k) + gravoturbulent-FDF clean-room rewrite + multi-component cluster arc (Engine A LIMEPY DF tables + -Engine B density-defined Eddington) complete - 16,957 LOC released-core source, 1485 tests -(released-core 1163: unit 895, integration 34, validation 234; experimental 322). King & EFF velocity DFs are true +Engine B density-defined Eddington) + 2026-06 pre-release audit hardening complete. +Released-core ~1243 tests (unit 956, integration 43, validation 244) + experimental 322; +see CI for the live count. King & EFF velocity DFs are true equilibria (lowered-Maxwellian / Eddington inversion). The binary-population engine is finalized: IMF→companion composition (`build_binary_cluster` over `primary_imf × companion_model × target`), faithful Moe & Di Stefano (2017) P–q–e coupling (`MoeCompanions`), the binary→spatial connector @@ -28,18 +29,18 @@ outer-venv clash and `--no-sync` runs against the installed env without re-locki # Install (released core + experimental extras: blackjax, optax for the inference layer) env -u VIRTUAL_ENV uv pip install -e ".[dev,experimental]" -# Released-core invariant (1163 tests). The multimass-LIMEPY equilibrium tests make the -# serial suite ~17 min; use pytest-xdist with XLA threads capped (one process per core). -# FAST GATE (inner loop, 1121 tests, ~4 min): excludes @pytest.mark.slow +# Released-core invariant (~1243 tests; see CI for the live count). The multimass-LIMEPY +# equilibrium tests make the serial suite slow; use pytest-xdist with XLA threads capped. +# FAST GATE (inner loop, ~1192 tests, ~4 min): excludes @pytest.mark.slow XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" \ env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit tests/integration tests/validation -q -m "not slow" -n auto -# FULL GATE (phase/commit gate, 1163 tests, ~9 min parallel): +# FULL GATE (phase/commit gate, ~1243 tests, ~9 min parallel): XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" \ env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit tests/integration tests/validation -q -n auto -env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit/ -q # 895 unit tests (released core) -env -u VIRTUAL_ENV uv run --no-sync pytest tests/integration/ -q # 34 integration tests -env -u VIRTUAL_ENV uv run --no-sync pytest tests/validation/ -q # 234 physics validation tests +env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit/ -q # ~956 unit tests (released core) +env -u VIRTUAL_ENV uv run --no-sync pytest tests/integration/ -q # ~43 integration tests +env -u VIRTUAL_ENV uv run --no-sync pytest tests/validation/ -q # ~244 physics validation tests # Experimental gravoturb_fdf subsystem (repo-only; needs src/experimental on the path): PYTHONPATH=src:src/experimental env -u VIRTUAL_ENV uv run --no-sync pytest tests/experimental -q # 322 tests @@ -113,7 +114,7 @@ jax.grad(loss)(1.0) # Fully differentiable! |--------|---------|-------------| | `profiles/` | Spatial density profiles | `PlummerProfile`, `KingProfile`, `MichieProfile`, `EFFProfile` | | `kinematics/` | Velocity DFs + transforms | `PlummerVelocityDF`, `KingVelocityDF`, `MichieVelocityDF`, `EFFVelocityDF` | -| `imf/` | Initial mass functions + binary stats | `PowerLawIMF`, `ChabrierIMF`, `IGIMF`, `BinaryIMF`, `MoeJointOrbit` | +| `imf/` | Initial mass functions + binary stats | `PowerLawIMF`, `ChabrierIMF`, `Maschberger`, `BinaryIMF`, `MoeJointOrbit`; environment-dependent IMF via `BirthEnvironment` + `env_to_imf_params()` | | `binaries/` | Orbital mechanics + connector + diagnostics | `KeplerElements`, `resolve_binary_components()`, `IndependentCompanions`, `MoeCompanions`, `binary_energy_budget()` | | `analytical/` | Test cases with exact solutions | `two_body_kepler()`, `three_body_figure_eight()` | | `diagnostics/` | Substructure + mass-segregation diagnostics | `compute_q_parameter()` (CW04 Q, A=πR²), `q_approx` (differentiable kNN), `energy_sorted_segregation()` | @@ -193,9 +194,9 @@ energy = compute_total_energy(positions, velocities, masses, G=PLANETARY.G) # W ## Test Structure ```text -tests/ 1163 released-core tests -├── unit/ 895 tests -│ ├── imf/ IMF tests (PowerLaw, Chabrier, IGIMF, Binary, Moe full P-q-e) +tests/ ~1243 released-core tests (see CI for the live count) +├── unit/ ~956 tests +│ ├── imf/ IMF tests (PowerLaw, Chabrier, environment, Binary, Moe full P-q-e) │ ├── profiles/ Profile tests (Plummer, King, EFF) │ ├── kinematics/ Velocity DF tests + anisotropy │ ├── analytical/ Analytical test case tests @@ -203,12 +204,12 @@ tests/ 1163 released-core tests │ ├── cluster/ MultiComponentCluster (Engine A + Engine B) + mass-segregation IC tests │ ├── dynamics/ Virial / energy utilities │ └── substructure/ CW04 Q diagnostic (compute_q_parameter, q_approx, baselines) -├── integration/ 34 tests +├── integration/ ~43 tests │ ├── test_jax_compatibility.py JIT/grad/vmap tests │ ├── test_units_through_pipeline.py G threading (audit C1) │ ├── test_binary_cluster.py build_binary_cluster (budgets + companions) │ └── test_end_to_end.py Full IC → energy checks -└── validation/ 234 tests +└── validation/ ~244 tests ├── test_plummer_physics.py Plummer equilibrium ├── test_king_physics.py King true-DF equilibrium + c(W0) ├── test_multimass_equilibrium_physics.py MultiComponentCluster shared-potential equilibrium @@ -250,7 +251,7 @@ All public symbols exported from `progenax.__init__`: **Velocity DFs**: `PlummerVelocityDF`, `KingVelocityDF`, `EFFVelocityDF` (Plummer/EFF take an optional `anisotropy_radius` for Osipkov-Merritt radial anisotropy, β(r)=r²/(r²+r_a²)), `MichieVelocityDF` (self-consistent anisotropic King = Michie 1963 + King 1966 cutoff; pairs with `MichieProfile`), `apply_solid_body_rotation()`, `apply_differential_rotation()` -**IMFs**: `PowerLawIMF`, `ChabrierIMF`, `Maschberger`, `TruncatedIMF`, `BinaryIMF`, `IGIMF`, `EnvironmentIMF`; mass-ratio: `FlatMassRatio`, `PowerLawMassRatio`, `TwinPeakedMassRatio`, `MoeDiStefano2017`, `MoeDiStefano2017Full`, `MoePeriod`, `MoeJointOrbit`; fractions: `ConstantBinaryFraction`, `MassDependentBinaryFraction` +**IMFs**: `PowerLawIMF`, `ChabrierIMF`, `Maschberger`, `TruncatedIMF`, `BinaryIMF`; mass-ratio: `FlatMassRatio`, `PowerLawMassRatio`, `TwinPeakedMassRatio`, `MoeDiStefano2017`, `MoeDiStefano2017Full`, `MoePeriod`, `MoeJointOrbit`; fractions: `ConstantBinaryFraction`, `MassDependentBinaryFraction`. The environment-dependent stellar IMF (Marks+2012 / Jeřábková+2018 α₃ relations) is the functional `BirthEnvironment` + `env_to_imf_params()` API — **not** a galaxy-wide IGIMF integration and **not** an `IGIMF`/`EnvironmentIMF` class (those never existed; audit R7). **Binaries**: `KeplerElements`, `BinaryOrbitalState`, `compute_period()`, `period_to_semimajor_axis()`, `LogUniformPeriod`, `LogNormalPeriod`, `SanaOBPeriod`, `ThermalEccentricity`, `UniformEccentricity`, `MoeEccentricity`; **connector/composition**: `resolve_binary_components()`, `ResolvedBinaries`, `CompanionElements`, `IndependentCompanions`, `MoeCompanions`; **diagnostics**: `relative_energy()`, `find_bound_pairs()`, `find_bound_multiples()`, `primordial_survival()`, `binary_energy_budget()`, `BinaryEnergyBudget` diff --git a/LICENSE b/LICENSE new file mode 100644 index 00000000..a454826b --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or Derivative + Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, excluding + those notices that do not pertain to any part of the Derivative + Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and do + not modify the License. You may add Your own attribution notices + within Derivative Works that You distribute, alongside or as an + addendum to the NOTICE text from the Work, provided that such + additional attribution notices cannot be construed as modifying + the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions for + use, reproduction, or distribution of Your modifications, or for any + such Derivative Works as a whole, provided Your use, reproduction, + and distribution of the Work otherwise complies with the conditions + stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2026 Anna Rosen + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md index ccdb48e6..9125be4d 100644 --- a/README.md +++ b/README.md @@ -2,11 +2,21 @@ Differentiable initial conditions for N-body simulations in JAX. -Part of the **jaxstro ecosystem** - providing IC generation that can be differentiated through for gradient-based inference. +Part of the **jaxstro ecosystem** — providing IC generation that can be differentiated through for gradient-based inference. ## Status -**Phase 1 + 2026-06 audit hardening + binaries SoTA arc (Batches 4f–4k) + gravoturbulent-FDF clean-room rewrite complete**: 14,731 LOC released-core source code, 965 tests passing (released-core 815: unit 653, integration 34, validation 128). King & EFF velocity DFs are true equilibria (lowered-Maxwellian / Eddington inversion). The binary-population engine is finalized — `build_binary_cluster` composes `primary_imf × companion_model × target` (Systems / Stars / TotalMass budgets), with the faithful Moe & Di Stefano (2017) P–q–e coupling (`MoeCompanions`) and energy-budget diagnostics. The gravoturbulent + fractal-density-field subsystem was rebuilt clean-room (2026-06) as the **experimental, repo-only `gravoturb_fdf` package** (1,400 LOC, 150 tests; `src/experimental/`, **not** in the released wheel). +**Research software, release-candidate core.** King, Michie, EFF, and multi-component +(LIMEPY / density-defined Eddington) velocity DFs are true equilibria (lowered-Maxwellian / +Eddington inversion, no external virial rescale). The binary-population engine is finalized — +`build_binary_cluster` composes `primary_imf × companion_model × target` with the faithful +Moe & Di Stefano (2017) P–q–e coupling (`MoeCompanions`). The gravoturbulent + fractal- +density-field subsystem lives in the **experimental, repo-only `gravoturb_fdf` package** +(`src/experimental/`, **not** shipped in the wheel). For current test/LOC counts see CI and +`tests/README.md` (counts are not duplicated here to avoid drift). + +> **Not on PyPI yet.** progenax depends on the sibling package `jaxstro` (also unpublished), +> resolved from a side-by-side checkout. See [Installation](#installation). ## Features @@ -14,9 +24,12 @@ Part of the **jaxstro ecosystem** - providing IC generation that can be differen | Profile | Class | Description | Reference | |---------|-------|-------------|-----------| -| **Plummer** | `PlummerProfile` | Plummer (1911) sphere ρ ∝ (1 + r²/a²)^(-5/2) | Plummer (1911) | +| **Plummer** | `PlummerProfile` | Plummer (1911) sphere ρ ∝ (1 + r²/a²)^(−5/2) | Plummer (1911) | | **King** | `KingProfile` | King (1966) lowered isothermal model | King (1966) | -| **EFF** | `EFFProfile` | Elson-Fall-Freeman (1987) truncated power law | EFF (1987) | +| **Michie** | `MichieProfile` | Anisotropic King (Michie 1963 + King 1966 cutoff) | Michie (1963) | +| **EFF** | `EFFProfile` | Elson–Fall–Freeman (1987) truncated power law | EFF (1987) | +| **LIMEPY** | `LIMEPYProfile` | Lowered-isothermal multi-mass family (single- or multi-component) | Gieles & Zocchi (2015) | +| **Uniform sphere** | `UniformSphereProfile` | Uniform-density ball (baselines/tests) | — | | **King ODE** | `solve_king_profile()` | Numerical King profile via Diffrax | King (1966) | ### Velocity Distribution Functions @@ -24,16 +37,27 @@ Part of the **jaxstro ecosystem** - providing IC generation that can be differen | DF | Class | Method | Notes | |----|-------|--------|-------| | **Plummer DF** | `PlummerVelocityDF` | Beta(3/2, 9/2) sampling | Exact, no rejection | -| **King DF** | `KingVelocityDF` | Lowered-Maxwellian f(E) ∝ e^{E/σ²}−1, σ self-consistent | True equilibrium (Q≈0.5 unscaled) | -| **EFF DF** | `EFFVelocityDF` | Isotropic Eddington inversion f(E) from ρ(Ψ) | True equilibrium (Q≈0.5 unscaled) | +| **King DF** | `KingVelocityDF` | Lowered-Maxwellian f(E), σ self-consistent | True equilibrium (Q≈0.5 unscaled) | +| **Michie DF** | `MichieVelocityDF` | Self-consistent anisotropic King | True equilibrium; pairs with `MichieProfile` | +| **EFF DF** | `EFFVelocityDF` | Isotropic Eddington inversion f(E) from ρ(Ψ) | γ=3 default ~8% sub-virial (documented); γ≥5 ~equilibrium | +| **LIMEPY DF** | `LIMEPYVelocityDF` | Lowered-isothermal f(E) (multi-mass) | Parity vs canonical LIMEPY code | + +`PlummerVelocityDF` and `EFFVelocityDF` take an optional `anisotropy_radius` (Osipkov–Merritt +radial anisotropy, β(r) = r²/(r²+r_a²)). Rotation overlays: `apply_solid_body_rotation()`, +`apply_differential_rotation()` (kinematic overlays — they inject energy/L_z and are NOT +stationary equilibria). + +### Multi-component clusters -**Velocity Transforms:** +`MultiComponentCluster` builds a self-consistent multi-component equilibrium in a shared +potential (replaces the deleted `TwoComponentConfig`/`generate_two_component_cluster`): -| Transform | Function | Description | -|-----------|----------|-------------| -| **Osipkov-Merritt** | `apply_osipkov_merritt()` | Radial anisotropy β(r) = r²/(r²+r_a²) | -| **Solid Body Rotation** | `apply_solid_body_rotation()` | v_rot = ω × r | -| **Differential Rotation** | `apply_differential_rotation()` | Peaked rotation curve | +| Constructor | Engine | Description | +|-------------|--------|-------------| +| `from_components()` | A | GZ15 multi-mass LIMEPY (one ψ ODE, Σ_j α_j ρ̂_j) | +| `from_imf()` | A | components binned from an IMF | +| `from_mass_segregation()` | A | primordial mass segregation | +| `from_density_profiles()` | B | prescribed Plummer/EFF/King densities → shared-Ψ Eddington/OM DFs | ### Initial Mass Functions @@ -43,10 +67,13 @@ Part of the **jaxstro ecosystem** - providing IC generation that can be differen | **Kroupa** | `PowerLawIMF.kroupa()` | 3-segment Kroupa (2001) | Kroupa (2001) | | **Maschberger** | `Maschberger` | Smooth transitions | Maschberger (2013) | | **Chabrier** | `ChabrierIMF` | Log-normal + power law | Chabrier (2003) | -| **Truncated** | `TruncatedIMF` | Mass limits wrapper | - | +| **Truncated** | `TruncatedIMF` | Mass-limit wrapper | — | | **Binary** | `BinaryIMF` | Primary IMF + mass-dependent companions (q, f_b) | Moe & Di Stefano (2017) | -| **IGIMF** | `IGIMF` | Integrated galactic IMF | Weidner+ (2013) | -| **Environment** | `EnvironmentIMF` | Metallicity/SFR dependent | Jeans mass theory | + +The **environment-dependent stellar IMF** (top-heavy α₃ from cluster density/metallicity) is +the functional `BirthEnvironment` + `env_to_imf_params()` API (Marks+2012 Fundamental Plane / +Jeřábková+2018 relations) — **not** a galaxy-wide IGIMF integration, and not an +`IGIMF`/`EnvironmentIMF` class. ### Binary Orbital Mechanics @@ -54,384 +81,264 @@ Part of the **jaxstro ecosystem** - providing IC generation that can be differen |-----------|----------------|-------------| | **Keplerian Elements** | `KeplerElements` | (a, e, i, Ω, ω, M₀) + masses | | **Orbital State** | `BinaryOrbitalState` | 6D position/velocity + masses | -| **Period Calculation** | `compute_period()` | T = 2π√(a³/GM) | +| **Period** | `compute_period()` | T = 2π√(a³/GM) | | **Inverse Period** | `period_to_semimajor_axis()` | a from T and M | -| **Batch Conversion** | `batch_elements_to_resolved()` | Vectorized element → state | - -**Period Distributions:** - -| Distribution | Class | Description | -|--------------|-------|-------------| -| Log-Uniform | `LogUniformPeriod` | Öpik's law | -| Log-Normal | `LogNormalPeriod` | Duquennoy & Mayor (1991) | -| Sana O/B | `SanaOBPeriod` | O-star periods (Sana+ 2012) | - -**Eccentricity Distributions:** - -| Distribution | Class | Description | -|--------------|-------|-------------| -| Thermal | `ThermalEccentricity` | f(e) = 2e (Heggie 1975) | -| Uniform | `UniformEccentricity` | f(e) = 1 | -| Moe+2017 | `MoeEccentricity` | p(e) ∝ e^η(logP, M₁) + Roche cap (Moe & Di Stefano 2017) | -| Logistic-thermal | `LogisticThermalEccentricity` | Circular→thermal heuristic (Duquennoy & Mayor 1991) | - -**Binary-cluster composition** (`build_binary_cluster`): composes a primary IMF with a -`CompanionModel` and a population-size budget. -| Component | Class | Description | -|-----------|-------|-------------| -| Budget | `Systems` / `Stars` / `TotalMass` | fix #systems (companions not counted) / #stars / total mass | -| Companions (versatile) | `IndependentCompanions` | independent f_b × q × P × e marginals | -| Companions (faithful) | `MoeCompanions` | Moe+2017 joint P–q–e; same q sets m₂ (self-consistent) | -| Connector | `resolve_binary_components()` | binary COMs → 2N components (COM-preserving) | +**Period distributions:** `LogUniformPeriod` (Öpik), `LogNormalPeriod` (Duquennoy & Mayor 1991), +`SanaOBPeriod` (Sana+ 2012), `MoePeriod` (Moe & Di Stefano 2017). +**Eccentricity distributions:** `ThermalEccentricity` (f(e)=2e), `UniformEccentricity`, +`MoeEccentricity` (p(e) ∝ e^η + Roche cap). -**Binary diagnostics:** `find_bound_pairs()`, `find_bound_multiples()`, -`primordial_survival()`, and `binary_energy_budget()` (separates the cluster COM virial -from the internal binary binding-energy reservoir). +**Binary-cluster composition** (`build_binary_cluster`): a primary IMF × a `CompanionModel` +(`IndependentCompanions` or the faithful `MoeCompanions`) × a population budget +(`Systems` / `Stars` / `TotalMass`). Connector `resolve_binary_components()` maps binary COMs +to 2N components. Diagnostics: `find_bound_pairs()`, `find_bound_multiples()`, +`primordial_survival()`, `binary_energy_budget()`. ### Analytical Test Cases -| System | Function | Description | -|--------|----------|-------------| -| **Kepler Orbit** | `two_body_kepler()` | Circular/eccentric two-body | -| **Figure-8** | `three_body_figure_eight()` | Moore (2001) periodic solution | -| **Earth-Sun** | `earth_sun_2body()` | JPL ephemerides | -| **Solar System** | `solar_system_inner_4()` | Sun + inner 4 planets | -| **Full Solar System** | `solar_system_full()` | Sun + 8 planets | -| **Harmonic Oscillator** | `harmonic_oscillator()` | 1D/2D SHO | +`two_body_kepler()`, `three_body_figure_eight()`, `earth_sun_2body()`, +`solar_system_inner_4()`, `solar_system_full()`, `harmonic_oscillator()`. ### Utilities -| Category | Functions | -|----------|-----------| -| **IC Builder** | `build_spatial_ic()`, `ICResult` | -| **Energetics** | `compute_kinetic_energy()`, `compute_potential_energy()` | -| **Transforms** | `to_com_frame()`, `virial_scale()` | -| **Stellar Radii** | `compute_stellar_radii()` (3-regime M-R relation) | -| **Tidal Physics** | `jacobi_radius()`, `jacobi_radius_isothermal()`, `apply_tidal_truncation()` | -| **Substructure diagnostic** | `compute_q_parameter()` (CW04 Q), `q_approx` (differentiable kNN) | -| **Two-Component** | `TwoComponentConfig`, `generate_two_component_cluster()` | -| **Mass Segregation** | `energy_sorted_segregation()` | +`build_spatial_ic()` / `ICResult`, `compute_kinetic_energy()`, `compute_potential_energy()`, +`to_com_frame()`, `virial_scale()`, `compute_stellar_radii()` (Demircan & Kahraman 1991 M–R), +`jacobi_radius()` / `apply_tidal_truncation()`, `energy_sorted_segregation()`, +substructure diagnostics `compute_q_parameter()` (CW04 Q) and `q_approx` (differentiable kNN; +needs the `[diagnostics]` extra for the exact scipy path). ## Installation -```bash -conda activate astro -pip install -e ".[dev]" -``` - -Or with UV (faster): +progenax is **not on PyPI**; it (and its sibling `jaxstro`) resolve from a standard +side-by-side checkout. Use **uv** (preferred): ```bash +# Clone progenax and its public sibling jaxstro next to each other +git clone https://github.com/jaxstro/jaxstro.git +git clone https://github.com/jaxstro/progenax.git +cd progenax + +# Install the released core + dev tools (uses the project .venv) uv pip install -e ".[dev]" ``` +Optional extras: `[experimental]` (the repo-only `gravoturb_fdf` inference layer — blackjax, +optax, …), `[diagnostics]` (numpy + scipy, for the exact non-differentiable CW04 Q path). + ## Quick Start -### Plummer Sphere IC +### Plummer sphere IC ```python import jax import jax.numpy as jnp from jaxstro.units import STELLAR -from progenax.profiles import PlummerProfile -from progenax.kinematics import PlummerVelocityDF +from progenax import PlummerProfile, PlummerVelocityDF -# Create profile and velocity DF -profile = PlummerProfile(r_h=1.0) # Half-mass radius = 1 pc +profile = PlummerProfile(r_h=1.0) # half-mass radius = 1 pc velocity_df = PlummerVelocityDF(r_h=1.0) -# Sample positions and velocities N = 1000 -masses = jnp.ones(N) # 1000 Msun total -key = jax.random.PRNGKey(42) -key_pos, key_vel = jax.random.split(key) - +masses = jnp.ones(N) # 1000 Msun total +key_pos, key_vel = jax.random.split(jax.random.PRNGKey(42)) positions = profile.sample_positions(masses, key_pos) velocities = velocity_df.sample_velocities(positions, masses, key_vel, G=STELLAR.G) ``` -### Two-Component Cluster +### One-call IC assembly ```python +import jax +import jax.numpy as jnp from jaxstro.units import STELLAR -from progenax import ( - PlummerProfile, PlummerVelocityDF, - TwoComponentConfig, generate_two_component_cluster, -) +from progenax import PlummerProfile, PlummerVelocityDF, build_spatial_ic -# Extended halo + concentrated core -config = TwoComponentConfig( - f_A=0.3, # 30% in extended halo - profile_A=PlummerProfile(r_h=2.0), - profile_B=PlummerProfile(r_h=0.5), - velocity_df_A=PlummerVelocityDF(r_h=2.0), - velocity_df_B=PlummerVelocityDF(r_h=0.5), +masses = jnp.ones(500) +ic = build_spatial_ic( + PlummerProfile(r_h=1.0), masses, PlummerVelocityDF(r_h=1.0), + key=jax.random.PRNGKey(0), G=STELLAR.G, ) +positions, velocities = ic.positions, ic.velocities +``` + +### Multi-component cluster (extended halo + concentrated core) -positions, velocities, pop_id = generate_two_component_cluster( - masses, config, key, G=STELLAR.G +```python +import jax +import jax.numpy as jnp +from jaxstro.units import STELLAR +from progenax import MultiComponentCluster + +# Two mass components in a shared self-consistent potential (Engine A, LIMEPY). +model = MultiComponentCluster.from_components( + alpha_j=jnp.array([0.6, 0.4]), # mass fractions + w_j=jnp.array([1.0, 1.0]), # relative velocity scales + m_j=jnp.array([0.4, 5.0]), # mean masses [Msun] + W0=7.0, g=1.0, r_c=1.0, ) +ic = model.sample_cluster(jax.random.PRNGKey(1), n_stars=2000, G=STELLAR.G) +positions, velocities, masses = ic.positions, ic.velocities, ic.masses ``` -### IMF Sampling +### IMF sampling ```python import jax -from progenax.imf import PowerLawIMF, ChabrierIMF - -# Kroupa IMF -imf = PowerLawIMF.kroupa() -key = jax.random.PRNGKey(42) -masses = imf.sample(key, 1000) # Sample 1000 stellar masses +from progenax import PowerLawIMF, ChabrierIMF -# Chabrier IMF -imf = ChabrierIMF() -masses = imf.sample(key, 1000) +k1, k2 = jax.random.split(jax.random.PRNGKey(42)) +kroupa_masses = PowerLawIMF.kroupa().sample(k1, 1000) # Kroupa (2001) +chabrier_masses = ChabrierIMF().sample(k2, 1000) # Chabrier (2003) ``` -### Binary Orbital Elements +### Binary orbital elements ```python from jaxstro.units import PLANETARY -from progenax.binaries import KeplerElements, compute_period - -# Create Keplerian elements -elements = KeplerElements( - a=1.0, # Semi-major axis [AU] - e=0.3, # Eccentricity - i=0.1, # Inclination [rad] - Omega=0.0, # Longitude of ascending node - omega=0.0, # Argument of periapsis - M0=0.0, # Mean anomaly at epoch -) - -# Convert to Cartesian state -M_total = 2.0 # Solar masses -state = elements.to_state(M_total, G=PLANETARY.G) -r, v = state['position'], state['velocity'] +from progenax import KeplerElements, compute_period -# Compute orbital period +elements = KeplerElements(a=1.0, e=0.3, i=0.1, Omega=0.0, omega=0.0, M0=0.0) +state = elements.to_state(M_total=2.0, G=PLANETARY.G) # CartesianState +r, v = state.position, state.velocity period = compute_period(a=1.0, M_total=2.0, G=PLANETARY.G) ``` -### Analytical Test Case +### Analytical test case (Earth–Sun) ```python from jaxstro.units import PLANETARY -from progenax.analytical import two_body_kepler - -# Create Earth-like orbit around Sun -ic = two_body_kepler( - a=1.0, # 1 AU - e=0.017, # Earth's eccentricity - m1=1.0, # 1 Msun - m2=3e-6, # ~1 Earth mass - G=PLANETARY.G, -) +from progenax import two_body_kepler -positions = ic.positions # (2, 3) -velocities = ic.velocities # (2, 3) -masses = ic.masses # (2,) +ic = two_body_kepler(M1=1.0, M2=3e-6, a=1.0, G=PLANETARY.G, e=0.017) +positions, velocities, masses = ic.positions, ic.velocities, ic.masses ``` -## Unit Systems - -progenax uses **jaxstro.units** for gravitational constant management: +### Differentiability ```python -from jaxstro.units import STELLAR, PLANETARY +import jax +import jax.numpy as jnp +from progenax import PlummerProfile -# STELLAR.G for star clusters (~0.00450 pc³ Msun⁻¹ Myr⁻²) -velocities = velocity_df.sample_velocities(positions, masses, key, G=STELLAR.G) +masses = jnp.ones(100) +key = jax.random.PRNGKey(0) -# PLANETARY.G for binaries/planets (~39.478 AU³ Msun⁻¹ yr⁻² = 4π²) -orbit = two_body_kepler(a=1.0, e=0.3, m1=1.0, m2=1.0, G=PLANETARY.G) +def loss(r_h): + pos = PlummerProfile(r_h=r_h).sample_positions(masses, key) + return jnp.mean(jnp.linalg.norm(pos, axis=1)) -# Convenience default (wrapper behavior only) -from progenax import DEFAULT_UNITS -velocities = velocity_df.sample_velocities(positions, masses, key, G=DEFAULT_UNITS.G) +gradient = jax.grad(loss)(1.0) # ∂⟨r⟩/∂r_h — fully differentiable ``` +## Unit Systems + +progenax uses **jaxstro.units** for explicit gravitational-constant management. Core APIs +require an explicit `G` (or `units`); convenience wrappers may resolve `G=None` to +`DEFAULT_UNITS` (= `STELLAR`). + | Unit System | G Value | Units | Use Case | |-------------|---------|-------|----------| | `STELLAR` | ~0.00450 | pc³ Msun⁻¹ Myr⁻² | Star clusters, galaxies | | `PLANETARY` | ~39.478 | AU³ Msun⁻¹ yr⁻² | Binaries, planets | -| `DEFAULT_UNITS` | = STELLAR | - | Default for progenax | - -## Architecture - -```text -progenax/ -├── src/progenax/ -│ ├── __init__.py # Public API (57+ exports) -│ ├── protocols.py # 3 runtime-checkable protocols -│ ├── builders.py # ICResult + build_spatial_ic -│ ├── populations.py # Two-component clusters -│ ├── tidal.py # Jacobi radius + truncation -│ ├── profiles/ -│ │ ├── plummer.py # PlummerProfile -│ │ ├── king.py # KingProfile + ODE solver -│ │ ├── eff.py # EFFProfile -│ │ └── mass_segregation.py -│ ├── kinematics/ -│ │ ├── plummer_df.py # PlummerVelocityDF (Beta sampling) -│ │ ├── king_df.py # KingVelocityDF -│ │ ├── eff_df.py # EFFVelocityDF -│ │ ├── anisotropy.py # Osipkov-Merritt -│ │ └── rotation.py # Solid body, differential -│ ├── imf/ -│ │ ├── base.py # BaseIMF with custom JVP -│ │ ├── power_law.py # PowerLawIMF (Kroupa, Salpeter) -│ │ ├── chabrier.py # ChabrierIMF -│ │ ├── environment.py # EnvironmentIMF -│ │ ├── binary.py # BinaryIMF -│ │ └── igimf.py # IGIMF -│ ├── binaries/ -│ │ ├── kepler.py # KeplerElements, Kepler solver -│ │ ├── orbital_state.py # BinaryOrbitalState -│ │ └── population.py # Period/eccentricity distributions -│ ├── diagnostics/ -│ │ ├── substructure.py # CW04 Q parameter (compute_q_parameter, A=πR²) -│ │ ├── q_approx.py # differentiable kNN Q approximation -│ │ └── mass_segregation.py -│ └── analytical/ -│ └── core.py # Solar system, Kepler orbits -└── tests/ - ├── unit/ # 653 unit tests - ├── integration/ # 34 integration tests - └── validation/ # 128 physics validation tests -``` - -> **Note (2026-06):** the gravoturbulent + fractal-density-field subsystem (and the -> `generate_fractal_positions` GW04 generator) was rebuilt clean-room as the experimental, -> repo-only `gravoturb_fdf` package under `src/experimental/` — not shipped in the wheel. +| `DEFAULT_UNITS` | = STELLAR | — | progenax default (wrappers only) | ## Key Patterns -### Protocol-Based Composition +### Protocol-based composition -Any `SpatialProfile` can pair with any `VelocityDF`: +Any `SpatialProfile` pairs with any `VelocityDF` (9 runtime-checkable protocols total): ```python -from progenax.protocols import SpatialProfile, VelocityDF - -profile: SpatialProfile = PlummerProfile(r_h=1.0) -df: VelocityDF = KingVelocityDF(W0=7.0, r_c=1.0, r_t=10.0) - -# Mix Plummer positions with King velocities! -positions = profile.sample_positions(masses, key_pos) -velocities = df.sample_velocities(positions, masses, key_vel, G=STELLAR.G) +import jax +import jax.numpy as jnp +from jaxstro.units import STELLAR +from progenax import PlummerProfile, KingVelocityDF + +masses = jnp.ones(500) +k_pos, k_vel = jax.random.split(jax.random.PRNGKey(0)) +profile = PlummerProfile(r_h=1.0) +df = KingVelocityDF(W0=7.0, r_c=1.0) # r_t derived from W0 (self-consistent) +positions = profile.sample_positions(masses, k_pos) +velocities = df.sample_velocities(positions, masses, k_vel, G=STELLAR.G) ``` -### Equinox Modules +### Equinox modules -All stateful classes are Equinox modules (immutable PyTrees): +All stateful classes are immutable Equinox PyTrees (e.g. `PlummerProfile` carries `r_h` and a +computed scale radius `a`), so sampling is differentiable through JAX and JIT/vmap-friendly. -```python -import equinox as eqx - -class PlummerProfile(eqx.Module): - r_h: Float[Array, ""] # Half-mass radius - a: Float[Array, ""] # Scale radius (computed) -``` - -### Differentiability - -All sampling is differentiable through JAX: +### Critical formula (Plummer scale radius) ```python -def loss(r_h): - profile = PlummerProfile(r_h=r_h) - positions = profile.sample_positions(masses, key) - return jnp.mean(jnp.linalg.norm(positions, axis=1)) - -grad_fn = jax.grad(loss) -gradient = grad_fn(1.0) # Fully differentiable! +import jax.numpy as jnp +r_h = 1.0 +a = r_h * jnp.sqrt(2 ** (2 / 3) - 1) # ≈ 0.7664 * r_h (NOT the inverse!) ``` -### Critical Formulas +Virial ratio convention: **Q = T/|V|**, with **Q = 0.5** the equilibrium (2T + V = 0). -**Plummer Scale Radius:** -```python -# From half-mass radius to scale radius -a = r_h * jnp.sqrt(2**(2/3) - 1) # ≈ 0.7664 * r_h -``` +## Architecture -**Virial Ratio:** -```python -Q = T / |V| # Q ≈ 0.5 for equilibrium (virial theorem: 2T + V = 0) +```text +src/progenax/ +├── __init__.py # public API +├── protocols.py # 9 runtime-checkable protocols +├── builders.py # ICResult + build_spatial_ic + build_binary_cluster +├── tidal.py # Jacobi radius + truncation +├── profiles/ # Plummer, King, Michie, EFF, LIMEPY, UniformSphere +├── kinematics/ # velocity DFs (+ Osipkov–Merritt anisotropy, rotation overlays) +├── imf/ # PowerLaw, Chabrier, Maschberger, Binary; environment/ (BirthEnvironment) +├── binaries/ # Kepler elements/solver, period & eccentricity distributions, connector +├── cluster/ # MultiComponentCluster (Engine A + B), mass segregation, turbulence +├── diagnostics/ # CW04 Q (compute_q_parameter), differentiable q_approx +└── analytical/ # solar system, Kepler orbits, figure-eight +src/experimental/ # gravoturb_fdf (repo-only; NOT in the wheel) ``` ## Testing ```bash -# All released-core tests -pytest tests/ -v # 815 tests, ~55s +# FAST gate (inner loop): released-core, excluding slow +XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" \ + uv run --no-sync pytest tests/unit tests/integration tests/validation -m "not slow" -n auto -# By tier -pytest tests/unit/ -v # 653 unit tests -pytest tests/integration/ -v # 34 integration tests -pytest tests/validation/ -v # 128 physics validation tests +# FULL gate (phase/commit gate): includes the slow trust anchors +XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" \ + uv run --no-sync pytest tests/unit tests/integration tests/validation -n auto # Experimental gravoturb_fdf subsystem (repo-only) -PYTHONPATH=src:src/experimental pytest tests/experimental -v # 150 tests - -# Specific modules -pytest tests/unit/imf/ -v -pytest tests/validation/test_plummer_physics.py -v +PYTHONPATH=src:src/experimental uv run --no-sync pytest tests/experimental ``` -### Physics Validation +See `tests/README.md` for the 3-tier architecture and current counts. -From `tests/validation/` (Q ≡ T/|V|; 0.5 = equilibrium). King & EFF DFs are -sampled in detailed equilibrium with no external virial rescale: +### Physics validation (selected) -- **Plummer virial ratio**: Q = 0.502 (expected 0.5) -- **King true-DF virial ratio**: Q ≈ 0.51 unscaled (lowered-Maxwellian DF) -- **EFF Eddington-DF virial ratio**: Q ≈ 0.50 unscaled (γ=5, mild truncation) -- **King concentration c(W₀)**: matches King (1966) Table II to Δ ≤ 0.02 -- **Bound particles**: 100% within escape velocity -- **Kepler orbits**: energy & angular momentum conserved to ~1e-16; period exact to 1e-10 +From `tests/validation/` (Q ≡ T/|V|; 0.5 = equilibrium), sampled with no external rescale: + +- **Plummer** virial Q = 0.502 (expected 0.5) +- **King** true-DF virial Q ≈ 0.51 unscaled (lowered-Maxwellian) +- **EFF** Eddington-DF virial Q ≈ 0.50 unscaled (γ=5, mild truncation) +- **King c(W₀)** matches King (1966) Table II to |Δlog₁₀c| ≤ 0.002 (hi-res) +- **Engine B** King A-vs-B σ_1d / radial KS ≤ 3e-4 (N=2e4) +- **Kepler** energy & angular momentum conserved to ~1e-16; period exact to 1e-10 ## Dependencies -- `jax>=0.4.20` -- `jaxlib>=0.4.20` -- `equinox>=0.11.0` -- `jaxtyping>=0.2.25` -- `diffrax>=0.4.0` (for King profile ODE integration) -- `jaxstro` (core utilities) +`jax>=0.4.20`, `jaxlib>=0.4.20`, `equinox>=0.11.0`, `jaxtyping>=0.2.25`, +`diffrax>=0.4.0` (King ODE), and the sibling `jaxstro` (core utilities; side-by-side checkout). ## References -**Density Profiles:** -- Plummer (1911), MNRAS 71, 460 -- King (1966), AJ 71, 64 -- Elson, Fall & Freeman (1987), ApJ 323, 54 - -**Velocity Distribution Functions:** -- Dehnen (1993), MNRAS 265, 250 - Exact Plummer DF -- Binney & Tremaine (2008), "Galactic Dynamics" - -**Initial Mass Functions:** -- Salpeter (1955), ApJ 121, 161 -- Kroupa (2001), MNRAS 322, 231 -- Chabrier (2003), PASP 115, 763 -- Maschberger (2013), MNRAS 429, 1725 - -**Binary Populations:** -- Duquennoy & Mayor (1991), A&A 248, 485 -- Sana+ (2012), Science 337, 444 -- Moe & Di Stefano (2017), ApJS 230, 15 - -**Substructure:** -- Goodwin & Whitworth (2004), A&A 413, 929 -- Kupper+ (2011), MNRAS 417, 2300 - McLuster - -**N-body Methods:** -- Aarseth (2003), "Gravitational N-Body Simulations" +**Profiles/DFs:** Plummer (1911); King (1966); Michie (1963); Elson, Fall & Freeman (1987); +Gieles & Zocchi (2015); Dehnen (1993); Binney & Tremaine (2008); Merritt (1985). +**IMFs:** Salpeter (1955); Kroupa (2001); Chabrier (2003); Maschberger (2013); +Marks+ (2012); Jeřábková+ (2018). +**Binaries:** Duquennoy & Mayor (1991); Sana+ (2012); Moe & Di Stefano (2017). +**Substructure/Methods:** Cartwright & Whitworth (2004); Goodwin & Whitworth (2004); +Küpper+ (2011); Aarseth (2003). ## License -MIT +Apache License 2.0 — see [LICENSE](LICENSE). Copyright 2026 Anna Rosen. diff --git a/STATUS.md b/STATUS.md index 504c4a57..2bb2e2a6 100644 --- a/STATUS.md +++ b/STATUS.md @@ -1,6 +1,7 @@ # progenax — status -next: **(1) Merge/push of `feat/multimass-limepy-validation` — AWAITS ANNA'S GO. (2) Batch B demos Tasks 6–9** (demo-only, independent): wrong-IMF bias curve + robustness grid, B3 halo+core recovery, website 60-science-demos section, close-out. **DONE (2026-06-11, branch `feat/multimass-limepy-validation`, commits 0bb5cc0..f0f1722): Multimass LIMEPY validation arc** (plan `docs/plans/2026-06-11-multimass-limepy-hardening-plan.md`; completion doc `.claude-work/MULTIMASS_LIMEPY_VALIDATION_COMPLETE.md`). Phase-1 PDF verification found+fixed TWO doc equation errors (Peuten eqs 8–9 ratio was inverted — paper reads (m̄*/m̄)^{2δ}; Bianchini A1 prefactor is μ^{−δ} and the paper itself omits the ^{1/2} on the E_γ ratio, confirmed via its own A2). **Task 1 parity**: Engine A ≡ canonical reference LIMEPY (SHA ef2a479, `meanmassdef='central'` ⇒ like-for-like, no eq-8/9 translation) to ≤1.5e-4 across 6 configs (single-mass g=1/1.5, 2-comp, B2 4-comp Maschberger, aniso η=0/0.5); gates measured-then-frozen ≤5e-4 (20× under the plan's 1% target); the reference runs ONLY in a pinned-env subprocess (py3.11/numpy1.26.4/scipy1.11.4 — project scipy 1.17 breaks it) with a committed npz cache + skip-if-absent slow test (3 tests). **Task 2 (Anna-reviewed → FIT-FREE)**: equipartition saturation validated by exact statements only — solver≡closed-form A1 to 1.6e-7; HEADLINE μ→0 identity m/(2η)≡derived m_eq=m̄(g+5/2)(g+7/2)/φ̂₀ to 0.11–0.19% (W0=5/7/9); η(m) shape via autodiff (flat 0.026–0.039 at m_min → ½ exactly deep-well). KEY SCIENCE: the DF saturates SHARPER than exponentially (Bianchini A2's quadratic coefficient is NEGATIVE, opposite eq-3's) ⇒ finite-window eq-3 fits recover ~0.5× the derived m_eq (0.50–0.53 over μ≈[0.14,1.7]; →1.0 as the window shrinks) — documented as a literature caution (window-match before comparing; future-work idea: fit the exact differentiable σ(m;g,W₀,m̄) to kinematic data). **Task 3 docs**: NEW theory page `10-theory/spatial-profiles/multimass-equipartition.md` + corrected per-paper notes wired into nav (myst 147 pages, 0 warnings); docstring-only src edit (GZ15 §4 √-update named). **Close-out**: FULL gate **1192 passed** (= main baseline 1189 + 3 parity tests; CLAUDE.md's "1163" is a stale count). meq/zeta stay DEFERRED. **DONE (2026-06-11, merged+pushed): Multicomponent Engine A IFT hardening (H1–H5)** — `find_alpha_for_masses` → hand-rolled `custom_vjp` (adaptive while_loop fwd + reverse-mode implicit-VJP backward), ~9× faster B2 demo, exact grads (AD-vs-FD iso 2e-7), jit(grad) pinned, FD-gate tightened 1e-6; review caught+fixed a Critical jit(grad) regression. **+ Batch B Tasks 1–5 DONE:** B1 cross-engine figure; B2 joint (α,δ,W₀) MLE recovering reference θ̂ bit-for-bit; Task 5 Fisher information panel (kin-only ρ(α,δ)=−0.26 mild, mass channel PINS α 4.7×) + NUTS corner (0 divergences, posterior-mean within 1σ of MLE, flat-in-θ target). **+ equation-verified per-paper notes** (Bianchini2016, Peuten2017 new; GZ15 updated) documenting the multimass methods/theory. Replaced the unrolled 30-iter eigenvalue scan in `find_alpha_for_masses` with a hand-rolled `jax.custom_vjp` (adaptive `while_loop` forward to tol=1e-6 + reverse-mode implicit-VJP backward of the √-map residual R; TWO solvers dispatched on `ra_hat is None`). ~3× faster per `value_and_grad`, flat-in-n_iter, gradients exact (AD-vs-FD iso 2e-7, aniso quadrature-oracle 5e-5; jit(grad)/jit(value_and_grad) pinned). The B2 demo (`demo_delta_recovery.py`) right-sized 4×600→3×300 → **5:42 vs ~50 min (~9×)**, recovers the gold reference θ̂=(α 2.2931±0.0041, δ 0.3972±0.0344, W₀ 4.9900±0.0209) **bit-for-bit**, OVERALL ALL PASS. Review cadence caught + fixed a Critical `jit(grad)` regression (`jnp.asarray(eta)` in a nondiff_argnums slot) at H2, two tasks before the demo that would have hit it. optimistix REJECTED (its JVP adjoint forces diffrax DirectAdjoint → slower). Fresh gates: **full released-core 1188 passed (7:17)**, memory 7/7 PASS, box-corner harness 3/3 (residual 5–8e-7), demo exit 0. Completion doc `.claude-work/MULTICOMPONENT_IFT_HARDENING_COMPLETE.md`. Future IFT candidates noted (Engine B trunc radius, King tidal radius). Commits `dd14713..ccf6ed1` (6) + harness/tests. **DONE (2026-06-10, Anna-approved): DF sampler-fusion micro-batch** (profile-driven: staged RSS decomposition showed construction ~1.7 GB dominates, eager draw +0.63 GB): speed-CDF tables cached at construction + jitted sampling cores (eqx.filter_jit) for LIMEPY/King/Michie standalone DFs — warm aniso draws **0.35 s → 15 ms** at N=2e4 (3–16 ms across N; 55× vs pre-batch quadrature at N=1e4), draw-chain peak +0.63→+0.20 GB, stage 2.37→2.02 GB, `limepy_df_aniso` gate tightened 3.0→2.5 GB, small-N build penalty eliminated; lax.map chunking of the table draw evaluated and rejected on measurement. Fidelity re-verified on the jitted path (identical KS/moment/β, ALL PASS). Gate: **1163** released-core (unit 895; fast 1121). Remaining DF-stage floor is model CONSTRUCTION (~1.8 GB, paid once per model) — future target if ever needed. **DONE (2026-06-10, Anna-approved): Batch A merged to local main** (fast-forward to `2263948`; FAST gate re-verified on merged main 1108; main NOT pushed — 22+ ahead of origin, awaits Anna's go) **+ the performance-evidence addendum** (Anna-requested, plan-approved): `scripts/benchmark_batch_a.py` measures BOTH code versions live (pre-batch worktree at `0dd1cd9` vs HEAD) — wall-clock: PE 2.4× faster at N=8k (0.169→0.070 s), N=2e4 feasible at 0.216 s (dense: 20.9 GB, infeasible), honest small-N regressions measured (PE 3.5× slower at N=2k; DF table-build dominates below ~5k); memory: 16× at N=8k (gated ≥5×); fidelity figures table-vs-quadrature ALL PASS (KS ≤ 6.7e-3); NEW `50-validation/performance-memory.md` page (myst 144 pages, 0 warnings) + King/Michie fidelity embeds. The multi-component cluster arc is DONE end-to-end: Engines A (LIMEPY DF tables) + B (density-defined Eddington) merged to main (`4f4f259`), the post-merge code-review fix batch landed (`0dd1cd9`), and **Batch A (memory + consolidation + docs)** is complete — blocked+remat'd virial kernels (forward AND grad O(block·N)), table-routed standalone LIMEPY/King/Michie DFs (quadrature kept as the `speed_method="quadrature"` oracle), `progenax.numerics` consolidation (cumtrap ×5 + inverse-CDF draws), `_EngineAState` grouping (A-only access on B raises informative AttributeError), shared angular conditional in `kinematics/_speed_kernels.py`, `lax.map(batch_size=2048)` bounded oracles, and the docs metrics refresh. **Memory headline (measured, `scripts/profile_cluster_memory.py`, all PASS): virial kernel at N=2e4 32.8 GB → 0.41 GB; standalone anisotropic DF sampling 10.9 GB → 2.3 GB; Engine A/B builds at N=1e5 ≤ 2.5 GB.** Gate: released-core **1150 passed** (unit 882 / integration 34 / validation 234; fast gate 1108). **Merge/push of `fix/memory-and-consolidation` awaits Anna's go.** +next: **Pre-release audit fixes COMPLETE (2026-06-12, branch `fix/prerelease-audit`, 27 commits `7fe4cec..86339d1`; plan `docs/plans/2026-06-12-prerelease-audit-fixes.md`, completion `.claude-work/PRERELEASE_AUDIT_FIXES_COMPLETE.md`).** All R1–R10 + S/J/A/T/D blockers fixed via strict RED→GREEN TDD. CI un-broken (lock decoupled from the private jaxstroviz sibling) + un-OOM'd (unit-tier sharding + CI-gated jax.clear_caches) + new wheel-smoke/nightly-version-matrix/strict-refs lanes; Moe F_twin q>0.3 (0.367→0.300), King/Michie/EFF core grid (+270%→<0.1%), sample_fixed_n guard, D&K91 radii, α=1 grad, J4/J5/J6/S1/S2/A1, honest equilibrium caveats; R7 README rewrite + executable-examples smoke test, phantom IGIMF/EnvironmentIMF + progenax-legacy claims killed, counts reconciled; Apache-2.0 LICENSE + wheel metadata, [diagnostics] extra + lazy scipy, CHANGELOG. FAST gate 1191 passed (exit 0); myst 151 pages 0 warnings; wheel metadata verified. **blocker: AWAITS ANNA — (1) R2 pip-installability strategy (publish jaxstro to PyPI / vendor units+jaxconfig / git-URL dep); (2) merge to local main + ONE final PR when CI is green. due: Anna's go.** +PRIOR next (audit-fix arc now complete): **Triage the 2026-06-11 pre-release adversarial audit** (`docs/reviews/2026-06-11-prerelease-adversarial-audit.md`, 6 parallel specialist agents): verdict **Beta, ~2 weeks from RC**. Top blockers: CI RED on main (missing jaxstroviz sibling checkout — zero tests enforced for ≥5 pushes); slow-marked headline validations (LIMEPY parity, A-vs-B) never run in CI; Moe F_twin normalized against q>0.1 instead of the paper's q>0.3 (+22% twins, measured); King/Michie CDF grid under-resolves the core at W0≳9 (+18% enclosed-mass error at 0.3 r_c, measured); `sample_fixed_n` silent 30% mass shortfall (zero tests); `compute_stellar_radii` exponents inverted vs MS homology; README teaches deleted `TwoComponentConfig` + phantom IGIMF/EnvironmentIMF; wheel uninstallable (unpublished jaxstro dep, no LICENSE). Then: **Merge/push decision — AWAITS ANNA'S GO** for `feat/batch-b-demos` (and the earlier `feat/multimass-limepy-validation`); then the gravax forward chain (render recovered ICs) OR paper assembly — Anna's call. **DONE (2026-06-11, branch `feat/batch-b-demos`, commits `760e086..15e112e` + close-out): Batch B science demos Tasks 6–9** (plan `docs/plans/2026-06-10-batch-b-science-demos.md`; completion doc `.claude-work/BATCH_B_SCIENCE_DEMOS_COMPLETE.md`). **B2d** wrong-IMF bias (dδ̂/dα = −0.0047 ± 0.0195, δ̂ PEAKS at the true α; REPORTED not gated — adjudicated the α=2.1 scatter as real realization noise, not optimizer artifact) + robustness grid (α_true 1.9/2.3/2.7 all within 3σ). **B3** halo+core Engine-B (t, r_a, r_h) recovery (0.602±0.008 / 3.051±0.053 / 2.015±0.029, all <1σ; 3 inits→identical loss 19.1479; realizability f_min 1.16e-4; binned-expectation σ + closed-form OM β). **Website** `60-science-demos/` (4 pedagogical pages + 6 figures wired after Validation; `myst build` exit 0, 0 warnings, **151 pages** — the plan's 148 assumed a stale 144 baseline). **Close-out**: FULL released-core gate **1192 passed (9:07), exit 0**; all four demos (B1/B2/B3/B2d) re-run exit 0 ALL PASS, fresh-captured. Integrity catches: B3 docstring had a FABRICATED "measured" runtime (the subagent never completed a run → corrected to measured 3.7 min, MLE-scan-compile dominated); B1 lowered-DF eq index is g (NOT g+½ — verified vs GZ15 eq 2 / limepy.py, King limit E₁=eˣ−1); independent B3 review → occupancy-STOP now exits before the figure. **DONE (2026-06-11, branch `feat/multimass-limepy-validation`, commits 0bb5cc0..f0f1722): Multimass LIMEPY validation arc** (plan `docs/plans/2026-06-11-multimass-limepy-hardening-plan.md`; completion doc `.claude-work/MULTIMASS_LIMEPY_VALIDATION_COMPLETE.md`). Phase-1 PDF verification found+fixed TWO doc equation errors (Peuten eqs 8–9 ratio was inverted — paper reads (m̄*/m̄)^{2δ}; Bianchini A1 prefactor is μ^{−δ} and the paper itself omits the ^{1/2} on the E_γ ratio, confirmed via its own A2). **Task 1 parity**: Engine A ≡ canonical reference LIMEPY (SHA ef2a479, `meanmassdef='central'` ⇒ like-for-like, no eq-8/9 translation) to ≤1.5e-4 across 6 configs (single-mass g=1/1.5, 2-comp, B2 4-comp Maschberger, aniso η=0/0.5); gates measured-then-frozen ≤5e-4 (20× under the plan's 1% target); the reference runs ONLY in a pinned-env subprocess (py3.11/numpy1.26.4/scipy1.11.4 — project scipy 1.17 breaks it) with a committed npz cache + skip-if-absent slow test (3 tests). **Task 2 (Anna-reviewed → FIT-FREE)**: equipartition saturation validated by exact statements only — solver≡closed-form A1 to 1.6e-7; HEADLINE μ→0 identity m/(2η)≡derived m_eq=m̄(g+5/2)(g+7/2)/φ̂₀ to 0.11–0.19% (W0=5/7/9); η(m) shape via autodiff (flat 0.026–0.039 at m_min → ½ exactly deep-well). KEY SCIENCE: the DF saturates SHARPER than exponentially (Bianchini A2's quadratic coefficient is NEGATIVE, opposite eq-3's) ⇒ finite-window eq-3 fits recover ~0.5× the derived m_eq (0.50–0.53 over μ≈[0.14,1.7]; →1.0 as the window shrinks) — documented as a literature caution (window-match before comparing; future-work idea: fit the exact differentiable σ(m;g,W₀,m̄) to kinematic data). **Task 3 docs**: NEW theory page `10-theory/spatial-profiles/multimass-equipartition.md` + corrected per-paper notes wired into nav (myst 147 pages, 0 warnings); docstring-only src edit (GZ15 §4 √-update named). **Close-out**: FULL gate **1192 passed** (= main baseline 1189 + 3 parity tests; CLAUDE.md's "1163" is a stale count). meq/zeta stay DEFERRED. **DONE (2026-06-11, merged+pushed): Multicomponent Engine A IFT hardening (H1–H5)** — `find_alpha_for_masses` → hand-rolled `custom_vjp` (adaptive while_loop fwd + reverse-mode implicit-VJP backward), ~9× faster B2 demo, exact grads (AD-vs-FD iso 2e-7), jit(grad) pinned, FD-gate tightened 1e-6; review caught+fixed a Critical jit(grad) regression. **+ Batch B Tasks 1–5 DONE:** B1 cross-engine figure; B2 joint (α,δ,W₀) MLE recovering reference θ̂ bit-for-bit; Task 5 Fisher information panel (kin-only ρ(α,δ)=−0.26 mild, mass channel PINS α 4.7×) + NUTS corner (0 divergences, posterior-mean within 1σ of MLE, flat-in-θ target). **+ equation-verified per-paper notes** (Bianchini2016, Peuten2017 new; GZ15 updated) documenting the multimass methods/theory. Replaced the unrolled 30-iter eigenvalue scan in `find_alpha_for_masses` with a hand-rolled `jax.custom_vjp` (adaptive `while_loop` forward to tol=1e-6 + reverse-mode implicit-VJP backward of the √-map residual R; TWO solvers dispatched on `ra_hat is None`). ~3× faster per `value_and_grad`, flat-in-n_iter, gradients exact (AD-vs-FD iso 2e-7, aniso quadrature-oracle 5e-5; jit(grad)/jit(value_and_grad) pinned). The B2 demo (`demo_delta_recovery.py`) right-sized 4×600→3×300 → **5:42 vs ~50 min (~9×)**, recovers the gold reference θ̂=(α 2.2931±0.0041, δ 0.3972±0.0344, W₀ 4.9900±0.0209) **bit-for-bit**, OVERALL ALL PASS. Review cadence caught + fixed a Critical `jit(grad)` regression (`jnp.asarray(eta)` in a nondiff_argnums slot) at H2, two tasks before the demo that would have hit it. optimistix REJECTED (its JVP adjoint forces diffrax DirectAdjoint → slower). Fresh gates: **full released-core 1188 passed (7:17)**, memory 7/7 PASS, box-corner harness 3/3 (residual 5–8e-7), demo exit 0. Completion doc `.claude-work/MULTICOMPONENT_IFT_HARDENING_COMPLETE.md`. Future IFT candidates noted (Engine B trunc radius, King tidal radius). Commits `dd14713..ccf6ed1` (6) + harness/tests. **DONE (2026-06-10, Anna-approved): DF sampler-fusion micro-batch** (profile-driven: staged RSS decomposition showed construction ~1.7 GB dominates, eager draw +0.63 GB): speed-CDF tables cached at construction + jitted sampling cores (eqx.filter_jit) for LIMEPY/King/Michie standalone DFs — warm aniso draws **0.35 s → 15 ms** at N=2e4 (3–16 ms across N; 55× vs pre-batch quadrature at N=1e4), draw-chain peak +0.63→+0.20 GB, stage 2.37→2.02 GB, `limepy_df_aniso` gate tightened 3.0→2.5 GB, small-N build penalty eliminated; lax.map chunking of the table draw evaluated and rejected on measurement. Fidelity re-verified on the jitted path (identical KS/moment/β, ALL PASS). Gate: **1163** released-core (unit 895; fast 1121). Remaining DF-stage floor is model CONSTRUCTION (~1.8 GB, paid once per model) — future target if ever needed. **DONE (2026-06-10, Anna-approved): Batch A merged to local main** (fast-forward to `2263948`; FAST gate re-verified on merged main 1108; main NOT pushed — 22+ ahead of origin, awaits Anna's go) **+ the performance-evidence addendum** (Anna-requested, plan-approved): `scripts/benchmark_batch_a.py` measures BOTH code versions live (pre-batch worktree at `0dd1cd9` vs HEAD) — wall-clock: PE 2.4× faster at N=8k (0.169→0.070 s), N=2e4 feasible at 0.216 s (dense: 20.9 GB, infeasible), honest small-N regressions measured (PE 3.5× slower at N=2k; DF table-build dominates below ~5k); memory: 16× at N=8k (gated ≥5×); fidelity figures table-vs-quadrature ALL PASS (KS ≤ 6.7e-3); NEW `50-validation/performance-memory.md` page (myst 144 pages, 0 warnings) + King/Michie fidelity embeds. The multi-component cluster arc is DONE end-to-end: Engines A (LIMEPY DF tables) + B (density-defined Eddington) merged to main (`4f4f259`), the post-merge code-review fix batch landed (`0dd1cd9`), and **Batch A (memory + consolidation + docs)** is complete — blocked+remat'd virial kernels (forward AND grad O(block·N)), table-routed standalone LIMEPY/King/Michie DFs (quadrature kept as the `speed_method="quadrature"` oracle), `progenax.numerics` consolidation (cumtrap ×5 + inverse-CDF draws), `_EngineAState` grouping (A-only access on B raises informative AttributeError), shared angular conditional in `kinematics/_speed_kernels.py`, `lax.map(batch_size=2048)` bounded oracles, and the docs metrics refresh. **Memory headline (measured, `scripts/profile_cluster_memory.py`, all PASS): virial kernel at N=2e4 32.8 GB → 0.41 GB; standalone anisotropic DF sampling 10.9 GB → 2.3 GB; Engine A/B builds at N=1e5 ≤ 2.5 GB.** Gate: released-core **1150 passed** (unit 882 / integration 34 / validation 234; fast gate 1108). **Merge/push of `fix/memory-and-consolidation` awaits Anna's go.** PRIOR (collapsed; superseded-arc detail in `docs/plans/` + `.claude-work/`): **UNIFIED MULTI-COMPONENT CLUSTER REDESIGN on branch `feat/multimass-limepy-equilibrium`.** An 8-dimension adversarial review of the two-component/segregation code found ONE real physics bug (`generate_two_component_cluster` feeds the FULL mass to each sub-population's velocity DF + no COM/virial rescale → non-self-consistent for unequal-mass IMFs; silent because tests use `jnp.ones(N)`) + an ad-hoc API; the `lambda_seg` blend + `energy_sorted_segregation` are sound. **Decision (Anna, pre-launch, no backwards-compat):** build ONE unified differentiable `MultiComponentCluster` (eqx.Module → ICResult) spanning mass-seg · GC 1G/2G · binaries-vs-singles · halo+core, via TWO engines — **Engine A** (DF-defined lowered-isothermal/LIMEPY, mostly built) + **Engine B** (density-defined Eddington/OM for Plummer/EFF/King). Keep `energy_sorted_segregation` as a labeled PRIMORDIAL generator; **retire the `lambda_seg` blend**. Subsume+retire `MultiMassLIMEPY` (mass-seg → a `.from_mass_segregation()` constructor). Plan: `~/.claude/plans/we-are-continuing-progenax-s-crispy-rain.md`. KEY INSIGHT: a component's single free Engine-A scale is the velocity-scale ratio `w_j=s_j/s` with `rescale_j=w_j^-2`; mass-seg is `w_j=μ_j^-δ`. **DONE: Phase 1a** — extracted `solve_multicomponent_limepy(alpha_j, rescale_j, W0, g, ra_hat_j=None)` core (direct per-component scales); `solve_multimass_limepy` is now a thin mass-wrapper (bit-identical, verified); +5 TDD unit tests; committed `47dc24a`. **DONE: Phase 1b (commit `13de746`, NOT pushed)** — `progenax.cluster.multicomponent.MultiComponentCluster` (eqx.Module): physics flows from `(w_j, rescale_j=w_j^-2, ra_hat_j)`, `m_j` decoupled (labels only); `.from_components` (general: GC 1G/2G, halo+core) / `.from_mass_segregation` (w_j=μ_j^-δ) / `.from_imf`; `sample_cluster → ICResult` with NEW optional `component_id` field; the NEW unequal-mass two-population regression passes (theory Q_j=0.5±2e-3, sampled global ±0.04, per-component ±0.07); `MultiMassLIMEPY` RETIRED (tests ported, validation file + both `validate_multimass_*` scripts rewired and re-run PASS, exports updated). Released-core **1065 passed, exit 0** (1056 at 1a − 14 ported-out + 23 new). **DONE: speed fixes** (commit `3f15479`): pytest-xdist + XLA thread caps + `slow` markers (FAST/FULL gates in CLAUDE.md); solver `rho_j_grid` reuse; jitted sampler core. **DONE: Phase 1.5 — LIMEPY DF tables** (Anna-approved plan `docs/plans/2026-06-09-limepy-df-tables-phase15.md`, subagent-driven with per-task independent review; commits `dda2887..f69f36d` + close-out, NOT pushed): three differentiable table primitives in `profiles/limepy_tables.py` (`AnisoDensityTable` (√W, asinh p) cubic-Lagrange; `SpeedCDFTable` 256×256; `AnisoSpeedCDFTable` 192×48×192 speed marginal with EXACT angular conditional) replace pointwise quadrature in the coupled-Poisson RHS + both sampler branches. **Measured: aniso construction 957→170 ms (5.6×); sampling at N=1e5: iso 67×, aniso 21.7× (0.48/2.9 µs/star); aniso small-N break-even ~3k stars (documented).** Every approximation budget-asserted with the exact quadrature retained as a selectable oracle (`aniso_method="quadrature"`) — density 6.05e-6 (1e-5 budget), solve |Δψ|≤1.93e-4 over 3 configs, CDF 4.5e-5, speed moments ≤0.28%/1.5% vs DF quadrature, β(r)≡DF <0.06 unchanged, AD-vs-FD 2.15e-4 + table-AD≡quadrature-AD 2.1e-4; `component_virial_ratios` deliberately quadrature-ONLY (oracle independence: table-built model proves Q_j=0.5001±1.5e-4). Reviews caught+fixed 3 real bugs pre-merge (W≤0 NaN gradient at the truncation boundary; high-g (≥2.5) CDF-row underflow corrupting outer-shell speeds, now gated g∈{0,2.5,3.5}; the plan's own √(2π) test bug). `scripts/validate_df_tables.py` ALL PASS; completion doc `.claude-work/TASK_1.5_DF_TABLES_COMPLETE.md`. Released-core at 1.5 close-out: 1098 (gate evidence in `.claude-work/TASK_1.5_DF_TABLES_COMPLETE.md`, commit `a85bada`). **DONE: Phase 1c (commits `b1caabb`, `8cda23d`, + 1c-3; NOT pushed)** — the Anna-approved OUTRIGHT DELETION (−3.4k LOC): `populations.py` (the full-mass-to-each-DF physics-bug module) + `cluster/core.py` (string-dispatch `generate_cluster_ic`/`ClusterState`/`SpatialStructureParams`/`MassSegregationLayer`) + `cluster/validation.py` + all `lambda_seg` machinery/scripts RETIRED; `test_segregation_equilibrium_physics.py` PORTED onto the primordial generator via the protocol API (λ=1 endpoint Q_j drift<0.08, orbit-reuse regression, Λ_MSR rises); `validate_cluster_ic.py` §2 rewired to `MultiComponentCluster` (cold w=0.7 median r 1.50 pc < hot 8.36 pc, oracle Q_j=[0.5,0.5], ALL PASS); `mass_segregation.py` relabeled PRIMORDIAL (documented departure from Baumgardt+2008 random per-bin draw; wrong MNRAS 384,1231 cite died with populations.py; dead `_mcluster_partial_shuffle` removed); **eqx migration complete — src/ is dataclass-free** (`ICResult`/`Systems`/`Stars`/`TotalMass`/`AnalyticalIC`/`RotationParams`/`VelocityModel` now eqx.Modules; physical scalars stay differentiable leaves, control-flow config static; jit-over-ICResult verified); `compute_profile_potential` docstring fixed (SPECIFIC potential, [length]²/[time]²). Full gate **1064 passed in 438s** (unit 802/integration 34/validation 228; fast gate 1029; exact count reconciliation 1098−28−7−4+4+1). Phase-3 note: 16 docs/website source pages still reference the deleted API (inventoried; the docs sweep is Phase 3 scope). **DONE: Phase 2 — Engine B (density-defined Eddington), commits `292b6b0..08e328b` + close-out, NOT pushed:** `MultiComponentCluster.from_density_profiles(profiles, mass_fractions, m_j, r_a_j=None, r_t=None)` — Plummer/EFF/King DENSITY components in ONE shared self-consistent potential (no ODE: one cumtrap Poisson pass, new `profiles/density_poisson.py` with derived-domain provenance + King-override conflict rule), per-component DFs via the generic `eddington_invert` (extracted from the EFF table BIT-IDENTICALLY, pinned; + Plummer E^{7/2} (untruncated zero point) and EXACT truncated closed-form oracles, 1.06e-4/8.86e-6), optional per-component OM `r_a_j`, f_j≥0 realizability gate (concrete builds refuse naming the component + f_min + remedy; traced builds store the `f_min_j` diagnostic). Anchors ALL PASS (`scripts/validate_multicomponent_eddington.py`, 11/11): King A-vs-B (TWO independent engines: σ_1d(r) dev 3e-4, radial KS 2e-4, both gates 0.02), EFF(γ=5)≡Plummer twins, halo+core headline (theory Q_j=0.5±4e-4; sampled global Q=0.4976 UNSCALED at N=30k; sampled per-component Q_j ≡ the hybrid predict-the-offset quadrature expectation — the hard-truncated halo's plateau 0.4953 predicted / 0.4947±0.0014 sampled (18 seeds) is VERIFIED truncation-edge physics, gated on the prediction, Anna-approved), OM β_halo(r)≡r²/(r²+r_a²) (max dev 0.028 < 0.05), DF-density fidelity 1.06e-3 (OM) / 2.4e-4 (iso) vs 5e-3 gate, AD≡FD 5.6e-9/7.8e-7/2.0e-8 (r_h, mass-fraction t, r_a). THREE physics findings: (1) the E^{7/2} oracle needs the UNTRUNCATED zero point (truncation shifts E by c=M/√(r_t²+a²), ~30% at 0.1Ψ₀); (2) the plan's default a_EFF=0.4 core is GENUINELY unrealizable (f_min −0.20; gate flip measured a∈(0.65,0.68); headline uses 0.8); (3) hybrid-sampling Q_j plateau predicted+verified. Plus the King dW/dr staircase bug (jnp.gradient of interpolated ψ → Eddington ringing, f_min −0.679) fixed via King's own Poisson identity → +5.1e-7. Full gate **1101 passed in 511s, exit 0** (unit 833, integration 34, validation 234; fast gate 1059); close-out `.claude-work/PHASE_2_ENGINE_B_COMPLETE.md`. **DONE: merge + push (2026-06-10, Anna-approved):** `feat/multimass-limepy-equilibrium` fast-forwarded into `main` (`18f116b..4f4f259`), merged-main gate re-verified **1101 passed**, pushed; local branch deleted (stale remote branch copy left — remote deletion not authorized). **DONE: Phase 3 docs sweep (2026-06-10):** comprehensive website update — `lowered-model-family.md` flipped roadmap→implemented Engine A methods page; populations section rewritten around the two-engine `MultiComponentCluster` (+ NEW `10-theory/populations/eddington-engine.md`); API reference regenerated (PUBLIC_MODULES −populations +cluster/diagnostics/dynamics; 11 module pages, 175 symbols; `30-api/populations.md` deleted); NEW `50-validation/{multimass-equilibrium,engine-b-eddington}.md` with Measured tables + figures (stale mass-seg numbers corrected by re-running validate_cluster_ic.py: Λ_MSR 4.74, ρ(m,E)=−0.92); all 16 inventoried stale-API pages fixed (dev-log: dated retirement admonitions, history preserved); myst build clean (142 pages, 0 warnings). **DONE: /code-review of the merged arc (7 finder angles + verification):** 3 CONFIRMED post-merge correctness bugs — (1) NaN-blind realizability gate (`nan < -1e-3` False; r_a_j=0 → silent all-NaN model), (2) no mass_fractions/m_j/r_a_j length validation vs profiles (phantom components via clamped gathers; potential built from wrong total mass), (3) absolute Ψ thresholds in `sample_speed_from_f_table` are scale-dependent (r_h≳1e4 pc → exact-zero speeds; 80% of stars at r_h=1e6) — plus `total_density()`/`rescale_j` silent-NaN on Engine B, `is_aniso=False` lie for OM-B models, `r_t_provenance` static field embedding float values (per-r_t recompile + vmap break), and a consolidation list (DF speed-moment kernel ×5 copies w/ inconsistent n_w, cumtrap Poisson pass ×3, King density/derivative split, union-by-NaN → _EngineAState grouping). All documented honestly in the new docs limitations admonitions. **RESUME AT: the code-review fix batch (findings 1–3 first; then 4–6; consolidations after), then cross-engine validation figures + science demos.** HITL: push only when Anna says. @@ -15,7 +16,7 @@ CAMPAIGN R2/R3 (2026-06-09, pushed): **analytical-test-cases** (9b96bc5): script CAMPAIGN R2 EXTENSION (2026-06-08, pushed): 3 more modules to pub standard, each 5 build-verified figures + Measured-column tables. **imf-statistics** (c0b8cfe): scripts/validate_imfs.py rewrite onto _plotstyle — recovered slopes via sample-MLE (Salpeter 2.347±0.004, Kroupa 0.294/1.297/2.301), PDF/CDF coverage+KS, mean_mass() vs 200k log-grid ref (≤1.2e-4) with an adversarial linear-grid failure-mode panel, inference-layer gradient (rel 1e-5); fixed classic.md stale Chabrier params (0.22/0.57 were the *system* IMF → code defaults single-star 0.08/0.69), wrong class name, non-existent API kwargs. **binary-imf** (d43e0c5): scripts/validate_binaries.py — Kepler/orbit oracles (machine-precision), Moe+17 q-dist+KS (mass-dependent twin excess 15.2%→6.3%), HEADLINE regenerated "confidently wrong" via fast differentiable MLE (no MCMC): naive single-star fit 17.8σ-wrong at N=1e5 (α̂=2.21 vs 2.30) while binary-aware recovers 2.298 (0.4σ); bias∝f_b →0 at f_b=0; grad 8e-6; fixed stale binary_aware_validation.py script name. **environment-IMF** (9c358d9, harden-then-validate): NEW tests/validation/test_environment_physics.py (12 tests vs published Marks+2012 Table 1 GC α3 + Table 4 slopes) → released-core 866→878; scripts/validate_environment.py — GC anchors (NGC 7078 most top-heavy 0.844 vs pub 0.76), Fundamental Plane (density beats metallicity 7:1), low-mass slopes exact, HONEST Marks-vs-Jeřábková cross-model (both top-heavy mechanism but ~0.7 zero-point divergence reported not asserted away), grad 9e-8. All myst-build-verified (figs in _build/site/public + page JSON); new 50-validation/environment-imf.md wired into nav+dashboard+audit-report. CAMPAIGN SUMMARY (2026-06-08): 6 modules at publication standard, each 5 build-verified figures + Measured-column tables: Plummer, King(retrofit), EFF, Michie, rotation/OM-anisotropy, substructure/CW04 Q; diagnostics (Λ_MSR + energy-sorted segregation) already done. Shared scripts/_plotstyle.py; LIMEPY reframed across 13 pages + roadmap (10-theory/spatial-profiles/lowered-model-family.md); validate_profiles.py deleted. Released-core 830→861 (+31 validation tests, all collect clean, build clean). Commits f964ebf, 4db6dbb, ee5f299, fcbfa32, 66cedaf on local main. Findings handled honestly: King W0 differentiable (corrected docs); EFF γ=5≡Plummer exact; Michie β suppressed below Osipkov-Merritt (validated vs DF moment oracle); OM stretch realizes β exactly (the Michie contrast); q_approx kNN-surrogate gradient (median 0.9%, cell-boundary worst-case). **rotation/anisotropy DONE (2026-06-08)**: new tests/validation/test_rotation_anisotropy_physics.py (10 tests, released-core 842→852) + scripts/validate_rotation_anisotropy.py (5 figs) + 50-validation/rotation-om-anisotropy.md. Solid-body v_φ=ΩR slope exact (resid 3e-14), L_z budget exact, differential curve <1e-6; OM anisotropy β(r)=r²/(r²+r_a²) realized EXACTLY for Plummer+EFF (dev 0.024/0.016) via Merritt-1985 direction-stretch — the clean CONTRAST to Michie's suppressed β. Ω/v_peak/r_a all differentiable (3e-9/7e-9/8e-6). **Michie DONE (2026-06-08)**: new tests/validation/test_michie_physics.py (12 tests, released-core 830→842) + scripts/validate_michie.py (5 figs) + 50-validation/michie-anisotropy.md. KEY FINDING: Michie-King β(r) is SUPPRESSED below the pure Osipkov-Merritt ceiling r²/(r²+r_a²) because the lowering term −exp(−J²/2r_a²σ²) breaks the f(Q) form — validated the sampler against the model's OWN DF β (2nd-moment oracle), not OM. γ/W0/r_c/M all differentiable. Per-module: status check → brainstorm plot set (HITL) → build pub figures → embed → BUILD-VERIFY → dashboard. **Plummer + King + EFF (2026-06-08) all done**: dedicated scripts/validate_{plummer,king,eff}.py on shared scripts/_plotstyle.py, each 5 pub figures printing expected-vs-measured; validation tables carry Measured columns; validate_profiles.py deleted (superseded). EFF: fixed false page claims (no γ>3 ValueError, no from_rh r_h-mapping, no KS test, "γ differentiable unlike King's W0" was wrong — both differentiable); γ=5→Plummer exact; γ AD-grad 2e-11. Entry point: docs/plans/2026-06-08-validation-audit-and-methods-figures-spec.md (per-module: brainstorm plot set → build pub-quality figures → embed on 50-validation page → BUILD-VERIFY render → update index dashboard). Done so far: **Plummer (2026-06-08, 5 pub figures)** + **King (32 tests, 5 figs)** + two-component + mass-segregation pages current+figure-bearing+build-verified; Λ_MSR validated; Tier C built. NEW shared scripts/_plotstyle.py (Okabe-Ito + ApJ rcParams); King refactored onto it. Plummer: dedicated scripts/validate_plummer.py (density+CDF, velocity-equilibrium, Beta(3/2,9/2), gradient AD-vs-FD, isotropy), each printing expected-vs-measured; Plummer+King validation tables now carry a Measured column (real regenerated numbers, fixed old overclaims). LIMEPY reframed across 13 pages (cite Gieles2015, own differentiable lowered-model family) + new roadmap 10-theory/spatial-profiles/lowered-model-family.md. EFF figure rewrite in validate_profiles.py still pending the EFF audit (validate_profiles.py is now Plummer-stale — its validate_plummer() superseded by scripts/validate_plummer.py; clean it during EFF). -blocker: none +blocker: CI red on main (jaxstroviz path dep not checked out in tests.yml) — blocks the planned PR-when-green workflow; see audit R1. due: ## Current focus diff --git a/docs/plans/2026-06-12-prerelease-audit-fixes.md b/docs/plans/2026-06-12-prerelease-audit-fixes.md new file mode 100644 index 00000000..149031a3 --- /dev/null +++ b/docs/plans/2026-06-12-prerelease-audit-fixes.md @@ -0,0 +1,1158 @@ +# Pre-Release Audit Fixes Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. Execute in batches (one phase = one batch); STOP at every ⛔ CHECKPOINT and report to Anna — she is the non-negotiable human-in-the-loop and must approve before you continue. Never weaken a test to make it pass; if a fix changes expected physics, derive the new expectation analytically inside the test. + +**Goal:** Fix and validate every release blocker from `docs/reviews/2026-06-11-prerelease-adversarial-audit.md` (findings R1–R10 + high-value mediums), moving progenax from Beta to Release Candidate. + +**Architecture:** Five phases ordered by dependency: (0) make CI green and enforcing, (1) TDD fixes for the five confirmed science bugs, (2) JAX/API hardening, (3) documentation honesty sweep, (4) packaging/metadata. Each science fix follows strict RED→GREEN: write the failing test that encodes the *paper's* convention or a hi-res numerical reference, watch it fail against current code, then fix. + +**Tech Stack:** JAX/Equinox (100% JAX-native core — never numpy/scipy in `src/progenax` except the documented `diagnostics` carve-out), pytest + pytest-xdist, uv, diffrax, GitHub Actions. + +**Source of truth:** `docs/reviews/2026-06-11-prerelease-adversarial-audit.md`. Finding IDs (R1–R10, S*, J*, A*, T*, D*) refer to that report. Read its §3 and §4 before starting. + +--- + +## Ground rules for the executing session + +- **Branch:** create `fix/prerelease-audit` off `main`. Commit per task. Do NOT push or merge without Anna's explicit go (HITL). +- **Test commands** (from `progenax/CLAUDE.md`): + ```bash + # FAST gate (inner loop, ~4 min). Baseline before any change: 1145 passed. + XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" \ + env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit tests/integration tests/validation -q -m "not slow" -n auto + # FULL gate (phase close-out, ~9 min). Baseline: 1192 passed. + XLA_FLAGS="--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1" \ + env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit tests/integration tests/validation -q -n auto + # Single test during TDD: + env -u VIRTUAL_ENV uv run --no-sync pytest :: -v + ``` +- **Verify RED before GREEN.** Every new test must be run and shown failing against unmodified code before the fix lands. If a test passes immediately, the test is wrong — stop and rethink. +- **Units:** STELLAR (M☉, pc, Myr); G explicit everywhere; never hardcode G. +- **JAX policy:** `jnp` only in core, `lax.scan`/`fori_loop` (no `while_loop`), the double-`where` `exp_safe` pattern for removable singularities (reference implementation: `src/progenax/imf/differentiable.py:47-54`). +- After each phase: run the FULL gate, report results + diffs to Anna at the ⛔ CHECKPOINT. + +--- + +# Phase 0 — CI: make the suite enforced again (R1, R8, T4, M5) + +CI has been red on `main` since 2026-06-10: every job dies in ~25 s at `uv sync`/`uv lock --check` with `Distribution not found at: .../jaxstroviz`, because `pyproject.toml` declares `[tool.uv.sources] jaxstroviz = { path = "../jaxstroviz" }` (and `uv.lock` pins it, dragging `fluxax` transitively) but `.github/workflows/tests.yml` only checks out `progenax` + `jaxstro`. Nothing in this repo is currently enforced anywhere. + +### Task 0.1: Check out the missing sibling packages in all CI jobs + +**Files:** +- Modify: `.github/workflows/tests.yml` (jobs `lock-check`, `released-core`, `experimental` — every job that runs `uv lock --check` or `uv sync`) + +**Step 1: Identify every sibling path dependency the lock graph needs** + +```bash +grep -n "editable" uv.lock | sort -u +grep -n "path =" pyproject.toml +``` +Expected: `../jaxstro`, `../jaxstroviz`, and (transitively via jaxstroviz) possibly `../fluxax`. Note exactly which appear — the checkout list below must match. + +**Step 2: Add checkout steps** + +In `tests.yml`, immediately after the existing `Checkout jaxstro` step in **each** of the three jobs, add (repeat per missing sibling found in Step 1): + +```yaml + - name: Checkout jaxstroviz (path dependency, sibling layout) + uses: actions/checkout@v4 + with: + repository: jaxstro/jaxstroviz + path: jaxstroviz + - name: Checkout fluxax (transitive path dependency of jaxstroviz) + uses: actions/checkout@v4 + with: + repository: jaxstro/fluxax + path: fluxax +``` + +If a repo name differs (check with `gh repo list jaxstro`), use the real name. If `fluxax` is NOT in `uv.lock`, omit it. + +**Step 3: Verify locally that the lock graph is the only problem** + +```bash +env -u VIRTUAL_ENV uv lock --check +``` +Expected: passes locally (siblings exist here). The CI failure is purely the missing checkout. + +**Step 4: Commit, push the branch, and watch CI** + +```bash +git add .github/workflows/tests.yml +git commit -m "fix(ci): check out jaxstroviz (+fluxax) siblings — unbreaks uv sync, CI red since 2026-06-10 (audit R1)" +git push -u origin fix/prerelease-audit # pushing the FIX BRANCH to run CI is authorized for this task only +gh run watch --exit-status || gh run list --branch fix/prerelease-audit --limit 3 +``` +Expected: `lock-check` passes; `released-core` and `experimental` shards run actual tests. If a shard fails on a *test* (not on sync), STOP and report — that's new information (tests that haven't run in CI for days may be stale). + +### Task 0.2: Make the LIMEPY reference cache a hard failure in strict mode (T4, H2) + +`tests/validation/test_limepy_reference_parity.py:62-64` silently `pytest.skip`s when `validation/data/limepy_reference/*.npz` is absent — a packaging mistake would convert the strongest external-reference gate into a silent pass. + +**Step 1: Write the failing test** — Create `tests/validation/test_strict_refs.py`: + +```python +"""Strict-mode guard: reference caches must exist when PROGENAX_STRICT_REFS=1.""" +import os +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +LIMEPY_CACHE = REPO_ROOT / "validation" / "data" / "limepy_reference" + + +@pytest.mark.skipif( + os.environ.get("PROGENAX_STRICT_REFS") != "1", + reason="strict reference mode not requested", +) +def test_limepy_reference_cache_present(): + """In strict (nightly/release) mode the committed LIMEPY cache must exist. + + Guards against a checkout/packaging mistake silently disabling the + reference-parity suite (audit T4/H2). + """ + npz_files = sorted(LIMEPY_CACHE.glob("*.npz")) + assert npz_files, f"no LIMEPY reference .npz under {LIMEPY_CACHE}" +``` + +**Step 2: Run it RED then GREEN** + +```bash +PROGENAX_STRICT_REFS=1 env -u VIRTUAL_ENV uv run --no-sync pytest tests/validation/test_strict_refs.py -v +``` +Expected: PASS (cache is committed). Now verify it actually guards: temporarily `mv validation/data/limepy_reference /tmp/`, rerun (expected FAIL), then `mv` it back and rerun (PASS). Show both outputs. + +**Step 3: Convert the silent skip in the parity test to a strict-mode failure** + +In `tests/validation/test_limepy_reference_parity.py`, replace the bare `pytest.skip(...)` at ~line 62-64 with: + +```python + if os.environ.get("PROGENAX_STRICT_REFS") == "1": + pytest.fail(f"reference-LIMEPY cache absent at {cache_path} (strict mode)") + pytest.skip(f"reference-LIMEPY cache absent at {cache_path}") +``` +(add `import os` if missing). + +**Step 4: Commit** + +```bash +git add tests/validation/test_strict_refs.py tests/validation/test_limepy_reference_parity.py +git commit -m "test: PROGENAX_STRICT_REFS=1 turns missing reference caches into failures (audit T4)" +``` + +### Task 0.3: Add the nightly/release full-physics CI lane (R8) + +The headline trust anchors (LIMEPY parity, multimass equilibrium, Engine A-vs-B, Engine B AD-vs-FD — 47 `@pytest.mark.slow` tests) currently run nowhere but Anna's laptop. + +**Files:** +- Create: `.github/workflows/physics-validation.yml` + +```yaml +name: physics-validation + +on: + schedule: + - cron: "17 9 * * *" # nightly + workflow_dispatch: + push: + tags: ["v*"] + +env: + JAX_ENABLE_X64: "1" + XLA_PYTHON_CLIENT_PREALLOCATE: "false" + XLA_PYTHON_CLIENT_ALLOCATOR: platform + OMP_NUM_THREADS: "1" + PROGENAX_STRICT_REFS: "1" + +jobs: + slow-physics: + runs-on: ubuntu-latest + timeout-minutes: 90 + steps: + # SAME checkout block as tests.yml (progenax + jaxstro + jaxstroviz [+ fluxax]) + # ... copy verbatim from tests.yml after Task 0.1 ... + - name: Setup uv + uses: astral-sh/setup-uv@v6 + with: + python-version: "3.13" + - name: Sync from lockfile (dev + experimental) + working-directory: progenax + run: uv sync --locked --extra dev --extra experimental + - name: Slow physics-validation suite (the advertised trust anchors) + working-directory: progenax + run: uv run --no-sync pytest tests -m slow -q + - name: Strict reference guard + working-directory: progenax + run: uv run --no-sync pytest tests/validation/test_strict_refs.py -q +``` + +**Verify:** `gh workflow run physics-validation.yml --ref fix/prerelease-audit` then `gh run watch`. The slow suite takes ~30–60 min on the 2-core runner; if it OOMs, shard it `-m "slow"` per tier like `tests.yml` does. Commit: + +```bash +git add .github/workflows/physics-validation.yml +git commit -m "ci: nightly + release-tag lane running the slow-marked headline validations (audit R8)" +``` + +### Task 0.4: Delete the ghost cache directories (M5) + +```bash +ls src/progenax/gravoturb/ src/progenax/cluster/fdf_density/ # confirm: ONLY __pycache__/.mypy_cache debris, zero .py +git ls-files src/progenax/gravoturb src/progenax/cluster/fdf_density # confirm: empty (nothing tracked) +rm -rf src/progenax/gravoturb src/progenax/cluster/fdf_density +``` +Nothing to commit (untracked). Note in the phase report that this was done. **Do NOT touch `src/progenax/cluster/turbulence.py`** — it is live core physics consumed by the environment-IMF (audit M5 note). + +### ⛔ CHECKPOINT 0 +Report to Anna: CI run URLs (green or what failed), the strict-refs RED/GREEN evidence, and FAST gate output (`1145 passed` + the 1 new strict test when env var unset → skipped). **Wait for approval before Phase 1.** + +--- + +# Phase 1 — Confirmed science bugs, strict TDD (R3, R4, R5, R6, R10) + +All five were verified by execution during the audit (see report §3 + the orchestrator re-verification). Each task: failing test first. + +### Task 1.1: Moe & Di Stefano F_twin normalization (R3) — **the most important fix in this plan** + +**The bug:** `MoeDiStefano2017Full.pdf` (`src/progenax/imf/binary/moe_di_stefano.py:279-289`) mixes the twin block with weight `ft` against the whole q ∈ [0.1, 1] population. The paper (MD17 p.5 + Fig. 2 caption; PDF held at `docs/core-papers/Moe_2017_ApJS_230_15.pdf`) defines F_twin as the excess-twin fraction **relative to q > 0.3 companions**. Measured impact: realized paper-convention F_twin = 0.367 vs Table-13 value 0.300 at (M1=1, logP=1) — +22% twin overweight, with a matching deficit at q < 0.3. + +**The math of the fix:** build the unnormalized mixture as the two-slope power law plus a twin block of mass `ft/(1−ft) · I_B` on [0.95, 1] (where `I_B` is the power-law mass on [q_break, 1]), then normalize the whole thing over [q_min, 1]. Check: excess-twin fraction among q > 0.3 = `(ft/(1−ft)·I_B) / (I_B + ft/(1−ft)·I_B)` = `ft` exactly. ✓ + +**Files:** +- Modify: `src/progenax/imf/binary/moe_di_stefano.py:279-289` (`pdf`; add a `_components` helper) +- Modify: `tests/unit/imf/test_moe_full.py` (the `mean(q>0.95) > 0.28` assertion at ~lines 92-100 encodes the OLD convention — replace with a derived expectation, do not just loosen) +- Test: `tests/unit/imf/test_moe_full.py` (new tests appended) + +**Step 1: Write the failing tests** + +Append to `tests/unit/imf/test_moe_full.py`: + +```python +class TestFTwinPaperConvention: + """MD17 p.5 / Fig.2: F_twin = excess-twin fraction of q > 0.3 companions. + + Audit finding R3: the pre-fix mixture realized F_twin = 0.367 instead of + 0.300 at (M1=1, logP=1). These tests pin the paper convention at four + Table-13 nodes via quadrature of the pdf's mixture components. + """ + + NODES = [(1.0, 1.0), (1.0, 3.0), (3.2, 1.0), (12.0, 3.0)] # (M1 [Msun], logP) + + def test_f_twin_excess_over_q_gt_03(self): + md = MoeDiStefano2017Full() + qs = jnp.linspace(md.q_min, 1.0, 200_001) + for m1, logP in self.NODES: + mass = jnp.asarray(m1) + period = jnp.asarray(10.0**logP) + ft_table = float(md.f_twin(period, mass)) + p_pl, p_twin = md._components(qs, mass, period) + mask = qs >= 0.3 + twin_mass = float(jnp.trapezoid(jnp.where(mask, p_twin, 0.0), qs)) + total_gt03 = float( + jnp.trapezoid(jnp.where(mask, p_pl + p_twin, 0.0), qs) + ) + realized = twin_mass / total_gt03 + assert abs(realized - ft_table) < 2e-3, ( + f"(M1={m1}, logP={logP}): realized paper-convention F_twin " + f"{realized:.4f} != Table 13 value {ft_table:.4f}" + ) + + def test_pdf_still_normalized(self): + md = MoeDiStefano2017Full() + qs = jnp.linspace(md.q_min, 1.0, 200_001) + for m1, logP in self.NODES: + p = jax.vmap( + lambda q: md.pdf(q, jnp.asarray(m1), jnp.asarray(10.0**logP)) + )(qs) + Z = float(jnp.trapezoid(p, qs)) + assert abs(Z - 1.0) < 1e-3 + + def test_sample_matches_pdf_twin_fraction(self): + """Sampled P(q >= 0.95) must match the quadrature of the FIXED pdf.""" + md = MoeDiStefano2017Full() + n = 200_000 + key = jax.random.PRNGKey(0) + m1 = jnp.full((n,), 1.0) + periods = jnp.full((n,), 10.0) + q = md.sample(key, m1, periods) + qs = jnp.linspace(md.q_min, 1.0, 200_001) + p = jax.vmap(lambda qq: md.pdf(qq, jnp.asarray(1.0), jnp.asarray(10.0)))(qs) + expected = float( + jnp.trapezoid(jnp.where(qs >= 0.95, p, 0.0), qs) + / jnp.trapezoid(p, qs) + ) + observed = float(jnp.mean(q >= 0.95)) + assert abs(observed - expected) < 0.01 # shot noise ~0.001 at n=2e5 +``` + +**Step 2: Run RED** + +```bash +env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit/imf/test_moe_full.py::TestFTwinPaperConvention -v +``` +Expected: `test_f_twin_excess_over_q_gt_03` FAILS — first with `AttributeError: _components` (helper doesn't exist), and after you add the helper as a pure refactor of the CURRENT (buggy) mixture it must fail with `realized 0.367 != 0.300` at the (1.0, 1.0) node. Show both failures. + +**Step 3: Implement** + +In `moe_di_stefano.py`, replace `pdf` (lines 279-289) with: + +```python + def _components(self, q, masses, periods): + """Normalized (power-law, twin-excess) mixture components. + + Paper convention (MD17 p.5, Fig. 2): F_twin is the excess-twin fraction + of q > 0.3 companions. The unnormalized twin block therefore carries + mass ft/(1-ft) * I_B (I_B = power-law mass on [q_break, 1]), so that + twin/(twin + I_B) = ft exactly after joint normalization. + """ + gs, gl, C, I_A, I_B = self._two_slope(periods, masses) + ft = self.f_twin(periods, masses) + p_lo = jnp.power(q, gs) + p_hi = C * jnp.power(q, gl) + p_pl_unnorm = jnp.where(q < self.q_break, p_lo, p_hi) + ft_safe = jnp.minimum(ft, 0.95) # Table 13 max is ~0.3; guard 1/(1-ft) + twin_mass = ft_safe / (1.0 - ft_safe) * I_B + twin_unnorm = twin_mass * jnp.where( + (q >= 0.95) & (q <= 1.0), 1.0 / 0.05, 0.0 + ) + Z_tot = I_A + I_B + twin_mass + in_range = (q >= self.q_min) & (q <= 1.0) + p_pl = jnp.where(in_range, p_pl_unnorm / Z_tot, 0.0) + p_twin = jnp.where(in_range, twin_unnorm / Z_tot, 0.0) + return p_pl, p_twin + + def pdf(self, q, masses, periods): + """Conditional mass-ratio pdf p(q | M1, P), normalized on [q_min, 1]. + + F_twin follows the paper convention: excess-twin fraction of the + q > 0.3 population (NOT of all q > q_min companions — audit R3). + """ + p_pl, p_twin = self._components(q, masses, periods) + return p_pl + p_twin +``` + +Update the class docstring (lines 233-242) to state the q>0.3 convention explicitly. + +**Step 4: Run GREEN + fix the stale assertion** + +```bash +env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit/imf/test_moe_full.py -v +``` +The old `mean(q > 0.95) > 0.28` test may now fail (the corrected solar-node twin share is ≈ 0.28, right at the gate). Replace its magic number with the quadrature-derived expectation (the same pattern as `test_sample_matches_pdf_twin_fraction`). Do NOT simply lower the constant. + +**Step 5: Check downstream consumers + gradients** + +```bash +env -u VIRTUAL_ENV uv run --no-sync pytest tests/unit/imf tests/unit/binaries tests/integration/test_binary_cluster.py tests/unit/imf/test_imf_gradients.py -q +``` +`MoeJointOrbit` and `MoeCompanions` consume this pdf — their distribution tests may have pinned the old convention; update each with derived expectations, documenting the R3 fix in a comment. + +**Step 6: Commit** + +```bash +git add src/progenax/imf/binary/moe_di_stefano.py tests/ +git commit -m "fix(imf): MoeDiStefano2017Full F_twin now follows the paper's q>0.3 convention (audit R3) + +Twin excess was mixed against the full q in [0.1,1] population; MD17 p.5/Fig.2 +define F_twin relative to q>0.3. Realized F_twin at (M1=1, logP=1) was 0.367 +for a Table-13 value of 0.300. Mixture rebuilt as PL + ft/(1-ft)*I_B twin block, +jointly normalized; pinned at four Table-13 nodes." +``` + +### Task 1.2: High-W0 King/Michie CDF core resolution (R4) + +**The bug:** `KingProfile.__init__` (`src/progenax/profiles/king.py:359`) builds the sampling CDF on `r_grid = jnp.linspace(0.0, r_t_arr, n_grid)` with `n_grid=1000`. Since r_t/r_c grows super-exponentially with W0 (131 at W0=9, 548 at W0=12), the core gets ~8 points at W0=9 and ~2 at W0=12. Measured enclosed-mass error at 0.3 r_c: **+18.3% (W0=9), +270% (W0=12)** (production sampler vs hi-res ODE reference, 2×10⁶ samples, shot ~2%). Same pattern at `michie.py:177` and `eff.py:83`. + +**The fix:** sqrt-stretched radial grid `r = r_t · u²` with `u = linspace(0, 1, n_grid)` — spacing dr ∝ √r concentrates points in the core (12+ points inside 0.3 r_c even at W0=12 with n_grid=1000) — plus the non-uniform trapezoid. Differentiability in W0/r_c is preserved (r_grid remains a smooth function of the traced `r_t_arr`). + +**Files:** +- Modify: `src/progenax/profiles/king.py:358-392` (grid + cumulative-mass block in `__init__`) +- Modify: `src/progenax/profiles/michie.py:~177` (same pattern — locate the `linspace(0, r_t, ...)` CDF grid and apply identically) +- Modify: `src/progenax/profiles/eff.py:~83` (same, for uniformity; EFF is currently fine at typical r_t/a but shares the pattern) +- Test: `tests/validation/test_king_physics.py` (new class) + +**Step 1: Write the failing test** — append to `tests/validation/test_king_physics.py`: + +```python +class TestHighW0CoreResolution: + """Audit R4: the linear 1000-pt CDF grid under-resolves the core at W0 >= 9. + + Reference = direct quadrature of rho_hat(psi) on a dense ODE solve + (xi_max=600, n_points=20000) — independent of the profile's internal CDF. + Measured pre-fix errors at 0.3 r_c: +18% (W0=9), +270% (W0=12). + """ + + @pytest.mark.parametrize("W0", [7.0, 9.0, 12.0]) + def test_sampled_core_mass_matches_dense_reference(self, W0): + from progenax.profiles.king import ( + solve_king_profile, + king_lowered_maxwellian_density, + ) + + xi, psi = solve_king_profile(W0, xi_max=600.0, n_points=20_000) + rho = king_lowered_maxwellian_density(jnp.maximum(psi, 0.0)) + integ = rho * xi**2 + M = jnp.concatenate( + [jnp.zeros(1), jnp.cumsum(0.5 * (integ[1:] + integ[:-1]) * jnp.diff(xi))] + ) + prof = KingProfile.from_W0_rc(W0=W0, r_c=1.0) + n = 2_000_000 + pos = prof.sample_positions(jnp.ones(n), jax.random.PRNGKey(3)) + r = jnp.linalg.norm(pos, axis=1) + for r_probe in (0.3, 1.0, 3.0): + m_ref = float(jnp.interp(r_probe, xi, M) / M[-1]) + m_samp = float(jnp.mean(r < r_probe)) + shot = 3.0 / (m_ref * n) ** 0.5 # 3 sigma binomial + tol = max(0.03, shot) # 3% grid budget or shot noise, whichever larger + assert abs(m_samp / m_ref - 1.0) < tol, ( + f"W0={W0}, r={r_probe} r_c: sampled M( 60 s; if so, ALSO add one fast non-slow case (W0=9, r_probe=0.3 only, n=5×10⁵, tol=max(0.05, shot)) so the regression is enforced in the PR lane. + +**Step 2: Run RED** + +```bash +env -u VIRTUAL_ENV uv run --no-sync pytest tests/validation/test_king_physics.py::TestHighW0CoreResolution -v +``` +Expected: FAIL at W0=9 (+~18% at 0.3 r_c) and W0=12 (+~270%); W0=7 passes. + +**Step 3: Implement** — in `king.py.__init__`, replace lines 358-360: + +```python + # Build radial grid for CDF — sqrt-stretched (r = r_t * u^2): spacing + # dr ∝ sqrt(r) concentrates points in the core. A LINEAR grid leaves + # <10 points inside the core at W0 >= 9 (xi_t grows super-exponentially: + # 131 r_c at W0=9, 548 at W0=12), giving +18%..+270% core-mass errors + # (audit R4, measured). Smooth in r_t -> differentiable in W0/r_c. + u_grid = jnp.linspace(0.0, 1.0, n_grid) + r_grid = r_t_arr * u_grid**2 + xi_grid_local = r_grid / r_c_arr +``` + +and replace the uniform-`dr` cumulative-mass block (lines ~386-392) with the non-uniform trapezoid: + +```python + integrand = 4.0 * jnp.pi * r_grid**2 * rho_grid + M_cum = jnp.concatenate([ + jnp.zeros(1), + jnp.cumsum(0.5 * (integrand[1:] + integrand[:-1]) * jnp.diff(r_grid)), + ]) +``` + +(Keep whatever normalization/interp follows; only the grid and the quadrature weights change.) Check whether `progenax.numerics` already exposes a non-uniform `cumulative_trapezoid` — if yes, call it instead of inlining (DRY). + +Apply the identical transformation to the CDF-grid blocks in `michie.py` (~line 177) and `eff.py` (~line 83). Read each block first; the variable names differ. + +**Step 4: Run GREEN + regression** + +```bash +env -u VIRTUAL_ENV uv run --no-sync pytest tests/validation/test_king_physics.py tests/validation/test_michie_physics.py tests/validation/test_eff_physics.py tests/unit/profiles -q +``` +All pass, including the existing W0=7 CDF quadrature gates (`test_cdf_quadrature.py`) — the stretch must not degrade the already-good low-W0 regime. + +**Step 5: Gradient regression** (the grid change touches the differentiable path): + +```bash +env -u VIRTUAL_ENV uv run --no-sync python - <<'EOF' +import jax, jax.numpy as jnp +from progenax import KingProfile +masses, key = jnp.ones(300), jax.random.PRNGKey(7) +def mean_r(W0): + pos = KingProfile.from_W0_rc(W0=W0, r_c=1.0).sample_positions(masses, key) + return jnp.mean(jnp.linalg.norm(pos, axis=1)) +g = float(jax.grad(mean_r)(5.0)); h = 1e-3 +fd = float((mean_r(5.0+h) - mean_r(5.0-h)) / (2*h)) +print(f"AD={g:.8f} FD={fd:.8f} ratio={g/fd:.6f}") +assert abs(g/fd - 1) < 1e-3 +EOF +``` +Expected: ratio within 1e-3 of 1 (pre-fix baseline was 0.999869). + +**Step 6: Commit** + +```bash +git add src/progenax/profiles/king.py src/progenax/profiles/michie.py src/progenax/profiles/eff.py tests/validation/ +git commit -m "fix(profiles): sqrt-stretched CDF grids resolve the core at high W0 (audit R4) + +Linear 1000-pt grids left <10 points in the core at W0>=9: measured core-mass +errors +18% (W0=9) and +270% (W0=12) vs a dense ODE reference. r = r_t*u^2 +grid + non-uniform trapezoid; AD-vs-FD W0 gradient regression preserved." +``` + +### Task 1.3: `sample_fixed_n` silent mass shortfall (R5) + +**The bug:** `BaseIMF.sample_fixed_n` (`src/progenax/imf/base.py:230-281`) promises "exactly n masses summing to m_total" but the one-sided quantile stretch (`q ≤ 1`) caps the achievable total at `Σ ppf((i+0.5)/n)`; the clipped Newton solver converges to the boundary silently. Measured: `Maschberger().sample_fixed_n(key, 1000, 500.0)` → **349.0 M☉**. None of the four sampling modes (`sample_m_total`, `sample_m_total_packed`, `sample_fixed_n`) has any test. + +**Scope decision (YAGNI):** guard + honest docs + tests. A two-sided stretch is a behavior redesign — defer it; note it in the phase report as a follow-up option for Anna. + +**Step 1: Write the failing tests** — create `tests/unit/imf/test_sampling_modes.py`: + +```python +"""Tests for the BaseIMF mass-target sampling modes (audit R5: previously zero tests).""" +import jax +import jax.numpy as jnp +import pytest + +from progenax.imf import Maschberger + + +class TestSampleFixedN: + def test_reachable_target_is_hit(self): + imf = Maschberger() + key = jax.random.PRNGKey(0) + n = 1000 + # natural total ~ n * E[m]; pick a comfortably reachable target + target = 250.0 + m = imf.sample_fixed_n(key, n, target) + assert m.shape == (n,) + assert abs(float(m.sum()) / target - 1.0) < 0.05 # ~target, not "exactly" + + def test_unreachable_target_raises(self): + """Pre-fix behavior: silently returned 349 Msun for a 500 Msun target.""" + imf = Maschberger() + key = jax.random.PRNGKey(0) + with pytest.raises(ValueError, match="unreachable"): + imf.sample_fixed_n(key, 1000, 500.0) + + def test_masses_within_imf_bounds(self): + imf = Maschberger() + m = imf.sample_fixed_n(jax.random.PRNGKey(1), 500, 120.0) + assert float(m.min()) >= imf.m_min + assert float(m.max()) <= imf.m_max + + +class TestSampleMTotalModes: + @pytest.mark.parametrize("seed", [0, 1]) + def test_sample_m_total_hits_target(self, seed): + imf = Maschberger() + out = imf.sample_m_total(jax.random.PRNGKey(seed), m_total=500.0) + # read the actual return signature first and unpack accordingly; + # assert realized total within one max-mass overshoot of the target + ... + + def test_sample_m_total_packed_consistent_with_unpacked(self): + ... +``` + +Before writing the two `sample_m_total` tests, READ `src/progenax/imf/base.py:120-228` for the actual signatures and overshoot semantics (prefix-cut, ≤1-star overshoot) and encode those semantics — not your guess. + +**Step 2: Run RED** — `test_unreachable_target_raises` must FAIL (no ValueError raised; total silently 349). + +**Step 3: Implement the guard** — in `sample_fixed_n` (base.py:230-251), after computing `u_base`: + +```python + u_base = (jnp.arange(n) + 0.5) / n + + # Reachability: the one-sided stretch (q <= 1) caps the total at + # sum(ppf(u_base)) — and stratification truncates the heavy tail, so + # this ceiling sits ~5% BELOW n*E[m]. A target above it used to return + # a silently short total (audit R5: asked 500, got 349). Eager inputs + # fail loudly; traced inputs cannot be checked (documented). + try: + m_ceiling = float(jnp.sum(self.ppf(u_base))) + if float(m_total) > m_ceiling: + raise ValueError( + f"m_total={float(m_total):.6g} Msun is unreachable for n={n}: " + f"the stratified-quantile ceiling is {m_ceiling:.6g} Msun. " + f"Increase n, lower m_total, or use sample_m_total()." + ) + except (jax.errors.ConcretizationTypeError, jax.errors.TracerArrayConversionError): + pass # traced m_total: caller owns reachability +``` + +Fix the docstring: replace "exactly n masses summing to m_total" with "n masses whose total approximates m_total (stratified quantile stretch; the realized total carries the residual of the final random jitter)" and document the ceiling + traced-input caveat. + +**Step 4: GREEN** — run the new file + `tests/unit/imf -q` (no regressions). + +**Step 5: Commit** + +```bash +git add src/progenax/imf/base.py tests/unit/imf/test_sampling_modes.py +git commit -m "fix(imf): sample_fixed_n raises on unreachable m_total instead of silently undershooting (audit R5) + +Measured: target 500 Msun, n=1000 returned 349 Msun with no warning. Adds the +stratified-ceiling guard (eager inputs), honest docstring, and first-ever tests +for the mass-target sampling-mode family." +``` + +### Task 1.4: `compute_stellar_radii` inverted exponents (R6) + +**The bug:** `src/progenax/builders.py:155-193` assigns R ∝ M^0.8 *above* 1 M☉ and M^0.57 *below* — the standard homology/empirical assignment is the opposite. 10 M☉ → 6.31 R☉ (ZAMS ≈ 4); 0.2 M☉ → 0.40 R☉ (observed ≈ 0.22). No citation anywhere; tests assert the code's own formula back at itself. + +**The fix:** adopt Demircan & Kahraman (1991, Ap&SS 181, 313) main-sequence fits — `R = 1.06·M^0.945` for M < 1.66 M☉, `R = 1.33·M^0.555` for M ≥ 1.66 M☉ — keeping the brown-dwarf plateau (0.1 R☉) below 0.08 M☉. Pleasant side effects: the low-mass branch meets the BD plateau almost continuously (1.06·0.08^0.945 = 0.0975 ≈ 0.1), and the D&K joint at 1.66 M☉ has only a ~3.5% step (their own fits; document it). + +> ⛔ **Mini-checkpoint inside this task:** before implementing, ask Anna whether she prefers Demircan & Kahraman (1991) (recommended: simple, cited, two-branch) or Tout et al. (1996) ZAMS (higher fidelity, rational-polynomial). The plan below assumes D&K91. + +**Step 1: Write the failing test** — replace the tautological block in `tests/unit/test_builders.py:11-31` with literature-anchored values: + +```python +class TestComputeStellarRadii: + """Anchored to Demircan & Kahraman (1991) Ap&SS 181, 313 + observed radii. + + Audit R6: the previous exponents were INVERTED vs MS homology + (M^0.8 above 1 Msun instead of below), giving 10 Msun -> 6.3 Rsun + (ZAMS ~4) and 0.2 Msun -> 0.40 Rsun (observed ~0.22). The old tests + asserted the code's own formula back at itself. + """ + + def test_solar(self): + r = compute_stellar_radii(jnp.array([1.0])) + assert abs(float(r[0]) - 1.06) < 0.01 # D&K91: 1.06 * 1^0.945 + + def test_massive_star_near_zams(self): + r = float(compute_stellar_radii(jnp.array([10.0]))[0]) + assert abs(r - 1.33 * 10.0**0.555) < 1e-6 # = 4.77 Rsun + assert 3.0 < r < 5.5 # within ~25% of ZAMS ~4 Rsun + + def test_m_dwarf(self): + r = float(compute_stellar_radii(jnp.array([0.2]))[0]) + assert abs(r - 1.06 * 0.2**0.945) < 1e-6 # = 0.231 Rsun + assert 0.17 < r < 0.28 # observed ~0.22 Rsun + + def test_brown_dwarf_plateau(self): + r = compute_stellar_radii(jnp.array([0.05, 0.02])) + assert jnp.all(jnp.abs(r - 0.1) < 0.05) + + def test_near_continuous_at_hydrogen_burning_limit(self): + r = compute_stellar_radii(jnp.array([0.079, 0.081])) + assert abs(float(r[1]) / float(r[0]) - 1.0) < 0.1 # no factor-2.4 jump (audit F7) +``` + +**Step 2: Run RED** — `test_solar` fails (old code: 1.0^0.57 = 1.0 ≠ 1.06), `test_massive_star_near_zams` fails (6.31). + +**Step 3: Implement** — replace the body of `compute_stellar_radii` (builders.py:155-193): + +```python +def compute_stellar_radii(masses: Float[Array, "N"]) -> Float[Array, "N"]: + """ + Main-sequence stellar radii from mass, in SOLAR RADII. + + Demircan & Kahraman (1991), Ap&SS 181, 313 (their Eqs. 5-6): + R/Rsun = 1.06 (M/Msun)^0.945 for 0.08 <= M < 1.66 Msun + R/Rsun = 1.33 (M/Msun)^0.555 for M >= 1.66 Msun + Brown dwarfs (M < 0.08): R ~ 0.1 Rsun plateau (electron degeneracy), + R = 0.1 (M/0.08)^0.08. The D&K branches meet at 1.66 Msun with their + own ~3.5% fit discontinuity; the low-mass branch meets the BD plateau + nearly continuously (1.06 * 0.08^0.945 = 0.0975). + + Used for collision radii in downstream N-body; ZAMS values, no evolution. + """ + + def radius_high_mass(m): + return 1.33 * jnp.power(m, 0.555) + + def radius_low_mass(m): + return 1.06 * jnp.power(m, 0.945) + + def radius_brown_dwarf(m): + return 0.1 * jnp.power(m / 0.08, 0.08) + + return jax.vmap( + lambda m: jax.lax.cond( + m >= 1.66, + radius_high_mass, + lambda mv: jax.lax.cond( + mv >= 0.08, radius_low_mass, radius_brown_dwarf, mv + ), + m, + ) + )(masses) +``` + +**Step 4: GREEN + downstream check** — run `tests/unit/test_builders.py -q` and grep for other consumers of `stellar_radii` expectations: `grep -rn "stellar_radii\|compute_stellar_radii" tests/ src/`. + +**Step 5: Commit** + +```bash +git add src/progenax/builders.py tests/unit/test_builders.py +git commit -m "fix(builders): stellar mass-radius relation per Demircan & Kahraman (1991) (audit R6) + +Previous exponents were inverted vs MS homology (uncited): 10 Msun -> 6.3 Rsun +instead of ~4.8, 0.2 Msun -> 0.40 instead of ~0.23. Tests now anchored to the +published fit, not the implementation." +``` + +### Task 1.5: PowerLawIMF NaN gradient at α = 1 (R10) + +**The bug:** four sites in `src/progenax/imf/power_law.py` use the bare `jnp.where(|e| < 1e-12, log_branch, (hi**e − lo**e)/e)` pattern; the untaken branch computes 0/0 and the `where` VJP propagates 0 × NaN = NaN. Confirmed: `jax.grad(mean_mass)(α=1.0)` → nan (interior values FD-exact). The correct double-`where` fix already exists at `src/progenax/imf/differentiable.py:47-54` — port it. + +**Sites** (verify each by reading; line numbers drift): `segment_integral` (~88-96), `_cdf_unnorm` (~200-216), `ppf` (~246-252), `mean_mass` (~276-288). + +**Step 1: Write the failing test** — append to `tests/unit/imf/test_imf_gradients.py`: + +```python +class TestAlphaOneGradients: + """Audit R10: bare where(|1-a|9.5 at the fixed fallback domain silently pins + xi_t to the boundary. Concrete high W0 auto-scales (fine); traced high + W0 must be an explicit, loud choice via xi_max.""" + import pytest + with pytest.raises(ValueError, match="xi_max"): + jax.jit(lambda w: KingProfile.from_W0_rc(W0=w, r_c=1.0))(12.0) + + +def test_traced_typical_w0_still_works(): + prof = jax.jit(lambda w: KingProfile.from_W0_rc(W0=w, r_c=1.0))(7.0) + assert jnp.isfinite(prof.r_t) +``` + +Wait — a ValueError cannot depend on a *traced value*. Correct design: `from_W0_rc` cannot know the traced W0. Instead, make the **pinning detectable and loud at solve time**: have `_find_tidal_radius` (king.py:241-283) compute `is_pinned = psi_grid[-1] > 0` (a traced bool) and store it on the profile as a diagnostic field `r_t_is_pinned`; for CONCRETE W0 > 9.5 with the fallback domain this cannot happen (auto-scaling), so additionally raise eagerly in `_auto_ode_domain`'s except-branch is impossible (W0 unknown). So the testable contract is: + +```python +def test_pinned_domain_is_flagged(): + """Force the failure mode concretely: explicit too-small xi_max.""" + prof = KingProfile.from_W0_rc(W0=12.0, r_c=1.0, xi_max=50.0, n_ode_points=500) + assert bool(prof.r_t_is_pinned) + +def test_healthy_solve_not_flagged(): + prof = KingProfile.from_W0_rc(W0=7.0, r_c=1.0) + assert not bool(prof.r_t_is_pinned) + +def test_concrete_pinned_solve_raises(): + """Eager construction with a domain that pins xi_t must refuse loudly.""" + import pytest + with pytest.raises(ValueError, match="pinned|xi_max"): + KingProfile.from_W0_rc(W0=12.0, r_c=1.0, xi_max=50.0, n_ode_points=500) +``` + +(The first test then becomes the traced-path behavior: flag stored, no raise — the Engine-B two-tier concrete/traced pattern already used in `cluster/eddington_engine.py:192-217`; read it and mirror it.) Drop the two draft tests above in favor of these three. + +**Step 2-4:** RED → implement (`r_t_is_pinned: Bool[Array, ""]` eqx field; eager `ValueError` when `jnp.ndim`-concrete and pinned; mirror the Engine-B `_is_concrete` guard) → GREEN → commit: + +```bash +git commit -m "fix(profiles): King r_t boundary-pinning is now flagged (traced) or refused (eager) (audit J4)" +``` + +### Task 2.2: `virial_scale` T=0 NaN + dedupe with its sibling (J5) + +`builders.py:256-258` divides by `Q_current` with no guard → cold input (v=0) yields NaN velocities silently; the near-duplicate `dynamics/virial.py:251 rescale_velocities_to_virial` has a different guard. Rescaling from T=0 is mathematically impossible — refuse loudly. + +**Step 1: Failing test** (`tests/unit/test_builders.py`): + +```python +def test_virial_scale_zero_velocities_raises(): + """Audit J5: cold input used to return all-NaN velocities silently.""" + import pytest + pos = jax.random.normal(jax.random.PRNGKey(0), (50, 3)) + vel = jnp.zeros((50, 3)) + m = jnp.ones(50) + with pytest.raises(ValueError, match="zero kinetic"): + virial_scale(pos, vel, m, Q_target=0.5, G=0.00449) +``` + +**Step 2-4:** RED → implement an eager concrete-input check (`try: float(T) ... except ConcretizationTypeError: pass`, raising `ValueError("cannot rescale from zero kinetic energy (T=0): velocities are all zero")`; for traced inputs keep a `jnp.where(T > 0, scale, jnp.nan)` — NaN is then the honest traced sentinel, document it) → GREEN. Then reconcile with `rescale_velocities_to_virial`: make `virial_scale` delegate to it or vice versa so there is ONE implementation (DRY; read both first — they differ in COM handling, keep the semantics of each public signature). Commit: + +```bash +git commit -m "fix(builders): virial_scale refuses T=0 input; dedupe with rescale_velocities_to_virial (audit J5)" +``` + +### Task 2.3: `q_approx` static-N branch (J6) + +`diagnostics/q_approx.py:270-275` (and `_compute_s_bar_subsampled:239`) use `lax.cond` on a *Python int* N — both branches compile every shape (~0.4 s wasted). Replace with Python `if N > 1000:` — identical semantics, half the compile. + +No new physics test needed (pure performance); the regression is the existing `tests/unit/substructure` + `tests/unit/diagnostics` suites passing unchanged, plus: + +```python +def test_auto_dispatch_small_n_equals_naive(): + pos = jax.random.normal(jax.random.PRNGKey(0), (200, 3)) + assert jnp.allclose(q_approx(pos, method="auto"), q_approx(pos, method="naive")) + +def test_auto_dispatch_large_n_equals_fast(): + pos = jax.random.normal(jax.random.PRNGKey(0), (1500, 3)) + assert jnp.allclose(q_approx(pos, method="auto"), q_approx(pos, method="fast")) +``` + +Commit: `perf(diagnostics): Python-if dispatch on static N in q_approx — avoids compiling both branches (audit J6)`. + +### Task 2.4: King r_t consistency guard (S1, A3, M4) + +`KingProfile`'s direct constructor accepts an arbitrary `r_t` inconsistent with c(W0) — `KingProfile(W0=7, r_c=1, r_t=10)` (the README's own example!) silently builds a non-self-consistent, non-equilibrium model. `KingVelocityDF.r_t` (`king_df.py:101,123`) is stored and never used. + +**Design (confirm at checkpoint if unsure):** in the direct `KingProfile.__init__`, when ALL inputs are concrete, warn (`warnings.warn`) if `r_t` deviates from `r_c · ξ_t(W0)` by > 5%; `from_W0_rc` (which derives r_t) is untouched. Remove the dead `r_t` field from `KingVelocityDF` — this repo's policy is no backwards-compat shims; update every callsite (`grep -rn "KingVelocityDF(" src tests docs README.md`). + +**TDD:** failing test asserting `pytest.warns(UserWarning, match="inconsistent")` for the bad triple and no warning for the consistent one; failing test asserting `KingVelocityDF(W0=7.0, r_c=1.0)` constructs (no r_t arg). RED → implement → GREEN → update callsites → commit: + +```bash +git commit -m "fix(api): warn on inconsistent KingProfile r_t; drop the dead KingVelocityDF.r_t field (audit S1)" +``` + +### Task 2.5: Pin the EFF γ=3 sub-virial offset (S2) + +The default EFF DF (γ=3, sharp truncation) is knowingly ~5-8% sub-virial (Eddington inversion of a truncated ρ with ρ(r_t)>0 — disclosed in the module docstring, but the default config is the broken case and nothing pins it, so a regression to −15% would pass). + +**Step 1: Failing-by-absence test** — append to `tests/validation/test_eff_physics.py`: + +```python +def test_gamma3_default_subvirial_offset_is_pinned(): + """Audit S2: gamma=3 + sharp truncation is KNOWN ~5-8% sub-virial by + construction (f(E) cannot represent rho(r_t)>0). Pin the band so a + regression (e.g. to -15%) cannot pass silently. This is a documented + limitation, not a target: see eff_df.py module docstring.""" + # build the DEFAULT EFFProfile/EFFVelocityDF pair (read the current + # defaults from eff.py / eff_df.py), sample N=20_000, seed 0, + # compute Q = T/|V| with compute_kinetic_energy / compute_potential_energy + # at G=STELLAR.G, and assert: + assert 0.42 < Q < 0.49, f"gamma=3 Q drifted out of the documented band: {Q}" +``` + +Fill in the construction by copying the existing γ=5 equilibrium test in the same file (lines ~227-249) and switching to defaults. Measure Q at 3 seeds first, set the band as measured ± 3σ (measured-first frozen gates — the project convention), THEN freeze. Also surface the caveat at the constructor docstring level in `eff_df.py` (currently module-level only). + +Commit: `test(eff): pin the documented gamma=3 sub-virial offset + constructor-level caveat (audit S2)`. + +### Task 2.6: Dedupe the `VelocityDF` protocol (A1) + +Two structurally identical `runtime_checkable` Protocols: `protocols.py:54-83` (canonical) and `kinematics/api.py:53-70` (drift hazard). Make `kinematics/api.py` import from `progenax.protocols`; `grep -rn "VelocityDF" src/ tests/` and fix every import that referenced the api.py copy. Regression: `tests/unit/test_protocols.py` + FAST gate. Commit: `refactor: single source of truth for the VelocityDF protocol (audit A1)`. + +### ⛔ CHECKPOINT 2 +FULL gate + report. Flag for Anna: the S1 design decision (warn vs raise on inconsistent r_t), and the deferred units-policy question (A2: should DF `G=None` defaults be removed for protocol-wide explicit G? — pre-1.0 is the time; recommend deciding now but it is a breaking API sweep, so it's HER call, not yours). + +--- + +# Phase 3 — Documentation honesty sweep (R7, D3–D5, S3, S4, S6, S10, S16, T5, L1) + +No physics changes in this phase; the verification is (a) every documented symbol imports, (b) every documented command runs, (c) the docs-example smoke test passes, (d) myst build clean. + +### Task 3.1: Reconcile the public API with the docs (R7 part 1) + +**Decision (recommended; confirm at checkpoint):** export the IMF/binary-statistics symbols from `progenax.__init__` — the docs already promise them and they are stable. + +**Step 1: Failing test** — create `tests/unit/test_documented_api.py`: + +```python +"""Every symbol the docs promise must import from the package root (audit R7). + +IGIMF and EnvironmentIMF are deliberately ABSENT: they never existed +(the environment-dependent IMF is the functional BirthEnvironment + +env_to_imf_params API); the docs are being fixed to stop advertising them. +""" +import pytest + +DOCUMENTED_ROOT_SYMBOLS = [ + "PowerLawIMF", "ChabrierIMF", "Maschberger", "TruncatedIMF", "BinaryIMF", + "FlatMassRatio", "PowerLawMassRatio", "TwinPeakedMassRatio", + "MoeDiStefano2017", "MoeDiStefano2017Full", "MoePeriod", "MoeJointOrbit", + "ConstantBinaryFraction", "MassDependentBinaryFraction", + "UniformSphereProfile", # audit L1: only profile missing from the root +] + +@pytest.mark.parametrize("name", DOCUMENTED_ROOT_SYMBOLS) +def test_symbol_importable_from_root(name): + import progenax + assert hasattr(progenax, name), f"progenax.{name} promised by docs, not exported" + +def test_phantom_classes_stay_dead(): + import progenax + for phantom in ("IGIMF", "EnvironmentIMF"): + assert not hasattr(progenax, phantom) +``` + +**Step 2-4:** RED (15 failures) → add the imports + `__all__` entries to `src/progenax/__init__.py` (group under an `# IMF / binary statistics` comment matching the existing style) → GREEN. + +**Step 5: Purge the phantoms from docs.** `grep -rn "IGIMF\|EnvironmentIMF" CLAUDE.md ../CLAUDE.md README.md docs/ src/` — replace every class-like mention with the real API ("environment-dependent IMF: `BirthEnvironment` + `env_to_imf_params()` (Marks+2012/Jeřábková+2018 α₃ relations)"). The audit (§4.3) is explicit: never use the word "IGIMF" for this machinery — it is not galaxy-wide IGIMF integration. + +Commit: `fix(api+docs): export the documented IMF symbols; kill the phantom IGIMF/EnvironmentIMF claims (audit R7)`. + +### Task 3.2: README rewrite (D3, H1, H2) + +Rewrite `README.md` against reality. Requirements checklist (verify each with a command, not by eye): + +- [ ] Quickstart uses ONLY current API — the `TwoComponentConfig`/`generate_two_component_cluster` block (lines ~154-175) is replaced by a `MultiComponentCluster` example (`from_components` or `from_density_profiles`; copy a working invocation from `tests/unit/cluster/`) plus the basic `build_spatial_ic` example. +- [ ] `MultiComponentCluster`, `MichieProfile`/`MichieVelocityDF`, `LIMEPYVelocityDF` appear in the feature tables. +- [ ] `apply_osipkov_merritt()` (line ~34) replaced by the real `anisotropy_radius` constructor argument. +- [ ] The King example `KingVelocityDF(W0=7.0, r_c=1.0, r_t=10.0)` (line ~320) fixed to the post-Task-2.4 signature. +- [ ] Installation: uv-first, sibling-checkout layout, only the extras that exist (`dev`, `experimental`, plus `diagnostics` after Task 4.2); the conda block deleted; an honest "not on PyPI yet — requires the jaxstro sibling checkout" banner. +- [ ] Stale metrics removed or regenerated (LOC/test counts — quote the FULL-gate number from CHECKPOINT 1, or drop counts entirely; audit T5/L3: three documents currently give three different numbers). +- [ ] "3 runtime-checkable protocols" → 9 (or just "protocol-based", uncounted). +- [ ] License section matches the actual LICENSE file (Task 4.1). +- [ ] No PRNG key reuse in any example (the installation smoke test currently passes the same key twice — split keys in every snippet). + +**Verification — the docs-example smoke test (audit test-matrix item 9).** Create `tests/integration/test_readme_examples.py`: + +```python +"""Execute every python code block in README.md (audit: would have caught R7). + +Skips blocks marked `# doctest: +SKIP` (e.g. GPU-only). Each block runs in a +fresh namespace with a 120 s budget. +""" +import re +from pathlib import Path + +import pytest + +README = (Path(__file__).resolve().parents[2] / "README.md").read_text() +BLOCKS = [ + b for b in re.findall(r"```python\n(.*?)```", README, re.DOTALL) + if "+SKIP" not in b +] + +@pytest.mark.parametrize("i", range(len(BLOCKS))) +def test_readme_block_executes(i): + exec(compile(BLOCKS[i], f"README.md:block{i}", "exec"), {}) +``` + +Run it RED against the CURRENT README first (it must fail on `TwoComponentConfig` — proof the harness works), then rewrite, then GREEN. Same treatment for `docs/website/00-getting-started/installation.md` code blocks if cheap; at minimum fix the `[all]`/`[io]`/`[viz]`/`[ml]` phantom extras and the monorepo-vs-split-repo clone instructions (H2), and the reused-key smoke test. + +Commit: `docs: README rewritten against the current API, with an executable-examples smoke test (audit R7/D3/D4)`. + +### Task 3.3: Kill the `progenax-legacy` fabrication (D5/H4) + +`docs/website/20-architecture/ic-redesign-history.md:110` claims `pip install progenax-legacy` works — it does not exist on PyPI. Rewrite that paragraph to: the pre-redesign code is preserved at the repo tag/path (find the real one: `git tag` / `ls ../legacy/`), no install claim. Then `grep -rn "progenax-legacy" docs/ README.md` → zero hits. Verify the website still builds: `env -u VIRTUAL_ENV uv run --no-sync myst build 2>&1 | tail -3` (expect 0 warnings, ~151 pages). Commit. + +### Task 3.4: Equilibrium-caveat docstrings (S3, S4, S6, S10, S16, S15) + +Pure docstring/comment fixes — one commit, no test changes, but each claim must be *accurate*, so read the audit entry + the function before writing: + +- `kinematics/rotation.py` (both transforms + module): add "kinematic overlay — injects kinetic energy and L_z; the output is NOT a stationary equilibrium (Q > 0.5)"; remove the misleading `target_Q` rescale suggestion in `kinematics/api.py:199-235` (an isotropic rescale does not restore stationarity); fix the stale zero-axis NaN comment (`rotation.py:49-50`) — behavior is a silent no-op; add an eager ValueError for a concrete zero axis (this one IS testable: 2-line test). +- `tidal.py:apply_tidal_truncation`: paragraph on the survivor set being super-virial w.r.t. its own reduced potential (some stars formally unbound near r_t); recommend `rescale_velocities_to_virial` or King/LIMEPY r_t-consistent models. +- `cluster/mass_segregation.py:energy_sorted_segregation`: the segregated mass-weighted density no longer matches the parent profile → the self-consistent potential differs; callers should finalize with a global virial rescale (as the validation suite already does). +- `binaries/diagnostics.py:222-223`: "inflated" → "deflated" (measured Q_resolved 0.31 < 0.5); add the one-line note that `softening` does not soften `E_internal`. +- `dynamics/virial.py:201-203`: per-group W_j is origin-dependent for subgroups; document "positions must be pre-centered" (or center internally — your call, but if you center, run the cluster validation tests). + +Commit: `docs(src): honest equilibrium caveats — rotation overlays, tidal truncation, segregation, energy-budget wording (audit S3/S4/S6/S10/S16)`. + +### Task 3.5: Reconcile the three stale test-count claims (T5/L3/D6) + +`tests/README.md` says 874; `progenax/CLAUDE.md` says 1163; STATUS.md says 1192. Update `tests/README.md` and `CLAUDE.md` to the post-Phase-1 FULL-gate count (quote the actual pytest output), and where counts appear in prose, prefer "the released-core suite (see CI)" over a number that rots. Commit. + +### ⛔ CHECKPOINT 3 +FULL gate + `myst build` output + the README smoke test green + `grep` proofs (zero hits for `TwoComponentConfig|progenax-legacy|IGIMF` outside the audit report and dev-log history pages). **Wait for approval.** + +--- + +# Phase 4 — Packaging & metadata (R2, R9, D1, D7, M6/H1-CI) + +### Task 4.1: LICENSE + project metadata (D1/C3) + +- Create `LICENSE` — MIT, copyright `2026 Anna Rosen` (⛔ confirm holder/year with Anna at the checkpoint if any doubt — the README claims MIT already). +- In `pyproject.toml [project]` add: + +```toml +license = "MIT" +license-files = ["LICENSE"] +authors = [{ name = "Anna Rosen", email = "alrosen@sdsu.edu" }] +keywords = ["astronomy", "n-body", "initial-conditions", "jax", "star-clusters", "differentiable"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Astronomy", +] + +[project.urls] +Repository = "https://github.com/jaxstro/progenax" +``` + +- `uv lock` (metadata change → lockfile refresh), then verify the wheel: + +```bash +env -u VIRTUAL_ENV uv build --wheel -o /tmp/progenax_wheel_check +python3 -c " +import zipfile, glob +w = glob.glob('/tmp/progenax_wheel_check/*.whl')[0] +meta = zipfile.ZipFile(w).read([n for n in zipfile.ZipFile(w).namelist() if n.endswith('METADATA')][0]).decode() +for field in ('License', 'Author', 'Classifier', 'Project-URL'): + assert field in meta, field +print('wheel metadata OK') +assert not any('experimental' in n or 'gravoturb' in n for n in zipfile.ZipFile(w).namelist()) +print('wheel content clean (no experimental leakage)') +" +``` + +Commit: `chore(packaging): LICENSE (MIT) + author/classifiers/URLs in wheel metadata (audit D1)`. + +### Task 4.2: `[diagnostics]` dependency story (R9/H3) + +`progenax.diagnostics` imports numpy + scipy eagerly (`substructure.py:20-22`, `mass_segregation.py:18-20`, both pulled in by `diagnostics/__init__.py:41-46`), but neither is a declared dependency → `from progenax.diagnostics import compute_q_parameter` explodes in a clean install. + +**Step 1: Failing test** — `tests/unit/diagnostics/test_lazy_imports.py`: + +```python +import subprocess, sys + +def test_diagnostics_import_without_scipy_gives_actionable_error(): + """Audit R9: in an env without scipy, importing the subpackage must not + crash at import time; calling an exact (scipy-backed) function must raise + an ImportError naming the [diagnostics] extra.""" + code = ( + "import sys;" + "sys.modules['scipy'] = None; sys.modules['scipy.sparse'] = None;" + "sys.modules['scipy.sparse.csgraph'] = None; sys.modules['scipy.spatial'] = None;" + "sys.modules['scipy.spatial.distance'] = None;" + "import progenax.diagnostics as d;" # must import fine + "from progenax.diagnostics import q_approx;" # JAX-native: must work + "import jax.numpy as jnp, jax;" + "q_approx(jax.random.normal(jax.random.PRNGKey(0), (50, 3)));" + "import pytest;" + "exc = None\n" + "try:\n" + " d.compute_q_parameter(jax.random.normal(jax.random.PRNGKey(0), (50, 3)))\n" + "except ImportError as e:\n" + " exc = e\n" + "assert exc is not None and 'diagnostics' in str(exc), exc\n" + ) + r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert r.returncode == 0, r.stderr +``` + +**Step 2-4:** RED → implement: move the numpy/scipy imports inside the functions that need them (`compute_q_parameter`, `compute_lambda_msr`, the MST builders), wrapped: + +```python + try: + import scipy.sparse.csgraph as csgraph + except ImportError as e: + raise ImportError( + "compute_q_parameter needs numpy+scipy (the exact, non-differentiable " + "diagnostics path). Install them with: uv pip install 'progenax[diagnostics]'" + ) from e +``` + +and add to `pyproject.toml`: + +```toml +diagnostics = ["numpy>=1.24", "scipy>=1.11"] +``` + +(`uv lock` after.) → GREEN → commit: `fix(packaging): [diagnostics] extra + lazy numpy/scipy imports with actionable errors (audit R9)`. + +### Task 4.3: Wheel-smoke + version-matrix CI (M6, T3, audit test-matrix items 4 & 10) + +Append two jobs to `.github/workflows/tests.yml`: + +```yaml + wheel-smoke: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + # same sibling checkouts as released-core + # ... + - name: Setup uv + uses: astral-sh/setup-uv@v6 + with: + python-version: "3.13" + - name: Build wheel + working-directory: progenax + run: uv build --wheel -o dist/ + - name: Install wheel in a clean venv and import + working-directory: progenax + run: | + uv venv /tmp/clean + uv pip install --python /tmp/clean/bin/python dist/*.whl ../jaxstro + /tmp/clean/bin/python -c "import progenax; import progenax.diagnostics; print(progenax.__name__, 'imports clean')" +``` + +and a nightly matrix job in `physics-validation.yml` (NOT the PR lane — keep PR latency flat): `python-version: ["3.10", "3.13"]` running the unit tier only. If 3.10 fails for an uninteresting reason (e.g. a dependency floor), the honest alternative is raising `requires-python` — flag it at the checkpoint rather than burning time. + +Add `wheel-smoke` to the `tests` aggregator's `needs` + result check. Commit. + +### Task 4.4: CHANGELOG + the jaxstro decision (D7, R2) — ⛔ DECISION TASK, do not implement unilaterally + +- Create `CHANGELOG.md` with an `## 0.1.0 (unreleased)` section summarizing the audit-fix arc (one line per R-finding fixed, with commit refs). +- **The pip-installability blocker (R2) requires Anna's strategy decision** (audit §9): (A — recommended) publish `jaxstro` to PyPI first and pin `jaxstro>=0.1,<0.2`; (B) vendor `units`+`jaxconfig` into `progenax._vendor`; (C) git-URL dep for a GitHub-only soft launch. Present the trade-offs from the audit verbatim and STOP. Whichever she picks is a separate follow-up arc (Option A is mostly work in the jaxstro repo). + +### ⛔ CHECKPOINT 4 (final) +1. FULL gate output (expect ≈ 1192 + ~25 new tests, 0 failures). +2. FAST gate output. +3. `PROGENAX_STRICT_REFS=1` slow lane locally once (`pytest tests -m slow -q`) — the trust anchors must pass with the Phase-1 physics fixes in (the LIMEPY parity and Engine A/B tests must be untouched by the CDF-grid change: Engine A samples via its own DF tables, but VERIFY, don't assume). +4. CI green on `fix/prerelease-audit` (all jobs incl. wheel-smoke), `gh run` URL. +5. Wheel metadata + content check output (Task 4.1). +6. Completion doc `.claude-work/PRERELEASE_AUDIT_FIXES_COMPLETE.md`: per-finding before/after evidence table (R1–R10 + each S/J/A/T/D handled), lessons learned, and the open Anna-decision list (jaxstro strategy, units-policy A2, two-sided stretch, D&K91-vs-Tout). +7. Update `STATUS.md` `next:`/`blocker:` lines. +8. **Merge/PR: Anna's call.** Per the project git workflow: merge to local main only on her go; ONE final PR when CI is green. + +--- + +## Explicitly OUT of scope (deferred, listed so they aren't silently dropped) + +- jaxstro PyPI publication / vendoring (Anna decision, Task 4.4). +- Units-policy A2 sweep (drop `G=None` from DF protocol) — breaking change, Anna decision. +- `q_approx` calibration sweep over (profile, N) (S5) — separate validation arc. +- Pre-paper PDF verification sweep (Moe Eqs. 17-18, CW04 Table 1, GZ15 §2.2) — requires Anna's eyes on the PDFs per project policy. +- 10-t_cross gravax stationarity runs per DF — paper arc, needs gravax. +- Two-sided quantile stretch for `sample_fixed_n` (R5 follow-up). +- Experimental `gravoturb_fdf` GRF β-gradient NaN (J3) — off the production inference path; fix opportunistically if touching that file. diff --git a/docs/reviews/2026-06-11-prerelease-adversarial-audit.md b/docs/reviews/2026-06-11-prerelease-adversarial-audit.md new file mode 100644 index 00000000..4cd822f5 --- /dev/null +++ b/docs/reviews/2026-06-11-prerelease-adversarial-audit.md @@ -0,0 +1,375 @@ +# progenax — Pre-Release Adversarial Audit (2026-06-11) + +**Reviewer:** Claude Fable 5 orchestrating six specialist audit agents (profiles/kinematics science, IMF/binaries science, cluster/builders/diagnostics science, JAX/differentiability, test suite, architecture/packaging). Every claim below carries file:line evidence from this session; agents independently re-derived formulas, executed gradient finite-difference cross-checks, transcription-checked Moe & Di Stefano (2017) Table 13 against the PDF, measured King CDF errors numerically, built and inspected the wheel, and pulled live CI run logs. + +--- + +## 1. Executive Summary + +The scientific core of progenax is in **unusually good shape for research software** — and that praise is specific: the Plummer inverse-CDF, the King ODE (factor-9 nondimensionalization re-derived), the Merritt-1985 OM coefficient 63/4, the Eddington inversion incl. the boundary term, the Kepler solver (machine precision to e=0.9999 *with correct gradients*), the Murray & Dermott element rotation matrix, the Moe Table 13 transcription (all 80 cells), Sana+2012 and Maschberger 2013 parameters, the GZ15 multimass LIMEPY engine (gated by parity against the canonical reference code at a pinned SHA with embedded provenance), and the CW04 Q diagnostic were each independently verified this session. The differentiability engineering (double-where NaN guards, straight-through estimators, checkpointed blocked kernels, AD-vs-FD validation tests) is state of the art; every headline gradient claim tested empirically verified against finite differences to ≤1e-5 relative. + +**But the repository is not releasable today**, for reasons that are mostly at the boundary rather than in the physics: + +1. **CI is red on main and has run zero tests for ≥5 pushes** (missing `jaxstroviz` sibling checkout). +2. Even when green, **CI never runs the slow-marked headline validations** (LIMEPY parity, Engine A-vs-B anchor, multimass equilibrium) that the docs advertise as the project's trust anchors. +3. **Nobody outside this machine can install the package** (path deps on unpublished `jaxstro`/`jaxstroviz`; no LICENSE; no metadata). +4. The **README teaches a deleted API** (`TwoComponentConfig`) and omits the flagship `MultiComponentCluster`; docs advertise phantom classes (`IGIMF`, `EnvironmentIMF`) and a nonexistent `[all]` extra and `progenax-legacy` PyPI package. +5. **Two confirmed science bugs**: the Moe & Di Stefano F_twin normalization (twins mixed against q>0.1 instead of the paper's q>0.3 → +22% twin overweight at solar logP=1), and an inverted stellar mass–radius relation in `compute_stellar_radii`. +6. **One quantified numerical defect**: King/Michie position-CDF grids under-resolve the core at W0 ≳ 9 (measured +18% enclosed-mass error at 0.3 r_c for W0=9, ~3.6× at W0=12). +7. `BaseIMF.sample_fixed_n` **silently returns 30% less mass than requested** when the target is unreachable, and the whole sampling-mode family has zero tests. + +All blockers are fixable in days, not months. Classification: **Release candidate** for the core (profiles/DFs/Kepler/builders), **Beta** for the binary-population layer until F_twin is fixed, **packaging/CI: pre-alpha**. + +--- + +## 2. Overall Release Readiness + +**Classification: Beta scientific software, close to Release Candidate.** + +Why not RC: a release candidate must be installable, have green CI enforcing its advertised validation claims, and have no known science bugs in features the docs call "faithful." None of those holds today. Why not Alpha: the validation culture is far beyond alpha — external-reference gating (King 1966 Table II, canonical LIMEPY at pinned SHA, CW04 Table 1, IAU constants, Moe Table 13), measured-first frozen tolerances with written noise budgets, AD-vs-FD gradient tests, and honest-scope documentation throughout. + +The gap to public release is ~2 weeks of mechanical work (CI, packaging, README, two science fixes) — **not** a rewrite. + +--- + +## 3. Highest-Risk Findings + +| # | Severity | Category | Finding | Files | Blocks release | +|---|---|---|---|---|---| +| R1 | **Critical** | CI | CI entirely red on main since ≥2026-06-10; `uv sync` dies on missing `../jaxstroviz`; zero tests enforced anywhere but the laptop | `pyproject.toml:43-46`, `.github/workflows/tests.yml:26-62` | **YES** | +| R2 | **Critical** | Packaging | Wheel not installable outside this machine: `jaxstro` unpinned + unpublished; `[tool.uv.sources]` paths don't ship; no LICENSE file, no author/classifier/URL metadata (README claims MIT with no license grant) | `pyproject.toml`, missing `LICENSE` | **YES** | +| R3 | **High** | Science | **Moe & Di Stefano F_twin mis-normalized**: twin excess mixed against the q∈[0.1,1] population; paper (p.5, Fig. 2) defines F_twin against q>0.3. Measured: realized paper-convention F_twin = 0.367 vs intended 0.30 at solar logP=1 (+22% twins), with companion deficit at q<0.3. Affects `MoeDiStefano2017Full`, `MoeJointOrbit`, `MoeCompanions` — everything labeled "faithful" | `src/progenax/imf/binary/moe_di_stefano.py:286-289` | **YES** (for the "faithful" label) | +| R4 | **High** | Science | **King/Michie/EFF CDF grid under-resolves the core at high W0**: linear 1000-pt grid over [0, r_t]; measured enclosed-mass error at 0.3 r_c: +0.9% (W0=7), **+17.7% (W0=9)**, **+264% (W0=12)**. Concentration tests pass (use the ODE, not the CDF); global Q passes (core mass tiny) — so the suite is blind to it | `src/progenax/profiles/king.py:359`, `michie.py:177`, `eff.py:83` | **YES** for advertised W0 ≳ 9 support; no for W0 ≤ 8 | +| R5 | **High** | Science | **`sample_fixed_n` silent mass shortfall**: docstring says "exactly n masses summing to m_total"; measured Maschberger n=1000, target 500 M☉ → returns **349 M☉** with no warning (one-sided quantile stretch hits its ceiling; Newton clips to bound silently). Zero tests for all four sampling modes | `src/progenax/imf/base.py:230-281` | **YES** (silent quantitative failure in documented API) | +| R6 | **High** | Science | **`compute_stellar_radii` mass–radius exponents inverted** vs standard MS homology (code: R∝M^0.8 above 1 M☉, M^0.57 below; literature: the opposite). 10 M☉ → 6.3 R☉ (should be ~4); 0.2 M☉ → 0.40 R☉ (observed ~0.22). No citation; tests assert the code's own formula back at itself. Propagates into every `ICResult.stellar_radii` (collision radii downstream) | `src/progenax/builders.py:155-193` | No (auxiliary), but fix before paper | +| R7 | **High** | Docs/API | README quickstart teaches deleted `TwoComponentConfig`/`generate_two_component_cluster`; `MultiComponentCluster` (the flagship) absent from README; docs advertise nonexistent `IGIMF`/`EnvironmentIMF` classes and 16 IMF symbols not exported from `progenax.__init__`; `installation.md` commands (`[all]` extra) fail | `README.md:154-175`, `progenax/CLAUDE.md`, `docs/website/00-getting-started/installation.md:21,40` | **YES** | +| R8 | **High** | CI/Testing | All headline external-reference validations (`test_limepy_reference_parity`, multimass equilibrium, Engine A-vs-B anchor, Engine B AD-vs-FD) are `@pytest.mark.slow` and CI runs `-m "not slow"` — the trust anchors are enforced by local convention only | `tests/validation/*:pytestmark`, `tests.yml:72` | **YES** | +| R9 | **High** | Packaging | `progenax.diagnostics` imports numpy + scipy eagerly; neither is a declared dependency — `from progenax.diagnostics import compute_q_parameter` raises ModuleNotFoundError in a clean install | `src/progenax/diagnostics/substructure.py:20-22`, `__init__.py:41-46` | **YES** for any README follower | +| R10 | **Medium** | JAX | `PowerLawIMF` NaN gradient at α=1.0 exactly (where-NaN; confirmed by execution). The correct `exp_safe` fix already exists in `imf/differentiable.py:47-54` — port it to 3 sites | `src/progenax/imf/power_law.py:92-96, 200-204, 249-252, 276-280` | Borderline; fix is 3 lines × 3 sites | + +--- + +## 4. Scientific Formula Audit + +### 4.1 Verified correct (independently re-derived this session — not taken on faith) + +| Formula | Where | How verified | +|---|---|---| +| Plummer a(r_h) = r_h√(2^{2/3}−1); inverse CDF r = a√(u^{2/3}/(1−u^{2/3})) | `profiles/plummer.py:48,122-123` | Algebraic re-derivation vs Aarseth+74 | +| Plummer speed g(q)∝q²(1−q²)^{7/2} ⇒ Beta(3/2, 9/2) | `kinematics/plummer_df.py:250` | Change-of-variables re-derived | +| Plummer OM DF coefficient 63/4 and bound ρ_a≤16/9 ⇔ r_a≥0.75a | `plummer_df.py:33,130` | Re-derived from Merritt 1985 Eq. 45–46 Γ-ratio | +| King ρ̂(W) = e^W erf(√W) − (2/√π)√W(1+2W/3) | `profiles/king.py:43-77` | Matches BT08 Eq. 4.131 | +| King Poisson ODE factor 9; r_c=√(9σ²/4πGρ₀); BCs + L'Hôpital center | `king.py:85-135` | Re-derived from ∇²W=−4πGρ/σ² | +| King DF: true lowered Maxwellian at local ψ, truncated at σ√(2ψ), σ²=GM/(9 r_c μ) | `kinematics/king_df.py:46-61,142-151,201-222` | Re-derived; no clipping; unscaled Q=0.5 tested at W0=7 and 12 | +| King c(W0) | `tests/validation/test_king_physics.py:266-274` | vs King (1966) Table II (paper-PDF-verified per project memory), \|Δ\|≤0.02–0.03 | +| Eddington inversion incl. (dρ/dΨ)\|₀/√E boundary term, u=√(E−Ψ) desingularization; OM via augmented density with full Q-variable inversion (NOT an approximation) | `kinematics/eddington.py:31-98` | Re-derived vs BT08 4.46a; three oracles (analytic Plummer f∝E^{7/2}, truncated closed form, rescale equivariance) | +| OM direction split: v_r=s cosθ, v_t=s sinθ/√(1+r²/r_a²) ⇒ β(r)=r²/(r²+r_a²) by construction | `eddington.py:133-177` | Jacobian consistency re-derived; a-posteriori \|Δβ\|<0.03–0.05 at 4 radii | +| Michie f∝exp(−s²u_t²/2)[exp(W−u²/2)−1]; closed-form u_r integral | `profiles/michie.py:26-60` | Re-derived; isotropic limit → King <1% | +| Virial PE = −G Σ_{i0 not representable by f(E)). Honest module docstring, but the *default constructor config* is the broken case, only γ=5 is in the headline table, and no test pins Q(γ=3) — a regression to −15% would pass | `kinematics/eff_df.py:15-21` | Constructor-level caveat + pinned Q(γ=3) band test | +| S3 | Med | Rotation transforms are additive kinematic overlays (`velocities + v_rotation`) that inject energy/L_z and break equilibrium; not stated at function level; `api.py` hint to rescale to Q=0.5 is physically misleading (isotropic rescale ≠ stationarity) | `kinematics/rotation.py:14-117`, `api.py:199-235` | Explicit "breaks equilibrium" caveat; optionally sign-flip (Lynden-Bell) equilibrium-preserving variant | +| S4 | Med | `apply_tidal_truncation` leaves survivor velocities drawn for the untruncated potential → super-virial survivor set, some formally unbound stars near r_t; undocumented, untested | `tidal.py:119-165` | Docstring paragraph + recommend re-virialization; optional validation test | +| S5 | Med | `q_approx` calibration (1.375696) from ONE point (uniform sphere, N=500); ~0.1 over-read in concentrated regime documented only in test docstrings, not API docs; N-dependence untested | `diagnostics/q_approx.py:36-39` | Calibration sweep over (profile, N); regime statement in API docstring | +| S6 | Med | `energy_sorted_segregation` docstring oversells "clean equilibrium": segregated mass-weighted density ≠ parent profile, so self-consistent potential differs; the validation quietly applies a global virial rescale that the docstring never tells users to do | `cluster/mass_segregation.py:6-11,56-58` | Add self-consistency caveat + "finalize with rescale" | +| S7 | Med | `MoePeriod` uses companion *frequency* per decade as a binary period *pdf* (folds outer tertiaries in; O-star f_mult=2.1), with flat clamped extrapolation outside logP∈[1,7]; undocumented | `moe_di_stefano.py:315-360` | Document; consider edge anchors | +| S8 | Low | `MoeEccentricity` extends p(e)∝e^η to full [0, e_max]; paper calibrates only to 0.8 e_max and M1 ranges 0.8–3, >7 M☉ | `binaries/eccentricity.py:254-280` | One-line caveat; optional taper | +| S9 | Low | `MoeDiStefano2017` (period-averaged) twin = invented Gaussian N(1, 0.03), eyeballed reductions — honestly disclosed, but it is the **default** `q_distribution` of `BinaryIMF`, whose docs say "full Moe+17 model" (oversell) | `moe_di_stefano.py:17-146`, `binary/imf.py:75,127-130` | Caveat in BinaryIMF docstring; never cite this class as MD17 in the paper | +| S10 | Low | `BinaryEnergyBudget` docstring claims Q_resolved "inflated"; measured: **deflated** (0.31 vs 0.50) — physics right, adjective wrong | `binaries/diagnostics.py:222-223` | Fix word | +| S11 | Low | `MassDependentBinaryFraction` steps are a Moe/Raghavan blend (disclosed) — cite as "derived from," not "from," Table 13 | `binary_fraction.py:40-62` | Citation hygiene | +| S12 | Low | King ρ̂(W) catastrophic cancellation: W=1e-6 rel. err 4.8e-4; **W=1e-8 returns negative** (−5.4e-20). Downstream clamps make impact ~nil (tidal-skin only) | `king.py:71-77` | Optional Taylor-series branch below W<1e-4 | +| S13 | Low | r_t non-differentiable in W0 (argmax+clamp, d ξ_t/dW0 ≡ 0) — documented & deferred; fine for shape inference, silently biases W0-inference through tidal observables | `king.py:259-266` | Keep warning in public docs | +| S14 | Low | Headline Engine B numbers partly single-seed: global Q=0.4976 is PRNGKey(0), N=30k, one seed (gate ±0.02); β-dev 0.028 is 4 seeds. Per-component machinery is properly multi-seed (18-seed evidence quoted) | `test_engine_b_physics.py:328-333` | Quote multi-seed mean ± scatter in the paper; archive the 18-seed run | +| S15 | Low | Zero rotation axis: comment claims NaN, code silently no-ops | `rotation.py:49-50` | Align comment/behavior; eager ValueError | +| S16 | Low | Per-group W_j is origin-dependent for subgroups; docstring claims origin independence; callers all pre-center but users won't know | `dynamics/virial.py:201-215` | Fix docstring; center internally | +| S17 | Low | Maschberger `_primitive` breaks at β=1 or α=1 (non-default), undocumented | `imf/maschberger.py` | Document | +| S18 | Low | `find_bound_pairs` softening param doesn't soften E_internal (correct choice, one docstring line missing) | `binaries/diagnostics.py` | Document | + +### 4.3 Ambiguous-conventions items +- Table 13 node masses (1.0/3.2/6.7/12/20 M☉) for the spectral-type bins are a progenax interpolation convention, not from the paper — document it. +- Moe Eqs. 17–18 literal text (η fits) verified only against Table 13 anchors, not the equation text — read pp. 41–43 of the PDF before the paper. +- "IGIMF" naming anywhere would be wrong: the implemented model is the Marks+2012/Jeřábková+2018 *environment-dependent stellar IMF*, not galaxy-wide IGIMF integration. + +--- + +## 5. JAX / Differentiability / Performance Audit + +### 5.1 Empirically verified gradient coverage (FD cross-checked this session) + +| Entry point | Param | Status | +|---|---|---| +| `PlummerProfile.sample_positions` | r_h | ✅ FD ratio 1−1e-11 | +| `build_spatial_ic` (Plummer, end-to-end incl. virial scale, COM) | r_h | ✅ FD ratio 1−1e-10 | +| `KingProfile.from_W0_rc` → positions (through diffrax ODE + CDF) | W0 | ✅ FD ratio 1.000006 (W0≲9.5) | +| `KingVelocityDF.sample_velocities` | W0 | ✅ FD ratio 0.999994 | +| `EFFVelocityDF` (Eddington) / OM | γ / r_a | ✅ FD ratios 1±6e-7 | +| `PowerLawIMF.sample` | α interior | ✅ FD ratio 1−1e-10; **α=1.0 exactly → NaN** (R10) | +| `KeplerElements.to_state` | e ∈ {0.05, 0.9, 0.99, 0.999} | ✅ FD ratios 1.00000000 | +| `q_approx` (naive + fast Morton path) | positions | ✅ finite, nonzero | +| `apply_tidal_truncation` (STE custom_jvp) | r_t | ✅ surrogate ≈ coarse FD of the step | +| Michie DF, LIMEPY tables, MultiComponentCluster | W0/g/w_j/r_a_j | Claimed; repo has AD-vs-FD tests; **spot-check recommended pre-release** (~2 min each) | +| `gaussian_random_field` (experimental) | β | ❌ **NaN** (where-NaN at DC mode; off the production inference path) | + +### 5.2 Findings + +| # | Sev | Class | Finding | Where | +|---|---|---|---|---| +| J1 | High | correctness | PowerLawIMF α=1 NaN grad (=R10); port `exp_safe` from `differentiable.py:47-54` | `power_law.py` ×4 sites | +| J2 | High | docs-correctness | 16 documented symbols missing from `progenax.__init__`; `IGIMF`/`EnvironmentIMF` don't exist (=R7) | `__init__.py` | +| J3 | Med | correctness (exp.) | GRF NaN grad w.r.t. β at kmag=0 dead branch; standard double-where fix | `gravoturb_fdf/field/field.py:63` | +| J4 | Med | correctness | Traced-W0 `_auto_ode_domain` fallback silently pins ξ_t to the domain edge for W0≳10 under jit/grad — wrong answer, no warning, exactly where the user can't see it. Add an `is_pinned` diagnostic or raise | `king.py:159-166,282` | +| J5 | Low | edge | `virial_scale` returns NaN for T=0 input; sibling `rescale_velocities_to_virial` has a (different) guard — the near-duplicates disagree | `builders.py:256-259` vs `virial.py:251` | +| J6 | Low | perf | `q_approx(method="auto")`: `lax.cond` on a *static* Python N compiles both branches every shape (measured 0.4 s); use Python `if` | `q_approx.py:239,270-275` | +| J7 | Low | perf | `_sample_cluster_arrays` keeps G static under filter_jit → recompile per distinct float G; `king_df.py:178` already shows the `jnp.asarray(G)` fix | `cluster/sampling.py:23-35` | +| J8 | Low | API | `inverse_cdf_draw` zero-weight clamps to grid[-1]; contract enforced by caller convention only (all 5 current sites comply) — add `fallback=` to make it structural | `numerics.py:37-55` | +| J9 | Low | edge | Kepler solver unguarded for e≥1 (silent garbage); document | `kepler.py` | +| J10 | Info | acceptable | O(N²) energy kernels: blocked scan + checkpoint, 32.8 GB→0.12 GB at N=2e4 measured by authors. **Leave them alone** — model implementation | +| J11 | Info | acceptable | Eddington d²ρ/dΨ² via FD on the Ψ grid — empirically clean (γ/r_a grads match FD to 5e-7); no action | +| J12 | Info | strength | **No rejection sampling anywhere** — all samplers exact inverse-CDF (Beta/analytic PPF/tabulated). No acceptance-rate pathology class exists in this codebase | + +**Hygiene verified clean:** no `lax.while_loop` (diffrax internals excepted), no numpy in traced paths, no PRNG key reuse, no `.item()`/`float()` on tracers, deliberate documented non-differentiable surfaces all correctly fenced (argsort diagnostics, Stars/TotalMass budgets, r_t). + +--- + +## 6. Architecture and API Review + +**Organization: strong.** Module layout maps cleanly onto physics; the 2026-06 redesign removed debt (string dispatch, buggy `populations.py`) instead of shimming it. Docstring sample (10 public symbols): genuinely excellent — units in brackets, equation-numbered references, gradient-behavior notes. + +**Findings:** + +| # | Sev | Finding | Where | +|---|---|---|---| +| A1 | Med | **Duplicate `VelocityDF` Protocol** — two structurally identical runtime-checkable definitions; will drift | `protocols.py:54-83` vs `kinematics/api.py:53-70` | +| A2 | Med | **Units-policy split**: the VelocityDF protocol bakes in `G=None→DEFAULT_UNITS` on *core* sampling entry points, while `build_spatial_ic` requires G — the package contradicts its own MANDATORY explicit-units policy. (Positive: zero `get_G()`/global-context violations found.) Decide once, pre-1.0 | `protocols.py:68` et al. | +| A3 | Med | KingVelocityDF over-parameterized (S1) — same r_t inconsistency at the API level | `king_df.py:107-115` | +| A4 | Med | Size-limit mandate violated: 8 functions >100 LOC (worst: `energy_sorted_segregation` 166, `two_body_kepler` 141, `env_to_imf_params` 137), 4 files >500 LOC (`limepy_multimass.py` 617, `king.py` 559, `multicomponent.py` 552, `builders.py` 502). Either split the two that genuinely read badly (`env_to_imf_params`, `energy_sorted_segregation`) or amend the rule | various | +| A5 | Low | Ghost cache dirs `src/progenax/gravoturb/` and `src/progenax/cluster/fdf_density/` (only `__pycache__`/mypy debris, git-ignored, **verified absent from the wheel**) — delete; they made this audit suspect stale code. Note `cluster/turbulence.py` is NOT stale (consumed by environment-IMF; correctly documented) | — | +| A6 | Low | Protocols advisory-only: `build_spatial_ic` never isinstance-checks; wrong objects die deep in sampling with AttributeError. One friendly boundary check would cost nothing | `builders.py` | +| A7 | Low | `UniformSphereProfile` exported from `progenax.profiles` but not top level (every other profile is); `CompanionModel.sample` missing return annotation; `ids=None` iff `id_offset==0` asymmetry; `TotalMass` top-up `while True` unbounded | `profiles/__init__.py:46`, `protocols.py:273-282`, `builders.py:317,364-375` | + +**Extensibility:** new SpatialProfile/VelocityDF/IMF/CompanionModel = easy (protocols + reusable Eddington/numerics kernels extracted in the right places). New multimass engine = hard but appropriately so. Experimental quarantine verified at the wheel level — no core leakage of `gravoturb_fdf`. + +--- + +## 7. Module-by-Module Completeness Table + +| Module | Status | Key evidence / caveats | +|---|---|---| +| Plummer profile | **Release-ready** | Exact algebra; CDF validated; u-clamp negligible & documented | +| King profile | **Release-ready W0≲8; under-verified W0≳9** | R4 CDF-grid core errors (measured); fix n_grid scaling + add core-ρ test | +| EFF profile | **Release-ready** (positions) | Trapezoid CDF gated + order-2 convergence measured; untested r_t/a≫30 | +| Michie profile | **Release-ready with doc caveats** | Closed-form integral re-derived; shares R4 grid pattern | +| Plummer DF (+OM) | **Release-ready** | Beta(3/2,9/2) exact; 63/4 re-derived; Merritt bound enforced | +| King DF | **Release-ready** (after S1 doc fix) | True lowered Maxwellian; unscaled Q=0.5 at W0=7 & 12 | +| EFF DF | **Release-ready with doc caveats** | S2: default γ=3 sub-virial by construction; pin it | +| Michie DF | **Release-ready** | Factorization re-derived; β self-oracle + g=1 LIMEPY identity; one external β point would close it | +| OM anisotropy | **Release-ready** | Full Q-variable inversion; β by construction + a-posteriori | +| Rotation transforms | **Technically functional, scientifically under-verified** | S3: equilibrium-breaking overlay undocumented | +| Multimass Engine A (LIMEPY) | **Release-ready** | Parity vs canonical LIMEPY (pinned SHA, provenance-stamped); per-component Q oracles; implicit-VJP grads vs FD | +| Engine B (Eddington, shared Ψ) | **Release-ready with doc caveats** | A-vs-B anchor; EFF≡Plummer identity; hybrid-sampling edge offset *predicted and gated*; S14 single-seed headline | +| Mass segregation | **Release-ready with doc caveats** | S6 self-consistency caveat missing; algorithm sound, regression-locked | +| Tidal truncation | **Release-ready with doc caveats** | r_J validated vs L1 root-find; S4 equilibrium break undocumented | +| Builders / IC composition | **Release-ready after R6 fix** | Budget semantics unambiguous & tested; G threading audited | +| IMF: PowerLaw/Chabrier/Maschberger/Truncated/IMFParams | **Release-ready** (R10 = 3-line fix; `differentiable.py` is the gold standard) | Paper-verified parameters throughout | +| `BaseIMF` mass-target sampling modes | **Do not release yet** | R5: silent 30% shortfall, zero tests | +| Environment-dependent IMF | **Release-ready with doc caveats** | Erratum-corrected; Eq.-9 deviation documented; kill "IGIMF" naming (R7) | +| Binary orbital mechanics (Kepler, elements, periods) | **Release-ready** | Machine-precision + correct grads to e=0.999 | +| Period/ecc distributions (Öpik/DM91/Sana/Thermal) | **Release-ready** | Paper-verified; α=−1 limit correct | +| MoeDiStefano2017Full / MoeJointOrbit / MoeCompanions | **Do not call "faithful" until R3 fixed** | Table 13 exact; F_twin normalization wrong | +| MoeDiStefano2017 (period-averaged) / MoePeriod | **Plausible, needs stronger validation** | Honest approximations; S7/S9 caveats | +| Binary connector + diagnostics | **Release-ready with doc caveats** | COM exact; S10 wording; softening scope note | +| Binary population synthesis (`build_binary_cluster`) | **Release-ready** | Primary-IMF semantics preemptively documented; budgets tested | +| CW04 Q exact / q_approx | **Release-ready** / **needs stronger validation** | Table 1 anchors / S5 single-point calibration | +| Analytical test cases | **Release-ready** | Digit-checked; self-validating; honest ephemeris provenance | +| gravoturb_fdf (experimental) | **Experimental only** (correct labeling) | J3 GRF grad NaN off-path; AC tolerances are honest smoke reductions; full-strength AC run not gated anywhere | +| Docs & examples | **Do not release yet** | R7: README teaches deleted API; phantom classes; broken install commands | +| CI / packaging | **Do not release yet** | R1/R2: red CI; uninstallable wheel; no LICENSE | + +--- + +## 8. Test Suite Audit + +**What the suite genuinely proves** (best-in-class for research code): analytic Plummer σ²(r); King 1966 Table II; LIMEPY parity against the canonical GZ15 code at pinned SHA ef2a479 with provenance embedded in the .npz (numpy/scipy/python versions + config — verified by reading `provenance.npy`); measured-first frozen gates with written noise budgets (e.g. `conftest.py:86-89` documents the 40-seed calibration that tightened VIRIAL_RATIO 0.20→0.05); AD-vs-FD gradient tests at rel<1e-5; CW04/Allison/Küpper anchors. **The 5 suspiciously-tight tolerances were checked: none are seed-hacks** — each has a written derivation. + +**What it does not prove / weaknesses:** + +| # | Sev | Finding | +|---|---|---| +| T1 | **Critical** | CI red on main (=R1); ≥5 runs failed at dependency sync; nothing enforced anywhere | +| T2 | **Critical** | Headline validations slow-marked and never run in CI (=R8): LIMEPY parity, multimass equilibrium, A-vs-B anchor, Engine B AD-vs-FD — 47/1192 tests including the advertised trust anchors | +| T3 | High | No version matrix: Python 3.13 only vs `requires-python>=3.10`; JAX floor 0.4.20 never tested (lock has 0.6.2) | +| T4 | High | LIMEPY parity **skips silently** if the cache is absent (`test_limepy_reference_parity.py:62-64`) — a packaging mistake converts the strongest gate into a silent pass. Make it fail in strict/release mode | +| T5 | High | Stale counts: `tests/README.md` says 874 tests; CLAUDE.md says 1163; measured 1192 collected / 1145 not-slow | +| T6 | Med | Gradient tests bimodal: strong FD-anchored ones exist, but 7/9 in `test_jax_compatibility.py` assert sign/finiteness only — a wrong-by-2× gradient passes | +| T7 | Med | Integration tier ~2/3 smoke/tautology (COM-sums-to-zero true by construction). The dropped-G regression IS real (verified it catches the catastrophic case) but doesn't test G-*scaling* | +| T8 | Med | `tests/unit/test_numerics.py:15-61` bit-for-bit mirrors the implementation — refactoring guard, not correctness; needs an analytic-integral oracle | +| T9 | Med | Single-seed stochastic tests (EFF ordering, Λ_MSR, multimass) — **probed at fresh seeds this session: all pass with large margin**, so flakiness is a hygiene smell not a hazard; convert to 3-seed means | +| T10 | Med | Untested edges: EFF γ→2, e≥0.95 distributions, IMF α≈1, N∈{0,1}, W0→0; all four `BaseIMF` mass-mode samplers (=R5) | +| T11 | Low | ~11 `scripts/validate_*.py` never invoked by pytest — the project's own Definition-of-Complete validation tier enforced by convention only; `benchmark_batch_a.json` referenced by nothing | + +**Loosest tolerances (for tightening):** `2.0<Λ<12.0` (mass-seg); `±0.20` on α₃ (environment); `0.40=0.1,<0.2`. Remove jaxstroviz from progenax's lock graph entirely (validation-side only; it's what broke CI). +- Option B: vendor the ~2 modules into `progenax._vendor` — fastest standalone wheel, costs ecosystem coherence (shared UnitSystem objects with gravax). +- Option C: git-URL dep — fine for a GitHub soft launch; PyPI rejects URL deps. + +--- + +## 10. Pre-Release Action Plan + +### Blockers (must do) +1. **Fix CI** (R1): check out jaxstroviz(+fluxax) siblings in all 4 jobs, or better, remove jaxstroviz from progenax's dependency graph. Get main green. +2. **Add the nightly/release-tag slow lane** (R8): LIMEPY parity, A-vs-B, multimass equilibrium, full-strength gravoturb acceptance; strict reference-cache mode. +3. **Fix Moe F_twin normalization** (R3): unnormalized p ∝ p_pl + (ft/(1−ft))·I_B·U[0.95,1], renormalize over [q_min,1]; regression test: excess-twin/(q>0.3) = Table 13 F_twin at grid nodes. +4. **Fix high-W0 CDF grids** (R4): scale n_grid with ξ_t (or log-stretched grid) in king/michie/eff; add sampled core-density tests at W0 ∈ {9, 12, 15}. +5. **Fix or fence `sample_fixed_n`** (R5): reachability guard + tests for all 4 modes; soften "exactly". +6. **Fix `compute_stellar_radii`** (R6): adopt a cited fit (Demircan & Kahraman 1991 or Tout+1996 ZAMS); literature-anchored test. +7. **Port `exp_safe` to PowerLawIMF** (R10): 3 sites; grad test at α=1.0. +8. **README + installation rewrite** (R7/D3/D4/D5): current API only, MultiComponentCluster featured, real install commands, kill `progenax-legacy` claim, fix reused-key example, fix phantom IGIMF/EnvironmentIMF everywhere (incl. both CLAUDE.md files), reconcile `__init__` exports with documented API (decide: export the 14 IMF symbols or fix docs). +9. **Packaging** (R2/D1): LICENSE, metadata, sibling-dep decision (Option A recommended), `[diagnostics]` extra or lazy numpy/scipy imports (R9), CHANGELOG, version tag. +10. **Wheel-smoke CI job**: build wheel, clean venv, import progenax + progenax.diagnostics, run README snippet. + +### Important, non-blocking +- S1 King r_t consistency guard; S2 pin Q(γ=3); S3/S4/S6 equilibrium-caveat docstrings; J4 traced-W0 pinning diagnostic; J5 virial_scale guard + dedupe with rescale_velocities_to_virial; A1 dedupe VelocityDF protocol; A2 units-policy decision; A5 delete ghost dirs; T6 upgrade finite-only grad tests; T9 3-seed conversions; S5 q_approx calibration sweep; M-dwarf multiplicity + DM91 constants PDF check; version matrix CI. + +### Pre-paper verification (cheap insurance) +- Eyeball King 1966 Table II once more against `test_king_physics.py:267,397`. +- Read Moe Eqs. 17–18 text (pp. 41–43) for η fits + 3–7 M☉ interpolation prescription. +- Regenerate the LIMEPY reference cache once in the pinned env and diff. +- CW04 Table 1 + area definitions against the actual PDF. +- GZ15 §2.2 m̄ convention + η default direct read. +- Archive the 18-seed Engine B run behind the quoted 0.4953/0.5007 numbers. +- Spot FD-check Michie/LIMEPY/MultiComponentCluster gradients (3 × 2-min harness runs). +- One long N-body run per DF (King/Michie/EFF/Engine B) in gravax over ~10 t_cross asserting profile stability — the gold-standard equilibrium demonstration for the paper. + +--- + +## 11. Methods Paper Plan + +**Title options:** +1. "progenax: Differentiable Initial Conditions and Population Synthesis for Star Cluster Modeling and Inference" +2. "Differentiable Equilibrium Initial Conditions for Collisional Stellar Dynamics with progenax" +3. "progenax: A Composable, Differentiable Engine for Star Cluster Initial Conditions, Binary Populations, and Gradient-Based Inference" + +**Abstract framing:** IC generation has been treated as a fire-and-forget preprocessing step; progenax reframes it as a differentiable model component, enabling gradients of any downstream observable with respect to structural (r_h, W0, c, r_t, r_a), population (IMF slopes, binary fractions, Moe parameters), and environmental parameters — while matching or exceeding the physical fidelity of standard generators (true lowered-Maxwellian/Eddington equilibria, no virial rescaling; LIMEPY-parity multimass models). + +**Claims currently justified (cite the named tests):** +- True detailed equilibria sampled without external virial rescale (Q=0.5 unscaled: Plummer/King/Michie/EFF γ=5/Engine A/Engine B). +- Multimass LIMEPY-family models with reference-code parity at 1e-8–5e-4. +- Two independent equilibrium engines agreeing on King (A-vs-B σ/KS at 2–3e-4). +- OM anisotropy by construction with a-posteriori β(r) verification. +- End-to-end differentiability with AD-vs-FD verification through ODE solves and Eddington inversions. +- Faithful Kepler machinery to e→1 with correct gradients. +- Moe Table 13 transcription exactness (after R3 fix, the full P–q–e claim). + +**Claims requiring more validation before being made:** +- "Faithful Moe & Di Stefano 2017" — blocked on R3 + Eqs. 17–18 verification. +- High-concentration (W0>9) sampled structure — blocked on R4. +- Long-term N-body stationarity (add the gravax 10-t_cross runs). +- q_approx as a general-purpose Q estimator (S5 regime limits). +- Anything about gravoturb_fdf beyond "experimental, repo-only" unless promoting it deliberately. + +**Section outline:** +1. Introduction — ICs as inference bottleneck; differentiable-programming context. +2. Architecture — protocols, Equinox PyTrees, explicit units, composability (profile × DF × IMF × binaries). +3. Spatial profiles & velocity DFs — Plummer/King/Michie/EFF; Eddington inversion; OM; rotation overlays (with honest equilibrium caveats). +4. Multi-component clusters — Engine A (GZ15 multimass) and Engine B (density-defined shared-Ψ Eddington); mass segregation; tidal truncation. +5. Stellar populations — IMFs (paper-grounded constants incl. the Jeřábková Eq.-9 correction), binary populations (Moe P–q–e), the connector. +6. Differentiability — fixed-iteration solvers, double-where guards, STE truncation, non-differentiable surfaces (honest table). +7. Validation — external-reference parity (King Table II, LIMEPY, CW04), equilibrium gates, AD-vs-FD, N-body stationarity. +8. Inference demonstrations (§12 below). +9. Performance & scaling; 10. Discussion/roadmap (stellax/startrax/fluxax coupling, gravoturb_fdf as future work). + +**Key figures:** (i) architecture/composability diagram; (ii) ρ(r), σ(r), β(r) sampled-vs-analytic 4-panel per DF family; (iii) c(W0) vs King 1966 + LIMEPY parity residuals; (iv) Engine A-vs-B overlay; (v) gradient-verification panel (AD vs FD across parameters/entry points); (vi) posterior corners from the inference demos; (vii) Moe P–q–e joint recovery vs Table 13; (viii) N-body stationarity over 10 t_cross. + +**Key tables:** validated formula↔reference provenance table; gradient coverage map (from §5.1 — publish it); equilibrium Q table (the existing CLAUDE.md table, multi-seed); timing/compile-cost table. + +--- + +## 12. Recommended Science and Inference Demonstrations + +| Demo | Demonstrates | Validates | Figure | Difficulty | Priority | +|---|---|---|---|---|---| +| 1. Recover (r_h) of a Plummer sphere from particle positions via gradient descent / HMC | The core differentiable-inference loop; bias/variance vs N | Plummer profile+DF, build_spatial_ic grads | posterior + convergence trace | Easy (exists in spirit in tests) | **Release-critical** | +| 2. Recover (W0, r_c) of a King model — gradients through the diffrax ODE | The hard differentiability claim | King profile/DF, implicit grads | corner plot | Easy–moderate | **Release-critical** | +| 3. Model comparison: fit Plummer vs King vs EFF to one synthetic cluster, gradient-based MAP + evidence proxy | Composability as scientific capability | all three families | profile residual panel | Moderate | Paper-enhancing | +| 4. IMF slope α from a synthetic resolved population (incl. α near break masses) | Differentiable population inference | IMF subsystem, IMFParams | posterior vs truth across α grid | Easy | **Release-critical** | +| 5. Unresolved-binary bias on inferred IMF slope (with/without MoeCompanions) | A real survey systematic, only easy with this tool | binary connector + Moe model | Δα vs binary fraction | Moderate | Paper-enhancing (flagship science result) | +| 6. Binary-population inference: recover F_twin/γ/period params from resolved binaries | Moe machinery post-R3-fix | MoeDiStefano2017Full | corner | Moderate–hard | Paper-enhancing | +| 7. Concentration + anisotropy joint (W0, r_a) recovery from projected kinematics | OM machinery; degeneracy structure | Engine B / OM DFs | 2D posterior with β(r) overlay | Moderate | Paper-enhancing | +| 8. Two-component mass-segregated cluster: recover δ (equipartition parameter) | Engine A multimass | MultiComponentCluster | δ posterior | Moderate | Paper-enhancing | +| 9. IC-sensitivity through N-body evolution: d(final r_h)/d(initial Q) via gravax SymplecticIntegrator | Ecosystem differentiability end-to-end | gravax coupling | gradient flow vs FD through evolution | Hard (gravax-side) | Paper-enhancing | +| 10. Differentiable CMD forward model via fluxax (flux_from_mass) → IMF+binary posterior from photometry | The LSST-facing pitch | fluxax coupling | synthetic CMD + posterior | Hard | Paper-enhancing, defer if fluxax not ready | +| 11. Substructure: Q-parameter trajectories of collapsing subvirial clusters; q_approx as differentiable summary | diagnostics | CW04 Q + q_approx (post-S5) | Q(t) tracks | Easy | Paper-enhancing | +| 12. gravoturb_fdf β-headline 2D inference | The experimental arc | gravoturb_fdf | SBC + posterior | Already in progress | **Separate paper**, not this one | + +Minimum release-critical set: demos 1, 2, 4 — each is a notebook-sized exercise on machinery already validated, and together they substantiate the paper's central claim ("differentiable inference workflows impossible with traditional IC generators"). + +--- + +## 13. Final Verdict + +**Beta scientific software, ~2 weeks of mechanical work from Release Candidate.** The physics core would survive expert review today — the audit *failed to find a single error* in the headline equilibrium machinery (Plummer/King/Michie/EFF/Engine A/Engine B), and the validation culture (external references, pinned provenance, measured-first gates, AD-vs-FD) is better than most published astronomy packages. What would embarrass you publicly is entirely at the boundary: a red CI, an uninstallable wheel, a README teaching deleted APIs, two real but fixable science bugs in the binaries/IMF periphery (R3, R5), one quantified numerical defect at high W0 (R4), and a factual error in an auxiliary stellar mass–radius relation (R6). + +### Top 10 actions, in order +1. Fix CI (jaxstroviz checkout / de-lock) and get main green — nothing else is enforced until then. +2. Add the nightly/release-tag slow-test lane so the advertised trust anchors actually run. +3. Fix the Moe F_twin normalization (R3) + Table-13 regression test. +4. Fix high-W0 CDF grids (R4) + core-density tests at W0 ∈ {9,12,15}. +5. Guard/fix/test the `BaseIMF` mass-target sampling modes (R5). +6. Swap the `compute_stellar_radii` exponents with a cited fit (R6). +7. Port `exp_safe` to PowerLawIMF (R10) + α=1 gradient test. +8. Rewrite README + installation docs; reconcile `__init__` exports vs documented API; kill IGIMF/EnvironmentIMF and `progenax-legacy` claims (R7/D5). +9. Packaging: LICENSE + metadata + jaxstro-first publication decision + `[diagnostics]` dependency story (R2/R9/D1). +10. Pre-paper verification sweep: Moe Eqs. 17–18 PDF read, CW04/GZ15 PDF checks, LIMEPY cache regen, Michie/multimass gradient spot-FD, and one 10-t_cross gravax stationarity run per DF. diff --git a/docs/website/00-getting-started/installation.md b/docs/website/00-getting-started/installation.md index c58965a8..3be58988 100644 --- a/docs/website/00-getting-started/installation.md +++ b/docs/website/00-getting-started/installation.md @@ -4,40 +4,37 @@ description: Install progenax via UV (preferred) or pip — single-environment s --- # Installation -progenax is a JAX-native Python package distributed through the -[jaxstro monorepo](https://github.com/drannarosen/jaxstro-dev). The -recommended installation tool is **UV**, which is 10–100× faster -than pip for this codebase. +progenax is a JAX-native Python package. It is **not on PyPI**; it +(and its public sibling `jaxstro`) resolve from a standard +side-by-side checkout. The recommended installation tool is **UV**, +which is 10–100× faster than pip for this codebase. ## Quick install (UV — recommended) -```bash -git clone https://github.com/drannarosen/jaxstro-dev.git -cd jaxstro-dev/progenax -uv pip install -e ".[all]" -``` - -The `[all]` extra pulls in development tooling (pytest, black, -mypy), I/O (orbax-checkpoint, h5py), visualisation (matplotlib, -seaborn), and ML (optax, blackjax, numpyro). For minimal install: +Clone progenax and `jaxstro` next to each other (the `[tool.uv.sources]` +path dependency expects the sibling layout): ```bash -uv pip install -e . +git clone https://github.com/jaxstro/jaxstro.git +git clone https://github.com/jaxstro/progenax.git +cd progenax +uv pip install -e ".[dev]" ``` -For the broader jaxstro ecosystem (gravax, fluxax, etc.): +The `[dev]` extra adds the test/lint tooling (pytest, black, isort, +flake8, mypy). For a minimal runtime-only install: ```bash -cd jaxstro-dev -uv pip install -e ./jaxstro -e ./gravax -e ./progenax -e ./fluxax +uv pip install -e . ``` ## Pip alternative ```bash -git clone https://github.com/drannarosen/jaxstro-dev.git -cd jaxstro-dev/progenax -pip install -e ".[all]" +git clone https://github.com/jaxstro/jaxstro.git +git clone https://github.com/jaxstro/progenax.git +cd progenax +pip install -e ".[dev]" ``` ## Smoke test @@ -52,11 +49,11 @@ from progenax.profiles import PlummerProfile from progenax.kinematics import PlummerVelocityDF masses = jnp.ones(100) -key = jax.random.PRNGKey(0) +key_pos, key_vel = jax.random.split(jax.random.PRNGKey(0)) # never reuse a key profile = PlummerProfile(r_h=1.0) df = PlummerVelocityDF(r_h=1.0) -positions = profile.sample_positions(masses, key) -velocities = df.sample_velocities(positions, masses, key, G=STELLAR.G) +positions = profile.sample_positions(masses, key_pos) +velocities = df.sample_velocities(positions, masses, key_vel, G=STELLAR.G) print(f"Generated {positions.shape[0]} particles, mean radius {jnp.linalg.norm(positions, axis=1).mean():.3f} pc") ``` @@ -97,20 +94,14 @@ chain that benefits from GPU. * - Extra - Provides * - `[dev]` - - pytest, black, isort, flake8, mypy -* - `[io]` - - orbax-checkpoint, h5py, pandas -* - `[viz]` - - matplotlib, seaborn -* - `[astro]` - - astropy, gala -* - `[ml]` - - optax, blackjax, numpyro -* - `[all]` - - All of the above + - pytest, pytest-cov, pytest-xdist +* - `[experimental]` + - the repo-only `gravoturb_fdf` inference layer (blackjax, optax, arviz, scipy, flowjax) +* - `[diagnostics]` + - numpy + scipy, for the exact (non-differentiable) CW04 `compute_q_parameter` path ``` -Install only what you need: `uv pip install -e ".[viz,ml]"`. +Install only what you need: `uv pip install -e ".[dev,experimental]"`. ## Next step diff --git a/docs/website/20-architecture/ic-redesign-history.md b/docs/website/20-architecture/ic-redesign-history.md index f0d62d92..50eca4b6 100644 --- a/docs/website/20-architecture/ic-redesign-history.md +++ b/docs/website/20-architecture/ic-redesign-history.md @@ -106,8 +106,8 @@ import it. Users had two paths: - When to use * - **Migrate** to the new API - All new work, all production code. Most pre-redesign user code rewrites in $\sim 30$ minutes given the new API's verbosity -* - **Keep using legacy** - - For reproducing old paper results bit-for-bit. `pip install progenax-legacy` still works +* - **Reproduce pre-redesign results** + - Check out the pre-redesign commit from git history (the redesign is documented in the [development log](../90-development-log/2026-02-12-ic-redesign.md)). There is **no** separate legacy package or install. ``` The migration guide (originally at `docs/migration_2026_02_12.md`, diff --git a/docs/website/50-validation/index.md b/docs/website/50-validation/index.md index 32a98957..452deac4 100644 --- a/docs/website/50-validation/index.md +++ b/docs/website/50-validation/index.md @@ -18,6 +18,10 @@ tolerances. See the [**validation audit report**](audit-report.md) for trustworthiness tiers, limits, recommendations, and the remaining/incomplete roadmap (point-in-time, 2026-06-08). +Where this section proves each model reproduces its ground truth, the +[**science demos**](../60-science-demos/index.md) run the models *backwards* — +recovering cluster parameters from mock observations by differentiable MLE / NUTS. + ## Status dashboard The single place to see the V&V state. ✅ = tests pass *and* verified figure(s) diff --git a/docs/website/60-science-demos/cross-engine.md b/docs/website/60-science-demos/cross-engine.md new file mode 100644 index 00000000..d2d352d8 --- /dev/null +++ b/docs/website/60-science-demos/cross-engine.md @@ -0,0 +1,159 @@ +--- +title: Cross-engine agreement (B1) +description: "One King model built two ways — Engine A (DF-defined lowered-isothermal + coupled Poisson ODE) and Engine B (density-defined Eddington inversion) — shown to agree on rho(r), sigma_1d(r) and f(E) with freshly measured KS and dispersion gates." +--- + +# Cross-engine agreement (B1) + +progenax builds multi-component equilibria two genuinely different ways, and the +first demo checks that they describe the *same physical cluster*. + +- **Engine A** (`MultiComponentCluster.from_components`) starts from a + **distribution function**: the lowered-isothermal / LIMEPY family + {cite:p}`Gieles2015`. The density is whatever that DF produces in its own + self-consistent potential, found by integrating the coupled Poisson ODE. +- **Engine B** (`MultiComponentCluster.from_density_profiles`) starts from a + **prescribed density** (here a `KingProfile`), builds the shared potential by + one cumulative-trapezoid Poisson pass, and recovers the DF by **Eddington + inversion**. + +If both engines are correct, the King model `(W₀=5, g=1, r_c=1 pc)` built either +way must sample the same radial profile, the same dispersion profile, and the +same energy distribution. Nothing forces this but the physics — the two code +paths share no quadrature, no ODE, and no DF evaluation. + +## The two constructions + +**Engine A — the lowered-isothermal DF.** For a single mass component with +truncation order $g$, the DF is a lowered (energy-shifted) Maxwellian +({cite:t}`Gieles2015`, eq. 2), + +```{math} +:label: b1-lowered-df +f(E) \;\propto\; E_g\!\left(\frac{E}{s^2}\right), +\qquad +E_g(x) = e^{x}\,P(g, x)\ \ (x>0),\quad 0\ \text{otherwise}, +``` + +with $E = \Psi(r) - \tfrac12 v^2$ the (positive, bound) relative energy, $s$ the +velocity scale, and $P(g,x)=\gamma(g,x)/\Gamma(g)$ the regularized lower +incomplete gamma. The index is **$g$ itself** (the density uses $g+\tfrac32$). +The single parameter $g$ slides continuously between classical models — +$g=0$ Woolley, **$g=1$ King**, $g=2$ Wilson — and at $g=1$ the lowered +exponential collapses to $E_1(x) = e^x - 1$, the textbook King DF. The density +$\hat\rho(\Psi)=\int f\,d^3v$ closes Poisson's equation, solved as a coupled ODE +in $(\Psi, d\Psi/dr)$ — full derivation on the +[lowered-model-family](../10-theory/spatial-profiles/lowered-model-family.md) +page. + +**Engine B — Eddington inversion.** Given the *same* King density as a +prescribed shape, the shared potential $\Psi(r)$ follows from Poisson's equation +by quadrature, and the ergodic DF is recovered by the Eddington formula (Binney +& Tremaine 2008, eq. 4.46b), + +```{math} +:label: b1-eddington +f(E) = \frac{1}{\sqrt8\,\pi^2} +\left[\int_0^{E}\frac{d^2\rho}{d\Psi^2}\,\frac{d\Psi}{\sqrt{E-\Psi}} + + \frac{1}{\sqrt E}\left.\frac{d\rho}{d\Psi}\right|_{\Psi=0}\right]. +``` + +The two routes invert each other in principle — A maps DF$\to$density, B maps +density$\to$DF — so agreement is a true round-trip cross-validation. + +## What is measured (freshly run, ALL PASS) + +Both engines are sampled at $N = 2\times10^4$ with the **same** PRNG key, so the +draws are strongly correlated and the residuals isolate genuine *model* +differences rather than shot noise. Measured 2026-06-11 (`key=0`): + +```{list-table} +:header-rows: 1 + +* - Check + - Measured + - Gate +* - Tidal radius $r_t$ (A vs B) + - $10.8054$ pc (both) + - consistency +* - Theory virial $Q_j = T_j/|W_j|$ — Engine A + - $0.50019$ + - $0.5 \pm 3\times10^{-3}$ +* - Theory virial $Q_j$ — Engine B + - $0.50000$ + - $0.5 \pm 3\times10^{-3}$ +* - Radial Kolmogorov–Smirnov distance (sampled CDF) + - $1.5\times10^{-4}$ + - $< 0.02$ +* - Max $|\sigma_B/\sigma_A - 1|$ (interior bins) + - $3.5\times10^{-4}$ + - $< 0.02$ +* - $\rho(r)$ shape max fractional diff (diagnostic) + - $7.95\times10^{-13}$ + - — +* - $f(E)$ peak-matched max diff (diagnostic) + - $5.05\times10^{-2}$ + - — +``` + +The two independent engines agree on the sampled radial CDF to a KS distance of +$1.5\times10^{-4}$ and on the dispersion profile to $3.5\times10^{-4}$ — two +orders of magnitude inside the $0.02$ gate. The density *shape* agrees to +$8\times10^{-13}$ (essentially the floating-point floor, since both reduce the +same King profile after re-normalizing to the central value). + +```{note} +The virial $Q_j = 0.5$ check is **necessary but not sufficient**: the Clausius +identity $2T_j + W_j = 0$ holds for *any* positive DF in a consistent +$(\Psi, d\Psi/dr)$ pair, so it tests self-consistency, not inversion +correctness. The KS and $\sigma$-dev anchors carry the cross-validation weight — +they would fail if either engine's DF were wrong even though $Q_j$ stayed $0.5$. +``` + +## Figure + +:::{figure} figures/cross_engine.png +:label: sci-cross-engine +:width: 100% + +**Cross-engine King agreement** (`scripts/demo_cross_engine.py`, ALL PASS). +**(a)** $\rho(r)/\rho_0$ for Engine A (DF-defined) and Engine B +(density-defined) with a residual strip — coincident to $8\times10^{-13}$ after +central re-normalization. **(b)** $\sigma_{1\mathrm D}(r)$ from the two +independent oracles plus both sampled profiles ($N=2\times10^4$, same key); +$\max|\sigma_B/\sigma_A-1| = 3.5\times10^{-4}$. **(c)** $f(E)$ overlay on the +dimensionless energy axis $\hat E = (W_0/\Psi_0)E$ — the lowered-isothermal DF +shape (A) and the Eddington $f$-row (B), peak-matched. +::: + +## Caveats + +```{warning} +- **The $f(E)$ overlay is matched by peak, and shows a $\sim 5\%$ edge + residual.** The two engines parametrize energy and normalize the DF + differently; after matching the peak on a common dimensionless axis, the + largest discrepancy ($5.05\times10^{-2}$) sits at the truncation edge / central + cusp, where the lowered-exponential cutoff and the Abel-inverted tail differ in + shape. It is reported as a **diagnostic**, not gated — the sampled-CDF (KS) and + dispersion gates carry the agreement claim, and a $5\%$ DF-shape edge mismatch + is invisible in those sampled statistics. +- **This is one model, not a sweep.** A single concentration ($W_0=5$) at one + resolution is checked. The wider parameter agreement is established separately + by the [Engine B validation suite](../50-validation/engine-b-eddington.md), + whose King A-vs-B anchor measures $0.0002$ / $0.0003$ across its own headline + mix. +``` + +## How to run + +```bash +env -u VIRTUAL_ENV uv run --no-sync python scripts/demo_cross_engine.py +``` + +## References + +The lowered-isothermal / LIMEPY family is {cite:t}`Gieles2015`; the King model +{cite:t}`King1966`; Eddington inversion and the ergodic DF follow Binney & +Tremaine (2008). Engine internals: the +[multi-component theory pages](../10-theory/populations/index.md) and the +[Engine B validation suite](../50-validation/engine-b-eddington.md). diff --git a/docs/website/60-science-demos/halo-core.md b/docs/website/60-science-demos/halo-core.md new file mode 100644 index 00000000..0b022cb2 --- /dev/null +++ b/docs/website/60-science-demos/halo-core.md @@ -0,0 +1,164 @@ +--- +title: Halo + core recovery (B3) +description: "Engine B inference: a Plummer halo (Osipkov-Merritt anisotropic) plus an EFF gamma=5 core, sampled in shared-potential equilibrium, with the mass split t, anisotropy radius r_a and halo scale r_h recovered jointly by MLE + Gauss-Newton Fisher from binned sigma(r) and the halo beta(r)." +--- + +# Halo + core recovery (B3) + +A real cluster is not one population. This demo builds a **two-family** cluster — +a spatially extended, kinematically **anisotropic halo** plus a compact, +**isotropic core** — in one shared self-consistent potential, and asks whether +three structural parameters can be recovered jointly from its kinematics: + +- $t$ — the **halo mass fraction** (the core carries $1-t$), +- $r_a$ — the halo's **Osipkov-Merritt anisotropy radius**, +- $r_h$ — the halo's **Plummer half-mass radius**. + +The truth model (the science-headline mix of the +[Engine B validation suite](../50-validation/engine-b-eddington.md)): + +```{math} +:label: b3-truth +\underbrace{\text{Plummer}(r_h=2\,\text{pc})}_{\text{halo, } 60\%,\ m_j=0.5\,\Msun,\ r_a=3\,\text{pc}} +\ +\ +\underbrace{\text{EFF}(a=0.8\,\text{pc},\,\gamma=5,\,r_t=9\,\text{pc})}_{\text{core, } 40\%,\ m_j=1.0\,\Msun,\ \text{isotropic}} +``` + +sampled with $N = 3\times10^4$ stars. The EFF core shape $(a,\gamma,r_t)$, the +stellar-mass labels $m_j$, and the core's isotropy are held fixed; only +$(t, r_a, r_h)$ are recovered. + +## The physics being recovered + +**Anisotropy (Osipkov-Merritt).** The halo's DF depends on energy and angular +momentum only through $Q = E - L^2/(2 r_a^2)$, which makes the velocity +ellipsoid isotropic in the centre and increasingly radial outside $r_a$. The +**Binney anisotropy** has the closed form + +```{math} +:label: b3-beta +\beta(r) \equiv 1 - \frac{\sigma_t^2}{2\sigma_r^2} + = \frac{r^2}{r^2 + r_a^2}, +``` + +so $\beta\to 0$ as $r\to 0$ (isotropic) and $\beta\to 1$ as $r\to\infty$ (radial). +The core, with $r_a\to\infty$, has $\beta\equiv 0$. The OM curve {cite:p}`Merritt1985` +is the second observable, and it is what pins $r_a$. + +**Dispersion.** The 1-D dispersion of each component is built from the speed +moments of its Eddington DF $f_j$ in the shared potential, with the OM stretch +folded in: + +```{math} +:label: b3-sigma +\sigma_{1\mathrm D, j}^2(r) = \frac{\langle v^2\rangle_j(r)}{3}, +\qquad +\langle v^2\rangle_j(r) = \frac{m_{2,j}(r)}{m_{0,j}(r)} +\left(\frac13 + \frac23\,\frac{1}{1 + r^2/r_{a,j}^2}\right), +``` + +where $m_{p,j}(r) = \int_0^{\sqrt{2\Psi}} w^{p} f_j\!\big(\Psi - \tfrac12 w^2\big)\,dw$ +are the stretched-frame speed moments. This is the **verbatim** oracle the Engine +B validation suite checks against sampled velocities. Physical units come from +the engine's own velocity scale $v = \sqrt{G M_{\rm sampled}/(4\pi\mu)}\,s$, +anchored to the *measured* total mass (never an input $M$). + +**Why $\sigma(r)$ and $\beta(r)$ together recover all three.** The halo scale +$r_h$ sets where $\sigma(r)$ falls off; the mass split $t$ sets the relative +amplitude and shape of the two components' dispersion profiles; $r_a$ is read +almost entirely from the rising $\beta(r)$. The three are nearly orthogonal — +borne out by the near-zero Fisher correlations below. + +## Result — joint MLE + Fisher (freshly run, ALL PASS) + +Measured 2026-06-11 ($N=3\times10^4$, three dispersed Adam starts; exit 0): + +```{list-table} +:header-rows: 1 + +* - Parameter + - Truth + - $\hat\theta \pm \hat\sigma$ + - Pull $(\hat\theta-\theta_{\rm true})/\hat\sigma$ +* - $t$ (halo mass fraction) + - $0.600$ + - $0.602 \pm 0.008$ + - $+0.22$ +* - $r_a$ (halo OM radius) + - $3.000$ pc + - $3.051 \pm 0.053$ pc + - $+0.96$ +* - $r_h$ (Plummer half-mass) + - $2.000$ pc + - $2.015 \pm 0.029$ pc + - $+0.51$ +``` + +All three within $1\sigma$, comfortably inside the $3\sigma$ recovery gate. +Supporting diagnostics, all measured in the same run: + +- **Robust optimum.** All three dispersed initializations converged to the + *identical* loss $19.1479$ — the recovered minimum is not an initialization + artifact. +- **Realizability.** Rebuilding Engine B at $\hat\theta$ gives DF positivity + margins $f_{\min,j} = [0.085,\ 1.16\times10^{-4}]$, both $\ge -10^{-3}$ — the + recovered model is a genuine equilibrium, not a negative-DF fiction. +- **Self-consistency.** The analytic prediction at truth matches the binned data + to $\max|{\rm dev}/{\rm SE}| = 2.34$ over all 48 populated cells — no + systematic oracle/binning bias. +- **Near-orthogonality.** Gauss-Newton Fisher correlations + $\rho(r_a, r_h) = +0.03$, $\rho(t, r_a) = -0.04$ (condition number $13.4$): + the three parameters are independently constrained. + +## Figure + +:::{figure} figures/halo_core.png +:label: sci-halo-core +:width: 100% + +**Halo + core recovery** (`scripts/demo_halo_core.py`, ALL PASS). **(a, b)** +Per-component $\sigma_{1\mathrm D}(r)$: mock data (points, finite-$N$ error bars) +with the best-fit binned-expectation curve and the truth curve, for the Plummer +halo and EFF core. **(c)** Halo Binney anisotropy $\hat\beta(r)$ rising along the +OM curve $r^2/(r^2+r_a^2)$ ($\hat r_a = 3.05$ pc vs truth $3.0$), with the +isotropic-core reference at $\beta=0$. **(d)** $2\sigma$ Fisher ellipses in +$(r_a, r_h)$ (and inset $(t, r_a)$): the truth ★ sits inside, near the MLE; the +ellipses are nearly axis-aligned, reflecting the near-zero correlations. +::: + +## Caveats + +```{warning} +- **Clean mock, shared family.** Truth and fit are the same Engine B model; the + data are the full 3-D $\sigma(r)$ and $\beta(r)$ with no projection, no + measurement error, and no incompleteness. This isolates the *inference + machinery*, not observational recovery. +- **The domain $r_t$ is fixed, not recovered.** The shared potential's outer + radius is set to the (fixed) EFF extent $r_t = 9$ pc — both at truth and inside + the traced fit — so a recovered $r_h$ never has to re-derive the domain. This + is the model the truth draw uses; it is a construction choice, stated plainly. +- **Hard-truncation edge.** The EFF core is hard-truncated ($\rho(r_t)>0$), an + edge offset no ergodic $f(E)$ carries exactly; the resulting small + sub-equilibrium is quantified (not eliminated) on the + [Engine B validation page](../50-validation/engine-b-eddington.md). It does not + bias the parameter recovery, which is gated on the $3\sigma$ pulls above. +- **Cost is MLE-compile-dominated.** End-to-end $\sim 3.7$ min: the warm + likelihood-and-gradient eval is $0.14$ s, but each of the three Adam runs is a + jit-compiled 400-step `lax.scan` whose body rebuilds Engine B and backprops, so + the three *scan compiles* dominate. +``` + +## How to run + +```bash +env -u VIRTUAL_ENV uv run --no-sync python scripts/demo_halo_core.py +``` + +## References + +The EFF profile is {cite:t}`ElsonFallFreeman1987`; the Plummer model +{cite:t}`Plummer1911`; Osipkov-Merritt anisotropy {cite:t}`Merritt1985`; +Eddington inversion follows Binney & Tremaine (2008). The Engine B construction +and its realizability/anisotropy validation are documented at +[Engine B (Eddington)](../50-validation/engine-b-eddington.md) and the +[multi-component theory pages](../10-theory/populations/index.md). diff --git a/docs/website/60-science-demos/imf-equipartition.md b/docs/website/60-science-demos/imf-equipartition.md new file mode 100644 index 00000000..d8003c0f --- /dev/null +++ b/docs/website/60-science-demos/imf-equipartition.md @@ -0,0 +1,268 @@ +--- +title: IMF + equipartition recovery (B2) +description: "The self-consistent joint demo: one IMF high-mass slope alpha drives both the observed mass histogram and the equipartition mass groups of a LIMEPY multimass cluster, so (alpha, delta, W0) are recovered jointly by MLE + Fisher + NUTS — with the kinematics-only degeneracy, a wrong-IMF bias curve and an alpha_true robustness grid." +--- + +# IMF + equipartition recovery (B2) + +This is the flagship demo: it couples a cluster's **mass function** to its +**internal kinematics** through a single physical parameter, and shows that +parameter is recoverable from both channels at once. + +## The self-consistent physics + +A multimass LIMEPY cluster {cite:p}`Gieles2015` is built from an +[initial mass function](../10-theory/imfs/index.md). progenax bins the IMF into +$J=4$ mass groups; each group $j$ has a representative mass $m_j$ and a population +weight set by the IMF. **Two-body relaxation** drives the cluster toward energy +equipartition — heavier stars sink and slow down — encoded in the multimass DF by +a per-group velocity scale, + +```{math} +:label: b2-equipartition +\mu_j = \frac{m_j}{\bar m}, +\qquad +s_j = s\,\mu_j^{-\delta}, +\qquad +\delta = \tfrac12 \Rightarrow m_j s_j^2 = \text{const (full equipartition)}, +``` + +with $\bar m$ the central-density-weighted mean mass and $\delta$ the +**equipartition degree** {cite:p}`Peuten2017`. Real clusters reach only +*partial* equipartition ($\delta < \tfrac12$); the heaviest stars approach +$\sigma\propto m^{-1/2}$ while the light stars saturate at an escape-speed +ceiling — the Bianchini relation, derived directly from this same DF on the +[multimass equipartition](../10-theory/spatial-profiles/multimass-equipartition.md) +theory page. + +**The high-mass IMF slope.** The observed masses follow a Maschberger (2013) IMF +whose high-mass behaviour is the power law + +```{math} +:label: b2-imf +\frac{dN}{dm}\Big|_{m\gg m_c}\ \propto\ m^{-\alpha}, +\qquad \alpha_{\rm true} = 2.3\ (\approx\text{Salpeter}), +``` + +(full smooth form and characteristic mass: {cite:t}`Maschberger2013`). **Here is +the coupling that makes the demo self-consistent:** the *same* $\alpha$ that sets +the slope of the observed mass histogram also sets the masses $m_j$ and weights of +the equipartition groups in @b2-equipartition — so $\alpha$ is constrained by the +**masses** (the histogram) *and* by the **kinematics** (the group-by-group +$\sigma_j(r)$). The three recovered parameters are +$\theta = (\alpha,\ \delta,\ W_0)$, with $W_0$ the central dimensionless +concentration. + +## The joint likelihood + +Two channels, summed (`scripts/demo_delta_recovery.py`): + +```{math} +:label: b2-joint +\ln\mathcal L(\alpha,\delta,W_0) += \underbrace{-\tfrac12\sum_{j,k} w_{jk} + \Big(\tfrac{\hat\sigma_{jk}-\sigma_{jk}^{\rm pred}}{\mathrm{SE}_{jk}}\Big)^2}_{\text{kinematics: per-group }\sigma_j(r)} +\ +\ +\underbrace{\sum_{i=1}^{N}\ln f_{\rm Masch}(m_i;\alpha)}_{\text{masses: the IMF histogram}} . +``` + +The kinematic predictor is the differentiable Engine A multimass oracle rebuilt +from $\theta$ inside the traced loss (the `find_alpha_for_masses` eigenvalue solve +is differentiable in $\alpha,\delta,W_0$); the mass term is the analytic, +normalized Maschberger log-pdf — fully differentiable in $\alpha$. + +```{note} +**Observed masses are a *global* IMF sample (clean-mock Option A).** The kinematic +groups carry only the $J=4$ representative labels $m_j$, so the per-star mass +likelihood is evaluated on a *separate* global draw $m_i \sim +\mathrm{Maschberger}(\alpha_{\rm true})$ over the full mass range, with **no** +per-star mass↔group correlation modeled. One $\alpha$ still drives both channels +at the population level; the demo does not claim to model the joint mass–velocity +covariance of individual stars. +``` + +## Result 1 — joint recovery (freshly run, ALL PASS) + +Measured 2026-06-11 ($N=10^5$, three dispersed Adam starts; exit 0): + +```{list-table} +:header-rows: 1 + +* - Parameter + - Truth + - $\hat\theta \pm \hat\sigma$ + - Pull +* - $\alpha$ (IMF high-mass slope) + - $2.300$ + - $2.293 \pm 0.004$ + - $-1.68$ +* - $\delta$ (equipartition degree) + - $0.400$ + - $0.397 \pm 0.034$ + - $-0.08$ +* - $W_0$ (central concentration) + - $5.000$ + - $4.990 \pm 0.021$ + - $-0.48$ +``` + +All within $3\sigma$ (the heaviest, sparsest group $j=3$ has occupancy $1153 \ge +300$, so the fit is not starved). The MLE is a robust optimum — the two best of +the three dispersed starts agree in loss to $3\times10^{-3}$. + +:::{figure} figures/imf_equip_fit.png +:label: sci-b2-fit +:width: 100% + +**Joint MLE fit** (`demo_delta_recovery.py`). **(a)** Per-group binned +$\hat\sigma_{1\mathrm D,j}(r)$ (points, finite-$N$ errors) with the best-fit +binned-expectation curves — the equipartition ordering ($\sigma$ decreasing with +group mass) is visible. **(b)** The observed mass histogram with the fitted +Maschberger pdf at $\hat\alpha$. +::: + +## Result 2 — the degeneracy the mass channel breaks + +Kinematics *alone* cannot cleanly separate $\alpha$ from $\delta$: a steeper +$\sigma(m)$ can come from a larger equipartition degree *or* from an $\alpha$ that +reweights the group masses. The Fisher information quantifies it (freshly +measured): + +```{list-table} +:header-rows: 1 + +* - Quantity + - Kinematics-only + - Joint (+ mass channel) +* - $\rho(\alpha,\delta)$ correlation + - $-0.265$ + - $-0.059$ +* - $\sigma_\alpha$ + - $0.0192$ + - $0.0041$ +* - $\sigma_\delta$ + - $0.0357$ + - $0.0344$ +* - $(\Delta\chi^2{=}4)$ ellipse area + - $8.30\times10^{-3}$ + - $1.78\times10^{-3}$ +``` + +Adding the mass histogram **pins $\alpha$ $4.7\times$ tighter** (ellipse area +ratio $4.67$) and collapses the $(\alpha,\delta)$ correlation from $-0.265$ toward +zero. Crucially the mass channel *pins* $\alpha$ rather than *rotating* the +ellipse — $\delta$'s width barely moves ($0.0357\to0.0344$); the gain is almost +entirely in $\alpha$. + +:::{figure} figures/imf_equip_fisher.png +:label: sci-b2-fisher +:width: 90% + +**Fisher degeneracy panel.** $\Delta\chi^2=4$ ($86.5\%$ in 2-D) ellipses in +$(\alpha,\delta)$: the broad, tilted kinematics-only ellipse vs the compact joint +ellipse. The mass channel removes the degeneracy along $\alpha$. +::: + +## Result 3 — full posterior (NUTS) + +A vendored blackjax No-U-Turn sampler (300 warmup + 600 samples) draws the full +posterior. Recorded measured run (`--run-nuts`, $\sim 52$ min wall): + +- **0 divergent transitions**; +- posterior mean within **$1\sigma$ of the MLE per parameter** + ($\alpha\ -0.05$, $\delta\ +0.27$, $W_0\ -0.38\ \sigma$); +- the target is made flat in $\theta$ by adding the $\sum\ln(d\theta/dz)$ + reparametrization Jacobian, so the corner is a faithful posterior, not a + box-edge artifact. + +:::{figure} figures/imf_equip_corner.png +:label: sci-b2-corner +:width: 85% + +**NUTS posterior corner** in $(\alpha,\delta,W_0)$ with the MLE and truth +overlaid; unimodal, no divergences, posterior mean on the MLE. +::: + +## Result 4 — wrong-IMF bias + robustness grid + +What if you **assume the wrong** $\alpha$ and refit only the kinematics? Freezing +$\alpha$ at a grid of wrong values and refitting $(\delta, W_0)$ from kinematics +alone (5 seeds each) gives the bias curve. The recovered $\hat\delta$ **peaks at +the true $\alpha$** ($\hat\delta = 0.42 \approx \delta_{\rm true}$ at +$\alpha_{\rm assumed}=2.3$) and is biased **low** on either side — assuming the +wrong IMF slope corrupts the equipartition measurement. The linear sensitivity +$d\hat\delta/d\alpha = -0.005 \pm 0.020$ (seed-ensemble SE) is near zero precisely +*because* the response is peaked, not monotonic — it is **reported, not gated** +(no published reference value to assert against). + +The companion **robustness grid** regenerates fresh truth datasets at +$\alpha_{\rm true}\in\{1.9, 2.3, 2.7\}$ and refits the full joint $(\alpha,\delta, +W_0)$ — every case recovers within $3\sigma$: + +```{list-table} +:header-rows: 1 + +* - $\alpha_{\rm true}$ + - $\hat\alpha$ + - $\hat\delta$ + - $\hat W_0$ + - max $|$pull$|$ +* - $1.9$ + - $1.903 \pm 0.003$ + - $0.411 \pm 0.033$ + - $4.984 \pm 0.019$ + - $0.97$ +* - $2.3$ + - $2.300 \pm 0.004$ + - $0.396 \pm 0.033$ + - $5.011 \pm 0.021$ + - $0.51$ +* - $2.7$ + - $2.700 \pm 0.005$ + - $0.428 \pm 0.030$ + - $5.006 \pm 0.024$ + - $0.96$ +``` + +:::{figure} figures/imf_equip_bias.png +:label: sci-b2-bias +:width: 100% + +**Wrong-IMF bias (left)** — $\hat\delta(\alpha_{\rm assumed})$ with seed scatter, +peaking at truth; **robustness grid (right)** — joint recovery across three +$\alpha_{\rm true}$, all within $3\sigma$. +::: + +## Caveats + +```{warning} +- **Clean-mock, Option A masses** — observed masses are a global IMF sample + uncorrelated with the kinematic group (see the note above); no measurement + errors, projection, or mass-dependent completeness. +- **The wrong-IMF curve is a sensitivity *measurement*, not a calibrated test** — + reported with seed-ensemble uncertainty, never gated. +- **NUTS is non-divergent but not SBC-calibrated** — the credible-interval + coverage has not been checked by simulation-based calibration; that is future + work. +- **$\delta$ is fit on $[0, 0.7]$** — $\delta\gtrsim 0.9$ approaches the Spitzer + instability (heavy stars decouple) and hard-crashes the equilibrium ODE solve, + so the box stops short of it by construction. +``` + +## How to run + +```bash +env -u VIRTUAL_ENV uv run --no-sync python scripts/demo_delta_recovery.py # MLE + Fisher (minutes) +env -u VIRTUAL_ENV uv run --no-sync python scripts/demo_delta_recovery.py --run-nuts # + NUTS corner (~52 min) +env -u VIRTUAL_ENV uv run --no-sync python scripts/demo_delta_recovery_bias.py # wrong-IMF curve + grid +``` + +## References + +The multimass LIMEPY DF is {cite:t}`Gieles2015` with the equipartition +$m$-convention of {cite:t}`Peuten2017`; the $\sigma(m)$ equipartition relation and +its derived equipartition mass are {cite:t}`Bianchini2016`; the IMF is +{cite:t}`Maschberger2013`. The equipartition physics is developed on the +[multimass equipartition theory page](../10-theory/spatial-profiles/multimass-equipartition.md), +and the underlying equilibrium is validated at +[multimass equilibrium](../50-validation/multimass-equilibrium.md). diff --git a/docs/website/60-science-demos/index.md b/docs/website/60-science-demos/index.md new file mode 100644 index 00000000..445b86f2 --- /dev/null +++ b/docs/website/60-science-demos/index.md @@ -0,0 +1,142 @@ +--- +title: Science demos +description: "Three end-to-end methods demonstrations that turn progenax's differentiable initial conditions into inference: cross-engine agreement, self-consistent IMF+equipartition recovery, and halo+core anisotropy recovery — each a gated CLI with measured recovery tables." +--- + +# Science demos + +The [validation section](../50-validation/index.md) proves that each progenax +model reproduces its analytic or published ground truth. This section answers a +different question: **can you run the models *backwards* — recover the +parameters of a cluster from mock observations of it — and does the recovery +land on truth?** That is the use case the whole package is built for: a +*differentiable* forward model whose gradients drive maximum-likelihood and +Hamiltonian-Monte-Carlo inference. + +Three demos, each a standalone gated CLI in `scripts/` (the same `validate_*` +house style — it exits nonzero if any recovery gate fails): + +```{list-table} +:header-rows: 1 + +* - Demo + - Question + - Engine + - Recovered +* - [Cross-engine agreement](cross-engine.md) (B1) + - Do the DF-defined and density-defined engines build the *same* cluster? + - A vs B + - (none — a consistency check) +* - [IMF + equipartition recovery](imf-equipartition.md) (B2) + - Can one IMF slope $\alpha$ be measured jointly from masses *and* kinematics? + - A (LIMEPY multimass) + - $(\alpha,\ \delta,\ W_0)$ +* - [Halo + core recovery](halo-core.md) (B3) + - Can a two-family cluster's mass split + anisotropy be recovered? + - B (Eddington) + - $(t,\ r_a,\ r_h)$ +``` + +## The shared method: physics-direct differentiable inference + +All three demos share one inference layer (`scripts/_demo_inference.py`). The +design avoids both expensive likelihood-free / simulation-based inference and +any non-differentiable resampling inside the loss. The recipe: + +1. **Forward-sample mock stars at truth.** Draw $N$ positions+velocities from a + progenax model (here $N = 3\times10^4$ to $10^5$). This is done *once*, to + make the data; it is **not** repeated inside the optimizer. + +2. **Compress to binned kinematic summaries.** Reduce the stars to per-group (or + per-component) binned 1-D velocity dispersions $\hat\sigma_{1\mathrm D,j}(r)$ + on a frozen set of radial bins, each carrying a finite-$N$ standard error. + For an isotropic pooled estimator over the three velocity components, + + ```{math} + :label: sci-sigma-estimator + \hat\sigma_{1\mathrm D}^2 = \frac{1}{3}\,\overline{|\mathbf v|^2} + = \frac{\sigma_r^2 + \sigma_t^2}{3}, + \qquad + \mathrm{SE}(\hat\sigma) = \frac{\hat\sigma}{\sqrt{6 n}}, + ``` + + where the $\sqrt{6n}$ (not $\sqrt{2n}$) comes from pooling all *three* + one-D components into a $\chi^2_{3n}$ inner sum. **This compression is what + makes the cost $N$-independent:** the likelihood sees a few hundred binned + numbers, never the $10^5$ stars. + +3. **Likelihood against analytic predictions.** Compare the binned data to the + model's *analytic* expectation with a Gaussian (χ²) likelihood, + + ```{math} + :label: sci-loglike + \ln \mathcal L(\theta) = -\tfrac12 \sum_{j,k} w_{jk} + \left(\frac{\hat\sigma_{jk} - \sigma_{jk}^{\rm pred}(\theta)}{\mathrm{SE}_{jk}}\right)^2 , + ``` + + with $w_{jk}\in\{0,1\}$ masking under-occupied bins. **Gradients flow only + through the analytic prediction** $\sigma^{\rm pred}(\theta)$ — the data side + is a fixed constant — so $\nabla_\theta \ln\mathcal L$ is exact and cheap, and + there is no reparametrization-gradient noise from resampling. + +4. **Optimize and quantify.** Adam (`mle_adam`, a fixed-step `lax.scan` — + deterministic and differentiable-friendly) from several dispersed starts; + the Fisher information from the Gauss-Newton $J^{\mathsf T}J$ of the residual + vector (`fisher_information_gn` via `jax.jacrev`, which is robust through the + ODE/quadrature `custom_vjp`s, unlike a full Hessian); and, where affordable, + full posterior sampling with a vendored blackjax NUTS. + +```{important} +**The predictor is the unbiased *binned expectation*, not a bin-centre +evaluation.** Where the dispersion and number density vary steeply across a wide +outer bin, evaluating $\sigma(\theta)$ at the bin centre is biased. Instead each +demo predicts the number-weighted bin average + +​ $$\mathbb E[\hat\sigma_{jk}^2] = \frac{\int_{\rm bin} n_j(r)\,\sigma_j^2(r)\,dr}{\int_{\rm bin} n_j(r)\,dr},$$ + +with $n_j(r)$ the model's own number density — a like-with-like comparison to +the binned data, computed by a cumulative-trapezoid-at-edges trick that stays +differentiable. +``` + +## Honest scope — these are clean-mock methods showcases + +```{warning} +The demos demonstrate the **inference machinery on mock data drawn from the same +model family**, with deliberately idealized observations. They do **not** claim +recovery from real data. Specifically, none of the three model: + +- **line-of-sight projection** — the kinematic observables are the full 3-D + $\sigma_{1\mathrm D}(r)$ and $\beta(r)$, not projected/proper-motion quantities; +- **measurement errors or selection / incompleteness** — every star is observed + with its exact phase-space coordinates; +- **model misspecification of the family** — truth and fit share the same + generative model (the B2 *wrong-IMF* curve is the one deliberate exception, and + it is reported, not gated); +- **SBC-calibrated posteriors** — the NUTS corner is shown to be non-divergent + with posterior mean ≈ MLE, but simulation-based calibration of the credible + intervals is future work. + +Adding observational realism (projection, errors, incompleteness) is the natural +next step and is out of scope here. +``` + +## Running the demos + +The inference layer needs the `experimental` optional dependencies (blackjax for +NUTS, optax for Adam): + +```bash +env -u VIRTUAL_ENV uv pip install -e ".[dev,experimental]" + +# each demo is a gated CLI (exits nonzero on a recovery-gate failure) +python scripts/demo_cross_engine.py # B1 (seconds) +python scripts/demo_delta_recovery.py # B2 headline MLE + Fisher (minutes) +python scripts/demo_delta_recovery.py --run-nuts # + the ~52 min NUTS corner +python scripts/demo_delta_recovery_bias.py # B2 wrong-IMF curve + robustness grid +python scripts/demo_halo_core.py # B3 (~4 min, MLE-compile-dominated) +``` + +Figures are written to `validation/plots/` (a gitignored, regeneratable +artifact); the copies embedded on these pages live in +`docs/website/60-science-demos/figures/`. diff --git a/docs/website/myst.yml b/docs/website/myst.yml index aa3acceb..a68afb6f 100644 --- a/docs/website/myst.yml +++ b/docs/website/myst.yml @@ -20,7 +20,7 @@ project: email: alrosen@sdsu.edu github: https://github.com/drannarosen/progenax license: - code: MIT + code: Apache-2.0 content: CC-BY-4.0 open_access: true subject: Astrophysics @@ -174,6 +174,12 @@ project: - file: 50-validation/analytical-test-cases.md - file: 50-validation/physics-tests.md - file: 50-validation/plot-gallery.md + - title: Science demos + children: + - file: 60-science-demos/index.md + - file: 60-science-demos/cross-engine.md + - file: 60-science-demos/imf-equipartition.md + - file: 60-science-demos/halo-core.md - title: Development log children: - file: 90-development-log/index.md diff --git a/pyproject.toml b/pyproject.toml index 6c47617a..253716ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,19 @@ version = "0.1.0" description = "Differentiable initial conditions for N-body simulations" readme = "README.md" requires-python = ">=3.10" +license = "Apache-2.0" +license-files = ["LICENSE"] +authors = [{ name = "Anna Rosen", email = "alrosen@sdsu.edu" }] +keywords = ["astronomy", "n-body", "initial-conditions", "jax", "star-clusters", "differentiable"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Astronomy", +] dependencies = [ "jax>=0.4.20", "jaxlib>=0.4.20", @@ -17,33 +30,46 @@ dependencies = [ "jaxstro", ] +[project.urls] +Repository = "https://github.com/jaxstro/progenax" + [project.optional-dependencies] dev = ["pytest>=7.0", "pytest-cov", "pytest-xdist>=3.0"] +# Exact (non-differentiable) CW04 Q / Lambda_MSR diagnostics path. progenax core +# is JAX-native; numpy+scipy are needed ONLY by progenax.diagnostics' scipy-backed +# estimators (compute_q_parameter, compute_lambda_msr), which import them lazily +# and raise an actionable error if this extra is missing (audit R9). +diagnostics = ["numpy>=1.24", "scipy>=1.11"] # Repo-only experimental subsystem (src/experimental/gravoturb_fdf; NOT in the wheel). # Differentiable-inference deps. blackjax 1.x `window_adaptation`/`nuts` API per # inference/hmc.py (verified against 1.5); arviz for HMC convergence diagnostics # (R-hat/ESS/divergences/BFMI, inference/diagnostics.py); scipy for the SBC rank-histogram -# chi^2 uniformity test (validation oracle side); jaxstroviz for the reusable validation -# figure gallery + the integer-aware SBC chi^2 helpers (AC18). All CPU-only, validation side. +# chi^2 uniformity test (validation oracle side). All CPU-only, validation side. +# +# NOTE (audit R1/D2): jaxstroviz is deliberately NOT declared here. It is a PRIVATE, +# unpublished sibling, and pinning its editable path in the lock made `uv lock --check` +# / `uv sync --locked` fail in CI (and would break any non-sibling install) because the +# lock graph then requires a repo the default GITHUB_TOKEN cannot check out. The only +# code that needs it is the experimental SBC figure gallery + integer-aware chi^2 helpers +# (AC18), which already `pytest.importorskip("jaxstroviz")` and are xfail. To use the +# figure gallery locally, install it editable from the sibling checkout: +# uv pip install -e ../jaxstroviz +# A nightly/release lane may check it out with a PAT and install it the same way. experimental = [ "blackjax>=1.2,<2", "optax>=0.2", "arviz>=0.18", "scipy>=1.11", - "jaxstroviz", "flowjax>=17", # conditional normalizing flow for the SBI/NPE low-N beta likelihood (Path C) # "numpyro>=0.13", # reserved for the hierarchical population model (Decision #7) ] -# jaxstro and jaxstroviz are sibling ecosystem packages (not published to PyPI). Resolve -# them from the standard side-by-side checkout so `uv lock` / `uv sync --locked` work. -# CI checks the siblings out as adjacent directories (see .github/workflows/tests.yml). -# jaxstroviz is the shared viz/validation package: progenax's validation scripts + figures -# live here and call it (like jaxstro for core utilities); it stays validation-side only -# (released progenax core never imports it; the experimental figure orchestrator import-guards). +# jaxstro is the only sibling ecosystem package in the lock graph. It is PUBLIC +# (github.com/jaxstro/jaxstro), so CI checks it out as an adjacent directory and +# `uv lock` / `uv sync --locked` resolve with no private-repo credentials. +# (jaxstroviz is intentionally absent — see the [experimental] note above.) [tool.uv.sources] jaxstro = { path = "../jaxstro", editable = true } -jaxstroviz = { path = "../jaxstroviz", editable = true } [tool.hatch.build.targets.wheel] packages = ["src/progenax"] diff --git a/scripts/demo_delta_recovery.py b/scripts/demo_delta_recovery.py index bdb55007..5a04fb5f 100644 --- a/scripts/demo_delta_recovery.py +++ b/scripts/demo_delta_recovery.py @@ -165,11 +165,6 @@ class bounds equal the draw bounds, so the normalization matches the sample). _TINY = 1e-30 -def _truth_imf(): - """The truth Maschberger over M_RANGE (class bounds == draw bounds).""" - return Maschberger(alpha=ALPHA_TRUE, m_min=M_RANGE[0], m_max=M_RANGE[1]) - - # --------------------------------------------------------------------------- # # g(W) = I4(W)/I2(W)/3 table (Engine A 1-D moment ratio; built once per eval) # --------------------------------------------------------------------------- # @@ -299,15 +294,22 @@ def _per_group(j): # --------------------------------------------------------------------------- # # Mock data construction (run once; R_CUT, R_EDGES then frozen) # --------------------------------------------------------------------------- # -def build_truth_data(): +def build_truth_data(key=None, alpha_true=ALPHA_TRUE): """Sample the truth cluster, truncate at R_CUT, bin sigma_1d,j(r), draw masses. Returns a dict of constants the loss closes over: R_EDGES (16 quantile bins on [0, R_CUT]), sig_hat, se, weight, n (kinematics), m_obs (Option A mass channel), M_FIXED, R_CUT. + + ``key`` (default ``PRNGKey(0)``) and ``alpha_true`` (default the module + truth) are exposed for the Task 6 ensembles (independent seed datasets and + the alpha_true robustness grid); the defaults reproduce the Task 3-5 + headline dataset bit-for-bit. delta / W0 truths stay at the module + constants in all cases. """ - imf = _truth_imf() - key = jax.random.PRNGKey(0) + imf = Maschberger(alpha=alpha_true, m_min=M_RANGE[0], m_max=M_RANGE[1]) + if key is None: + key = jax.random.PRNGKey(0) k_kin, k_mass = jax.random.split(key) model_true = MultiComponentCluster.from_imf( @@ -506,9 +508,14 @@ def run_mle(negloglike, data): return z_hats[i_best], traces[i_best], finals, i_best, z0s -def recovery_table(theta_hat, sigma_theta): - """Print param | truth | theta_hat | sigma_hat | (theta_hat-truth)/sigma_hat.""" - truths = (ALPHA_TRUE, DELTA_TRUE, W0_TRUE) +def recovery_table(theta_hat, sigma_theta, truths=None): + """Print param | truth | theta_hat | sigma_hat | (theta_hat-truth)/sigma_hat. + + ``truths`` defaults to the module truth ``(ALPHA_TRUE, DELTA_TRUE, W0_TRUE)``; + the Task 6 robustness grid passes its per-dataset ``alpha_true``. + """ + if truths is None: + truths = (ALPHA_TRUE, DELTA_TRUE, W0_TRUE) names = ("alpha", "delta", "W0") print(f"\n{'param':>6} {'truth':>8} {'theta_hat':>10} {'sigma_hat':>10} " f"{'(hat-truth)/sigma':>18}") diff --git a/scripts/demo_delta_recovery_bias.py b/scripts/demo_delta_recovery_bias.py new file mode 100644 index 00000000..de1f5038 --- /dev/null +++ b/scripts/demo_delta_recovery_bias.py @@ -0,0 +1,672 @@ +r"""B2d science demo (Task 6): wrong-IMF bias curve + alpha_true robustness grid. + +Sibling of ``demo_delta_recovery.py`` (Tasks 3-5): that script's default run is +the budgeted single-dataset MLE headline, so the Task 6 ENSEMBLES live in this +separate gated CLI (the plan's "sibling script" branch). It reuses the B2 +machinery verbatim -- truth constants, ``build_truth_data``, ``predict_binned``, +the expit boxes, and the ``_demo_inference`` Adam/Fisher helpers -- nothing is +re-derived here. + +Two ensembles +------------- +1. **Wrong-IMF bias curve** (REPORTED, not gated). For each + ``alpha_assumed in {1.9, 2.1, 2.3, 2.5, 2.7}``: a KINEMATICS-ONLY refit of + ``(delta, W0)`` -- the mass channel DROPPED, alpha FROZEN at + ``alpha_assumed`` -- on mock data generated at the truth + ``(alpha=2.3, delta=0.4, W0=5)``. ``N_SEEDS`` independent truth datasets + (different sampling keys, SAME dataset reused across the 5 alpha_assumed + within a seed, so the per-seed slope is a paired measurement). The quoted + slope ``d delta_hat / d alpha_assumed`` is a per-seed linear fit, reported + as mean +/- SE over the seed ensemble. It is REPORTED rather than gated + because it is a sensitivity MEASUREMENT of this mock configuration -- there + is no published reference value for the equipartition-vs-IMF-slope bias of + a 4-component LIMEPY fit to assert against; the demo's claim is the number + itself, with an honest seed-scatter uncertainty. +2. **Robustness grid** (GATED). ``alpha_true in {1.9, 2.3, 2.7}`` with a FRESH + truth dataset at each ``alpha_true`` (same ``delta=0.4, W0=5``), then the + FULL joint ``(alpha, delta, W0)`` refit WITH the mass term (3 dispersed + inits, exactly the Task 4 recipe). Gate: each recovery within 3 sigma_hat + componentwise (Gauss-Newton Fisher errors via ``fisher_information_gn`` + + ``constrained_cov``). The script exits nonzero if any component misses. + +Runtime budget (STOP checkpoint, executed BEFORE the ensembles run) +------------------------------------------------------------------- +``N_SEEDS`` is scaled to the MEASURED warm eval cost so the whole bias panel +stays under ``BIAS_BUDGET_MIN = 30`` min. The script measures (i) one truth +dataset build and (ii) one warm ``jit(value_and_grad)`` of the 2-parameter +kinematics-only loss, then projects + + T_bias(N_SEEDS) ~ N_SEEDS * [ t_data + N_ALPHA * N_ADAM_STEPS * t_warm ] + +and picks ``N_SEEDS = clamp(floor(budget / per-seed), 3, 5)``. The arithmetic +is PRINTED before any ensemble fit runs; if even 3 seeds do not fit the +budget the script STOPs (exit 2) and reports options instead of degrading. + +Captured headline run (2026-06-11, CPU/float64; authoritative run record +/tmp/b2d_run.log): t_data = 6.1 s; one warm 2-param kin-only FIT (300 fused +Adam steps) = 64.3 s (the projection uses the measured warm-FIT cost, not +300 x a lone warm value_and_grad, which understates the per-step cost). +Per-seed ~5.46 min -> N_SEEDS = 5, PROJECTED 27.3 min < 30 min budget; +bias-ensemble ACTUAL wall 30.3 min (post-hoc machine-load variance -- the +projection BEFORE the run is the budget gate, and it passed). Robustness +grid: measured 13.6 min. NOTE that ``N_SEEDS`` is LOAD-ADAPTIVE: the budget +arithmetic picks 3-5 seeds from the warm-fit cost measured at run time, so +a rerun on a loaded machine may legitimately produce a 4-seed figure (the +captured headline run measured 64.3 s -> 5 seeds). + +Peaked-shape caveat (captured run): the measured ensemble-mean +``delta_hat(alpha_assumed)`` = (0.138, 0.393, 0.420, 0.285, 0.187) is +NON-MONOTONIC and peaked AT the truth alpha=2.3, so the linear slope is not +an adequate summary by itself -- wrong alpha in EITHER direction biases +delta_hat LOW (by ~0.2-0.3 at |Delta alpha| = 0.4). The script therefore +also prints the peak-to-end drops alongside the slope and annotates them on +figure panel (a). + +CLI ``--diag-21`` (valley diagnostic; runs NO ensemble and NO gate) +------------------------------------------------------------------- +The captured bias curve has its largest seed scatter at alpha_assumed=2.1 +(seed std 0.084 vs <=0.016 elsewhere; seeds 0 and 1 split delta_hat ~0.47 +vs ~0.27). ``--diag-21`` probes whether that split is an optimizer artifact +or a real likelihood feature: (i) it profiles the kinematics-only +negloglike along delta (25 points over ``DELTA_BOX``) at alpha_assumed=2.1 +with W0 FIXED at each seed's fitted value, for seeds 0 and 1, printing the +profile values so the valley shape (single-well vs flat/double-well) is in +the output; (ii) it refits alpha_assumed=2.1 for all 5 headline seed +datasets from 3 dispersed inits (``b2.dispersed_inits`` adapted to 2 +params) with 600 Adam steps, printing per-seed per-init +(delta_hat, W0_hat, final loss). Interpretation rule: same-dataset +init-dependence => optimizer artifact; init-independence with +seed-dependence => genuinely flat/bistable likelihood direction (a +reportable finding). The headline ensemble numbers and gates are untouched. + +Bias-curve refit design notes +----------------------------- +* alpha is frozen by building ``theta = (alpha_assumed, expit(z[0], DELTA_BOX), + expit(z[1], W0_BOX))`` inside the loss -- the optimizer sees only the + 2-vector ``z``; no gradient flows to alpha. +* Single init ``z = 0`` (delta=0.35, W0=5.5): Task 4 showed the joint 3-param + landscape is init-insensitive (3 dispersed inits, identical optimum); the + kinematics-only 2-param problem is a slice of it. Each fit's Adam trace is + plateau-checked (reported; a non-plateau run is flagged in the table). +* The fit driver is ONE jitted function over (alpha_assumed, data arrays), so + XLA compiles once and all 5 x N_SEEDS fits reuse the executable; the + robustness grid likewise compiles its joint driver once for 3 x 3 fits. +""" + +import os +import sys +import time + +import jax +import jax.numpy as jnp + +import progenax # noqa: F401 -- enables float64 at import +from progenax.imf.smooth import Maschberger + +sys.path.insert(0, os.path.dirname(__file__)) +import _demo_inference as di # noqa: E402 +import demo_delta_recovery as b2 # noqa: E402 (Tasks 3-5 machinery, reused) + +# --------------------------------------------------------------------------- # +# Task 6 configuration +# --------------------------------------------------------------------------- # +ALPHAS_ASSUMED = (1.9, 2.1, 2.3, 2.5, 2.7) # bias curve (truth alpha = 2.3) +ALPHAS_TRUE_GRID = (1.9, 2.3, 2.7) # robustness grid truths +SEED_BASE_BIAS = 100 # PRNGKey(100 + s), s = 0..N_SEEDS-1 +SEED_BASE_GRID = 200 # PRNGKey(200 + i) per alpha_true +N_SEEDS_MIN, N_SEEDS_MAX = 3, 5 +BIAS_BUDGET_MIN = 30.0 # whole-bias-panel budget (STOP gate) + + +# --------------------------------------------------------------------------- # +# Kinematics-only (delta, W0) refit with alpha FROZEN (bias curve) +# --------------------------------------------------------------------------- # +def _kin_negloglike_2p(z2, alpha_assumed, sig_hat, se, weight, r_edges, m_fixed): + """Kinematics-only negloglike in z2 = (z_delta, z_W0); alpha frozen, mass + term DROPPED. Same standardized chi^2 as the joint loss's kinematic part.""" + theta = (alpha_assumed, + di.expit(z2[0], *b2.DELTA_BOX), + di.expit(z2[1], *b2.W0_BOX)) + sig_model = b2.predict_binned(theta, r_edges, m_fixed) + safe_se = jnp.where(se > 0, se, 1.0) + resid = jnp.sqrt(weight) * (sig_hat - sig_model) / safe_se + return 0.5 * jnp.sum(resid * resid) + + +def _fit_kin_only(alpha_assumed, sig_hat, se, weight, r_edges, m_fixed): + """One kinematics-only (delta, W0) Adam fit (Task 4 step count / lr). + Jitted ONCE at module scope (below); all 5 x N_SEEDS fits reuse it.""" + def negll(z2): + return _kin_negloglike_2p(z2, alpha_assumed, sig_hat, se, weight, + r_edges, m_fixed) + + return di.mle_adam(negll, jnp.zeros(2), + n_steps=b2.N_ADAM_STEPS, lr=b2.ADAM_LR) + + +fit_kin_only = jax.jit(_fit_kin_only) +kin_negll_2p = jax.jit(_kin_negloglike_2p) + + +# --------------------------------------------------------------------------- # +# --diag-21: valley diagnostic at alpha_assumed = 2.1 (see module docstring) +# --------------------------------------------------------------------------- # +DIAG_ALPHA = 2.1 +DIAG_N_DELTA = 25 +DIAG_PROFILE_SEEDS = (0, 1) # the seeds that split: delta_hat ~0.47 vs ~0.27 +DIAG_N_SEEDS = 5 # the captured headline run's seeds (keys 100-104) +DIAG_N_STEPS = 600 # 2x the ensemble's 300 (rules out slow convergence) + + +def dispersed_inits_2p(): + """``b2.dispersed_inits`` adapted to the 2-param (delta, W0) refit: + z0 = 0 plus ``N_INITS - 1`` draws from N(0, INIT_SCALE^2 I_2), same + INIT_KEY / INIT_SCALE as the joint-refit recipe.""" + key = jax.random.PRNGKey(b2.INIT_KEY) + draws = jax.random.normal(key, (b2.N_INITS - 1, 2)) * b2.INIT_SCALE + return jnp.concatenate([jnp.zeros((1, 2)), draws], axis=0) + + +def _kin_negloglike_theta(delta, w0, alpha_assumed, sig_hat, se, weight, + r_edges, m_fixed): + """The SAME kinematics-only negloglike as ``_kin_negloglike_2p`` but in + PHYSICAL (delta, W0) -- used to profile the likelihood on a delta grid.""" + sig_model = b2.predict_binned((alpha_assumed, delta, w0), r_edges, m_fixed) + safe_se = jnp.where(se > 0, se, 1.0) + resid = jnp.sqrt(weight) * (sig_hat - sig_model) / safe_se + return 0.5 * jnp.sum(resid * resid) + + +profile_negll = jax.jit(jax.vmap( + _kin_negloglike_theta, + in_axes=(0, None, None, None, None, None, None, None))) + + +def _fit_kin_only_diag(z0, alpha_assumed, sig_hat, se, weight, r_edges, + m_fixed): + """Kinematics-only fit from an EXPLICIT init with DIAG_N_STEPS Adam steps + (the ensemble fit is z0=0 with N_ADAM_STEPS).""" + def negll(z2): + return _kin_negloglike_2p(z2, alpha_assumed, sig_hat, se, weight, + r_edges, m_fixed) + + return di.mle_adam(negll, z0, n_steps=DIAG_N_STEPS, lr=b2.ADAM_LR) + + +fit_kin_only_diag = jax.jit(_fit_kin_only_diag) + + +def diag_alpha21(): + """--diag-21 driver: (i) delta profile at fixed W0_hat for seeds 0 and 1; + (ii) 3-dispersed-init x 600-step refits for all 5 headline seed datasets. + Runs NO ensemble and NO gate; headline numbers untouched.""" + t_all0 = time.perf_counter() + print("=" * 72) + print(f"DIAG-21: alpha_assumed = {DIAG_ALPHA} valley diagnostic " + "(kinematics-only refit)") + print("=" * 72) + print(f"context: the captured bias curve has its largest seed scatter at " + f"alpha_assumed={DIAG_ALPHA} (seed std 0.084 vs <=0.016 elsewhere; " + f"seeds 0/3 delta_hat~0.47 vs seeds 1/2/4 ~0.27-0.37).") + print("question: optimizer artifact or genuine likelihood feature?") + + datasets = {} + + def get_data(s): + if s not in datasets: + datasets[s] = b2.build_truth_data( + key=jax.random.PRNGKey(SEED_BASE_BIAS + s)) + return datasets[s] + + def args_of(data): + return (data["sig_hat"], data["se"], data["weight"], + data["r_edges"], data["M_fixed"]) + + # ------------------------------------------------------------------ # + # (i) negloglike profile along delta, W0 FIXED at each seed's fit + # ------------------------------------------------------------------ # + print("\n" + "-" * 72) + print(f"(i) negloglike profile along delta ({DIAG_N_DELTA} points over " + f"DELTA_BOX={b2.DELTA_BOX}), W0 FIXED at each seed's fitted value") + print("-" * 72) + deltas = jnp.linspace(b2.DELTA_BOX[0], b2.DELTA_BOX[1], DIAG_N_DELTA) + for s in DIAG_PROFILE_SEEDS: + data = get_data(s) + args = args_of(data) + # The ensemble fit itself (300 steps, z0 = 0) -> this seed's W0_hat. + z2_hat, _ = fit_kin_only(jnp.asarray(DIAG_ALPHA), *args) + d_hat = float(di.expit(z2_hat[0], *b2.DELTA_BOX)) + w_hat = float(di.expit(z2_hat[1], *b2.W0_BOX)) + vals = profile_negll(deltas, jnp.asarray(w_hat), + jnp.asarray(DIAG_ALPHA), *args) + v = [float(x) for x in vals] + print(f"\n seed {s} (key {SEED_BASE_BIAS + s}): ensemble fit " + f"delta_hat={d_hat:.4f}, W0_hat={w_hat:.4f} (W0 frozen below)") + print(f" {'delta':>8} {'negloglike':>14}") + for k in range(DIAG_N_DELTA): + mark = "" + if 0 < k < DIAG_N_DELTA - 1 and v[k] < v[k - 1] and v[k] < v[k + 1]: + mark = " <- local min" + print(f" {float(deltas[k]):>8.4f} {v[k]:>14.4f}{mark}") + i_min = int(jnp.argmin(vals)) + print(f" grid minimum at delta = {float(deltas[i_min]):.4f} " + f"(negloglike {v[i_min]:.4f}); profile span max-min = " + f"{max(v) - min(v):.4f}") + + # ------------------------------------------------------------------ # + # (ii) dispersed-init refits, all headline seeds + # ------------------------------------------------------------------ # + print("\n" + "-" * 72) + print(f"(ii) refit alpha_assumed={DIAG_ALPHA} for all {DIAG_N_SEEDS} seed " + f"datasets: {b2.N_INITS} dispersed inits x {DIAG_N_STEPS} Adam steps") + print("-" * 72) + z0s = dispersed_inits_2p() + print(f" inits (unconstrained z = (z_delta, z_W0)): " + f"{[[round(float(x), 3) for x in z0] for z0 in z0s]}") + spreads = [] + best_d = [] + for s in range(DIAG_N_SEEDS): + data = get_data(s) + args = args_of(data) + d_hats, finals = [], [] + print(f"\n seed {s} (key {SEED_BASE_BIAS + s}):") + for i in range(b2.N_INITS): + z2_hat, trace = fit_kin_only_diag(z0s[i], jnp.asarray(DIAG_ALPHA), + *args) + d = float(di.expit(z2_hat[0], *b2.DELTA_BOX)) + w = float(di.expit(z2_hat[1], *b2.W0_BOX)) + final = float(kin_negll_2p(z2_hat, jnp.asarray(DIAG_ALPHA), *args)) + plat, _, _ = b2.plateau_ok(trace) + d_hats.append(d) + finals.append(final) + print(f" init {i}: delta_hat={d:.4f}, W0_hat={w:.4f}, " + f"final loss={final:.4f}, plateau " + f"{'PASS' if plat else 'FAIL'}") + spread = max(d_hats) - min(d_hats) + spreads.append(spread) + best_d.append(d_hats[int(jnp.argmin(jnp.array(finals)))]) + print(f" across-init delta_hat spread = {spread:.4f} " + f"(best-loss init's delta_hat = {best_d[-1]:.4f})") + + print(f"\n across-INIT spread per seed (max) = {max(spreads):.4f}") + print(f" across-SEED spread of best-loss delta_hat = " + f"{max(best_d) - min(best_d):.4f}") + print("\n" + "=" * 72) + print("INTERPRETATION RULE") + print("=" * 72) + print(" same-dataset INIT-dependence of delta_hat => optimizer artifact") + print(" init-INDEPENDENCE with seed-dependence => genuinely flat / " + "bistable likelihood direction at alpha_assumed=2.1 " + "(a reportable finding, not an optimizer bug)") + print(f"\n diag wall-time = {(time.perf_counter() - t_all0) / 60.0:.1f} " + "min") + print(" (headline ensemble numbers and gates untouched -- this mode runs " + "NO ensemble and NO gate)") + + +# --------------------------------------------------------------------------- # +# Full joint (alpha, delta, W0) refit (robustness grid) -- Task 4 recipe +# --------------------------------------------------------------------------- # +def _joint_negloglike(z, sig_hat, se, weight, r_edges, m_fixed, m_obs): + """The Task 3 joint negloglike (kinematics + Option A mass channel), with + the data arrays as arguments so the jitted fit compiles once.""" + theta = b2._theta_of_z(z) + sig_model = b2.predict_binned(theta, r_edges, m_fixed) + safe_se = jnp.where(se > 0, se, 1.0) + resid = jnp.sqrt(weight) * (sig_hat - sig_model) / safe_se + alpha = theta[0] + mass_nll = -jnp.sum( + Maschberger(alpha=alpha, m_min=b2.M_RANGE[0], + m_max=b2.M_RANGE[1]).logpdf(m_obs)) + return 0.5 * jnp.sum(resid * resid) + mass_nll + + +def _fit_joint(z0, sig_hat, se, weight, r_edges, m_fixed, m_obs): + def negll(z): + return _joint_negloglike(z, sig_hat, se, weight, r_edges, m_fixed, m_obs) + + return di.mle_adam(negll, z0, n_steps=b2.N_ADAM_STEPS, lr=b2.ADAM_LR) + + +fit_joint = jax.jit(_fit_joint) + + +def joint_refit(data): + """Full joint refit on one dataset: 3 dispersed inits -> best z_hat, + theta_hat, Gauss-Newton Fisher sigma_theta, plateau flag. Task 4 recipe via + the b2 builders (make_residual_fn / make_mass_negloglike / fisher_*).""" + args = (data["sig_hat"], data["se"], data["weight"], + data["r_edges"], data["M_fixed"], data["m_obs"]) + z0s = b2.dispersed_inits() + z_hats, traces, finals = [], [], [] + for i in range(b2.N_INITS): + z_hat, trace = fit_joint(z0s[i], *args) + z_hats.append(z_hat) + traces.append(trace) + finals.append(float(_joint_negloglike(z_hat, *args))) + i_best = int(jnp.argmin(jnp.array(finals))) + z_hat, trace = z_hats[i_best], traces[i_best] + plat_ok, _, _ = b2.plateau_ok(trace) + + residual_fn = b2.make_residual_fn(data) + mass_negloglike = b2.make_mass_negloglike(data) + F_z = di.fisher_information_gn(residual_fn, z_hat, + extra_negloglike=mass_negloglike) + cov_theta = di.constrained_cov(F_z, b2._dtheta_dz(z_hat)) # raises if not PD + sigma_theta = jnp.sqrt(jnp.diag(cov_theta)) + theta_hat = b2._theta_of_z(z_hat) + return theta_hat, sigma_theta, plat_ok, finals, i_best + + +# --------------------------------------------------------------------------- # +# Figure: (a) bias curve (b) robustness-grid pulls +# --------------------------------------------------------------------------- # +def make_bias_figure(bias, grid, out_dir): + """Two-panel Task 6 figure via _plotstyle. + + Panel (a): delta_hat(alpha_assumed) per seed (light points) + ensemble + mean +/- seed std, truth lines, quoted slope. Panel (b): robustness-grid + pulls (theta_hat - truth)/sigma_hat per parameter vs alpha_true, with the + +/- 3 sigma gate band. + """ + import matplotlib.pyplot as plt + import numpy as np + + import _plotstyle as ps # noqa: E402 (path already inserted) + + ps.apply_pub_style() + fig, (axa, axb) = plt.subplots(1, 2, figsize=(8.5, 3.6)) + + # ---- panel (a): bias curve --------------------------------------------- # + alphas = np.asarray(bias["alphas"]) # (A,) + d_hat = np.asarray(bias["delta_hat"]) # (S, A) + d_mean = d_hat.mean(axis=0) + d_std = d_hat.std(axis=0, ddof=1) + for s in range(d_hat.shape[0]): + axa.plot(alphas, d_hat[s], "o", ms=3, color=ps.OI["sky"], alpha=0.6, + mec="none", zorder=2, + label="per-seed fits" if s == 0 else None) + axa.errorbar(alphas, d_mean, yerr=d_std, fmt="s-", ms=4.5, + color=ps.OI["blue"], capsize=2.0, elinewidth=1.0, lw=1.4, + zorder=3, label=r"seed mean $\pm$ std") + axa.axhline(b2.DELTA_TRUE, color=ps.OI["orange"], lw=1.0, ls=":", + label=rf"truth $\delta={b2.DELTA_TRUE}$") + axa.axvline(b2.ALPHA_TRUE, color=ps.OI["orange"], lw=1.0, ls=":") + axa.set_xlabel(r"$\alpha_{\rm assumed}$ (frozen IMF slope)") + axa.set_ylabel(r"$\hat\delta$ (kinematics-only refit)") + axa.legend(loc="best") + axa.text(0.97, 0.05, + rf"$\mathrm{{d}}\hat\delta/\mathrm{{d}}\alpha = " + rf"{bias['slope_mean']:+.4f} \pm {bias['slope_se']:.4f}$", + transform=axa.transAxes, ha="right", va="bottom", fontsize=9) + axa.text(0.97, 0.13, + rf"peaked at truth: wrong $\alpha$ either way biases " + rf"$\hat\delta$ low" "\n" + rf"($-{bias['drop_lo']:.2f}$ at $\alpha={alphas[0]:.1f}$, " + rf"$-{bias['drop_hi']:.2f}$ at $\alpha={alphas[-1]:.1f}$)", + transform=axa.transAxes, ha="right", va="bottom", fontsize=7.5) + ps.panel_label(axa, "(a)") + + # ---- panel (b): robustness-grid pulls ---------------------------------- # + a_true = np.asarray(grid["alphas_true"]) # (G,) + pulls = np.asarray(grid["pulls"]) # (G, 3) + names = (r"$\alpha$", r"$\delta$", r"$W_0$") + colors = (ps.OI["blue"], ps.OI["green"], ps.OI["vermilion"]) + markers = ("o", "s", "^") + off = (-0.02, 0.0, 0.02) + axb.axhspan(-3.0, 3.0, color=ps.OI["sky"], alpha=0.15, zorder=0) + axb.axhline(0.0, color=ps.OI["black"], lw=0.7, ls="-", zorder=1) + for p in range(3): + axb.plot(a_true + off[p], pulls[:, p], markers[p], ms=5, + color=colors[p], mec=ps.OI["black"], mew=0.4, zorder=3, + label=names[p]) + for y in (-3.0, 3.0): + axb.axhline(y, color=ps.OI["vermilion"], lw=0.8, ls="--", zorder=1) + axb.set_xlabel(r"$\alpha_{\rm true}$ (fresh truth dataset)") + axb.set_ylabel(r"pull $(\hat\theta - \theta_{\rm true})/\hat\sigma$") + axb.set_xticks(a_true) + axb.set_ylim(-4.0, 4.0) + axb.legend(loc="best", ncol=3) + ps.panel_label(axb, "(b)") + + cap = (rf"(a) wrong-IMF bias of the kinematics-only $(\delta, W_0)$ refit, " + rf"{d_hat.shape[0]} seeds; (b) joint refit pulls vs $\alpha_{{\rm true}}$ " + rf"($\pm3\hat\sigma$ gate band)") + fig.text(0.5, -0.02, cap, ha="center", va="top", fontsize=8.5) + fig.tight_layout() + ps.save_fig(fig, out_dir, "demo_delta_recovery_bias") + print(f"\nbias figure -> {out_dir}/demo_delta_recovery_bias.png (+ .pdf)") + + +# --------------------------------------------------------------------------- # +# Driver +# --------------------------------------------------------------------------- # +def main(): + print("=" * 72) + print("B2d demo (Task 6): wrong-IMF bias curve + alpha_true robustness grid") + print("=" * 72) + print(f"truth: alpha={b2.ALPHA_TRUE}, delta={b2.DELTA_TRUE}, W0={b2.W0_TRUE}; " + f"N_COMP={b2.N_COMP}, N_STARS={b2.N_STARS}, M_RANGE={b2.M_RANGE}") + print(f"bias curve: alpha_assumed in {ALPHAS_ASSUMED} (kinematics-only, " + f"alpha frozen, mass term dropped)") + print(f"robustness grid: alpha_true in {ALPHAS_TRUE_GRID} (fresh datasets, " + f"full joint refit; GATED at 3 sigma)") + + # ------------------------------------------------------------------ # + # RUNTIME BUDGET CHECKPOINT (before any ensemble fit runs) + # ------------------------------------------------------------------ # + print("\n" + "=" * 72) + print(f"RUNTIME BUDGET CHECKPOINT (bias panel < {BIAS_BUDGET_MIN:.0f} min)") + print("=" * 72) + t0 = time.perf_counter() + data0 = b2.build_truth_data(key=jax.random.PRNGKey(SEED_BASE_BIAS)) + t_data = time.perf_counter() - t0 + print(f" one truth-dataset build (sample + bin) t_data = {t_data:.1f} s") + + args0 = (data0["sig_hat"], data0["se"], data0["weight"], + data0["r_edges"], data0["M_fixed"]) + # Probe with TWO REAL fits (reused below as the first two seed-0 ensemble + # members, so no probe time is wasted). The projection uses the measured + # warm FIT cost rather than N_ADAM_STEPS x (one warm eval): measured here, + # the fused Adam scan costs slightly MORE per step than a lone + # value_and_grad (0.28 vs 0.23 s on this machine), so the per-eval + # arithmetic would UNDERSTATE the panel. + t0 = time.perf_counter() + z2_p0, trace_p0 = fit_kin_only(jnp.asarray(ALPHAS_ASSUMED[0]), *args0) + z2_p0.block_until_ready() + t_fit_cold = time.perf_counter() - t0 + t0 = time.perf_counter() + z2_p1, trace_p1 = fit_kin_only(jnp.asarray(ALPHAS_ASSUMED[1]), *args0) + z2_p1.block_until_ready() + t_fit_warm = time.perf_counter() - t0 + probe_fits = {(0, 0): (z2_p0, trace_p0), (0, 1): (z2_p1, trace_p1)} + print(f" 2-param kin-only fit ({b2.N_ADAM_STEPS} Adam steps): " + f"cold (incl. compile) {t_fit_cold:.1f} s, warm {t_fit_warm:.1f} s") + + n_alpha = len(ALPHAS_ASSUMED) + per_seed = t_data + n_alpha * t_fit_warm + print(f" per-seed ~ t_data + {n_alpha} x warm-fit = {t_data:.1f} + " + f"{n_alpha} x {t_fit_warm:.1f} = {per_seed:.1f} s " + f"= {per_seed / 60.0:.2f} min") + n_seeds_fit = int((BIAS_BUDGET_MIN * 60.0) // per_seed) + n_seeds = max(N_SEEDS_MIN, min(N_SEEDS_MAX, n_seeds_fit)) + proj_min = n_seeds * per_seed / 60.0 + print(f" budget arithmetic: floor({BIAS_BUDGET_MIN:.0f} min / " + f"{per_seed / 60.0:.2f} min) = {n_seeds_fit} seeds -> " + f"N_SEEDS = clamp(.., {N_SEEDS_MIN}, {N_SEEDS_MAX}) = {n_seeds}") + print(f" PROJECTED bias-panel wall-time = {n_seeds} x {per_seed / 60.0:.2f} " + f"= {proj_min:.1f} min (budget {BIAS_BUDGET_MIN:.0f} min; " + f"data build + 2 probe fits already banked)") + if n_seeds_fit < N_SEEDS_MIN: + print(f"\n STOP: even {N_SEEDS_MIN} seeds project to " + f"{N_SEEDS_MIN * per_seed / 60.0:.1f} min > " + f"{BIAS_BUDGET_MIN:.0f} min. NOT degrading the solve. Options: " + "reduce N_ADAM_STEPS for the 2-param refit (with a plateau " + "re-check), reduce the alpha_assumed grid to 3 points, or raise " + "the budget with Anna's ok.") + sys.exit(2) + # Informational projection for the (un-budgeted) robustness grid: the + # joint loss adds the (cheap) mass channel + a 3rd parameter; ~1.3x the + # kin-only fit cost is a generous per-fit allowance. + grid_proj = (len(ALPHAS_TRUE_GRID) + * (t_data + b2.N_INITS * 1.3 * t_fit_warm)) / 60.0 + print(f" robustness-grid projection (informational, ~1.3x kin fit cost " + f"x {b2.N_INITS} inits x {len(ALPHAS_TRUE_GRID)} datasets): " + f"~{grid_proj:.0f} min + one compile + 3 Fisher evals") + + # ------------------------------------------------------------------ # + # Ensemble 1: wrong-IMF bias curve (REPORTED, not gated) + # ------------------------------------------------------------------ # + print("\n" + "=" * 72) + print(f"BIAS CURVE: kinematics-only (delta, W0) refit, {n_seeds} seeds") + print("=" * 72) + t_bias0 = time.perf_counter() + delta_hat = [] # (S, A) + w0_hat = [] # (S, A) + n_no_plateau = 0 + for s in range(n_seeds): + data = data0 if s == 0 else b2.build_truth_data( + key=jax.random.PRNGKey(SEED_BASE_BIAS + s)) + args = (data["sig_hat"], data["se"], data["weight"], + data["r_edges"], data["M_fixed"]) + d_row, w_row = [], [] + for ia, a in enumerate(ALPHAS_ASSUMED): + if (s, ia) in probe_fits: # banked during the budget probe + z2_hat, trace = probe_fits[(s, ia)] + else: + z2_hat, trace = fit_kin_only(jnp.asarray(a), *args) + plat, _, _ = b2.plateau_ok(trace) + if not plat: + n_no_plateau += 1 + d_row.append(float(di.expit(z2_hat[0], *b2.DELTA_BOX))) + w_row.append(float(di.expit(z2_hat[1], *b2.W0_BOX))) + print(f" seed {s} (key {SEED_BASE_BIAS + s}) " + f"alpha_assumed={a:.1f}: delta_hat={d_row[-1]:.4f}, " + f"W0_hat={w_row[-1]:.4f}" + + ("" if plat else " [NO PLATEAU]")) + delta_hat.append(d_row) + w0_hat.append(w_row) + t_bias = time.perf_counter() - t_bias0 + delta_hat = jnp.array(delta_hat) + w0_hat = jnp.array(w0_hat) + + alphas = jnp.asarray(ALPHAS_ASSUMED) + print(f"\n{'alpha_assumed':>13} {'delta_hat mean':>14} {'seed std':>9} " + f"{'W0_hat mean':>11} {'seed std':>9}") + for i, a in enumerate(ALPHAS_ASSUMED): + print(f"{a:>13.1f} {float(delta_hat[:, i].mean()):>14.4f} " + f"{float(delta_hat[:, i].std(ddof=1)):>9.4f} " + f"{float(w0_hat[:, i].mean()):>11.4f} " + f"{float(w0_hat[:, i].std(ddof=1)):>9.4f}") + + # Per-seed linear slope d delta_hat / d alpha (paired across alpha_assumed). + a_c = alphas - alphas.mean() + slopes = (delta_hat @ a_c) / jnp.sum(a_c * a_c) # (S,) LSQ slopes + slope_mean = float(slopes.mean()) + slope_std = float(slopes.std(ddof=1)) + slope_se = slope_std / float(jnp.sqrt(n_seeds)) + print(f"\n per-seed slopes d delta_hat/d alpha = " + f"{[round(float(x), 4) for x in slopes]}") + print(f" SLOPE d delta_hat/d alpha = {slope_mean:+.4f} +/- {slope_se:.4f} " + f"(SE of mean over {n_seeds} seeds; seed scatter {slope_std:.4f})") + print(" [REPORTED, not gated: a sensitivity measurement of this mock " + "configuration -- no published reference value to assert against.]") + # Peaked-shape summary: the response delta_hat(alpha_assumed) peaks AT the + # truth, so a linear slope alone is not an adequate summary. + i_truth = ALPHAS_ASSUMED.index(b2.ALPHA_TRUE) + d_mean_a = jnp.mean(delta_hat, axis=0) + drop_lo = float(d_mean_a[i_truth] - d_mean_a[0]) # vs alpha_assumed = 1.9 + drop_hi = float(d_mean_a[i_truth] - d_mean_a[-1]) # vs alpha_assumed = 2.7 + print(f" peak-to-end drops (ensemble mean): " + f"delta_hat({b2.ALPHA_TRUE}) - delta_hat({ALPHAS_ASSUMED[0]}) = " + f"{drop_lo:+.4f}; delta_hat({b2.ALPHA_TRUE}) - " + f"delta_hat({ALPHAS_ASSUMED[-1]}) = {drop_hi:+.4f} " + f"(max |Delta delta_hat| = {max(abs(drop_lo), abs(drop_hi)):.4f})") + print(" the measured response is PEAKED at the truth alpha, so a linear " + "slope is NOT an adequate summary -- wrong alpha in EITHER " + "direction biases delta_hat LOW (by ~0.2-0.3 at |Delta alpha| = 0.4).") + print(f" plateau: {n_seeds * n_alpha - n_no_plateau}/{n_seeds * n_alpha} " + f"fits plateaued") + print(f" bias-ensemble wall-time = {t_bias / 60.0:.1f} min " + f"(projected {proj_min:.1f} min)") + if t_bias / 60.0 > BIAS_BUDGET_MIN: + print(f" NOTE: actual wall-time exceeded the {BIAS_BUDGET_MIN:.0f}-min " + f"target by {t_bias / 60.0 - BIAS_BUDGET_MIN:.1f} min -- " + "post-hoc machine-load variance; the PROJECTION before the run " + "is the budget gate (and it passed).") + + # ------------------------------------------------------------------ # + # Ensemble 2: robustness grid (GATED at 3 sigma componentwise) + # ------------------------------------------------------------------ # + print("\n" + "=" * 72) + print("ROBUSTNESS GRID: full joint (alpha, delta, W0) refit per alpha_true") + print("=" * 72) + t_grid0 = time.perf_counter() + grid_pulls, grid_rows = [], [] + grid_ok = True + for i, a_true in enumerate(ALPHAS_TRUE_GRID): + print(f"\nalpha_true = {a_true} (key {SEED_BASE_GRID + i}; " + f"delta={b2.DELTA_TRUE}, W0={b2.W0_TRUE})") + data = b2.build_truth_data(key=jax.random.PRNGKey(SEED_BASE_GRID + i), + alpha_true=a_true) + n = data["n"] + print(f" top-group occupancy = {int(n[-1].sum())}; " + f"M_FIXED = {data['M_fixed']:.1f}; R_CUT = {data['r_cut']:.3f}") + theta_hat, sigma_theta, plat, finals, i_best = joint_refit(data) + print(f" init finals = {[round(f, 4) for f in finals]} " + f"(best: init {i_best}); plateau {'PASS' if plat else 'FAIL'}") + truths = (a_true, b2.DELTA_TRUE, b2.W0_TRUE) + pulls = b2.recovery_table(theta_hat, sigma_theta, truths=truths) + ok = all(abs(p) < 3.0 for p in pulls) + grid_ok = grid_ok and ok + print(f" 3-sigma componentwise: {'PASS' if ok else 'FAIL'} " + f"(max |pull| = {max(abs(p) for p in pulls):.3f})") + grid_pulls.append(pulls) + grid_rows.append((a_true, theta_hat, sigma_theta, pulls, ok)) + t_grid = time.perf_counter() - t_grid0 + print(f"\n robustness-grid wall-time = {t_grid / 60.0:.1f} min") + + print(f"\n{'alpha_true':>10} {'alpha_hat':>16} {'delta_hat':>16} " + f"{'W0_hat':>16} {'max|pull|':>9} {'gate':>5}") + for a_true, th, sg, pulls, ok in grid_rows: + print(f"{a_true:>10.1f} " + f"{float(th[0]):>9.4f}+/-{float(sg[0]):.4f} " + f"{float(th[1]):>9.4f}+/-{float(sg[1]):.4f} " + f"{float(th[2]):>9.4f}+/-{float(sg[2]):.4f} " + f"{max(abs(p) for p in pulls):>9.3f} " + f"{'PASS' if ok else 'FAIL':>5}") + + # ------------------------------------------------------------------ # + # Figure + verdict + # ------------------------------------------------------------------ # + out_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), + "validation", "plots") + os.makedirs(out_dir, exist_ok=True) + bias = dict(alphas=alphas, delta_hat=delta_hat, slope_mean=slope_mean, + slope_se=slope_se, drop_lo=drop_lo, drop_hi=drop_hi) + grid = dict(alphas_true=jnp.asarray(ALPHAS_TRUE_GRID), + pulls=jnp.array(grid_pulls)) + make_bias_figure(bias, grid, out_dir) + + print("\n" + "=" * 72) + print("GATES") + print("=" * 72) + print(f" robustness grid 3-sigma (all alpha_true, componentwise): " + f"{'PASS' if grid_ok else 'FAIL'}") + print(" bias curve: REPORTED (slope quoted with seed-ensemble " + "uncertainty; not gated by design)") + print("\n" + "=" * 72) + print(f"OVERALL {'ALL PASS' if grid_ok else 'FAIL'}") + print("=" * 72) + if not grid_ok: + print("\nNOTE: the 3-sigma robustness gate is REAL. A >3-sigma miss at " + "some alpha_true is a PHYSICS finding -- do NOT widen the gate; " + "report the table above.") + sys.exit(0 if grid_ok else 1) + + +if __name__ == "__main__": + if "--diag-21" in sys.argv[1:]: + diag_alpha21() + else: + main() diff --git a/scripts/demo_halo_core.py b/scripts/demo_halo_core.py new file mode 100644 index 00000000..21ff9643 --- /dev/null +++ b/scripts/demo_halo_core.py @@ -0,0 +1,784 @@ +r"""B3 science demo (Task 7): halo+core Engine-B recovery of (t, r_a, r_h). + +A two-family prescribed-density cluster -- a Plummer HALO (Osipkov-Merritt +anisotropic) plus an EFF(gamma=5) CORE (isotropic) -- is sampled in detailed +shared-potential equilibrium (Engine B: prescribed rho_j -> one quadrature Psi +-> per-component Eddington inversion; NO ODE). Three parameters are recovered by +joint MLE + Gauss-Newton Fisher from two binned observables: + + theta = (t, r_a, r_h) + t -- halo (Plummer) mass fraction (core = 1 - t), truth 0.6 + r_a -- halo Osipkov-Merritt anisotropy radius, truth 3.0 + r_h -- Plummer half-mass radius, truth 2.0 + + fixed: EFF core (a=0.8, gamma=5, r_t=9), m_j=[0.5, 1.0], core isotropic + (r_a=inf), domain r_t=9 (the EFF extent -- concrete by construction). + +Observables (clean-mock; no LOS projection, no errors, no incompleteness): + + * sigma channel -- per-component binned 1-D dispersion sigma_1d,j(r) for BOTH + components (di.binned_sigma_beta with component_id, n_min=50). + * beta channel -- the HALO (component 0) Binney anisotropy beta_hat(r). The + core is isotropic (truth beta ~ 0), gated on the halo only (the plan's + "halo beta_hat(r)"). + + Bins: 16 equal-count radial bins over [0, R_CUT] with R_CUT = max sampled + radius (~ the EFF edge r_t=9). R_EDGES is FROZEN after one truth draw. + +predict_fn (model side; plan amendment 4 -- BINNED EXPECTATION, no bin-center +oracle). Inside the traced loss the Engine-B state is rebuilt from theta: + + m = from_density_profiles([Plummer(r_h), EFF(a=0.8,gamma=5,r_t=9)], + [t, 1-t], m_j=[0.5,1.0], r_a_j=[r_a, inf], r_t=9) + st = m.engine_b # r_poisson, Psi_poisson, dPsi_dr_poisson, E_grid, + # f_j_grid [J, n_e], rho_j_poisson [J, n_r], r_a_j, f_min_j + + r_t=9 is passed EXPLICITLY: the domain extent is the (fixed) EFF r_t, NOT a + function of the recovered theta, and derive_r_t concretizes its max-extent + pick (float()), so a traced r_h alone would raise ConcretizationTypeError. The + explicit override gives the IDENTICAL model the truth draw uses (derived r_t + there is also 9.0) while keeping the build fully traceable/differentiable. + + sigma recipe per component j (VERBATIM from _predicted_component_Q, + tests/validation/test_engine_b_physics.py:146-161, replicated as JAX): + + moments(Psi_r, f_row): w = linspace(0, sqrt(2 Psi_r), n_w); + f_at = max(interp(Psi_r - w^2/2, E_grid, f_row), 0); + m0 = trapz(w^2 f_at, w), m2 = trapz(w^4 f_at, w) + v2(r) = (m2/(m0+1e-300)) * (1/3 + (2/3) * inv_st2(r)), + inv_st2 = 1/(1 + (r/r_a)^2) (=1 for the isotropic core) + sigma_j^2(r) = v2(r) / 3 # 1-D dispersion squared + + BINNED EXPECTATION (number-weighted by the prescribed density rho_j_poisson[j], + the 1/m_j cancels in the per-bin ratio): + + E[sigma_hat^2_jk] = (int_bin rho_j sigma_j^2 dr) / (int_bin rho_j dr) + + via a cumulative-trapezoid-at-bin-edges trick (cumsum of sigma^2 dN ratioed to + cumsum of dN at the bin edges, jnp.interp at edges) -- the same construction as + demo_delta_recovery.predict_binned. predict returns sqrt of that, shape (J, K), + with the double-`where` safe-sqrt (clamp BEFORE sqrt so empty bins keep a + finite gradient). + + beta channel predict -- beta(r) = r^2/(r^2 + r_a^2) closed-form in r_a (the OM + definition, amendment 3; NOT from the f-moments), number-weighted the SAME way + per bin for a like-with-like comparison to the binned beta_hat (small + correction; beta is smooth). + +Loss: residual_fn(z) stacks the weighted sigma residuals (both components) and +the weighted halo-beta residuals; negloglike = 0.5 * sum residual^2. The Fisher +information is the Gauss-Newton J^T J (di.fisher_information_gn, jacrev) mapped to +theta-space by di.constrained_cov with the dtheta/dz box Jacobian. MLE is +di.mle_adam from dispersed inits with a plateau check. + +Gates (REAL -- exit nonzero on any failure; never weaken): + 1. 3-sigma recovery: |theta_hat - theta_true| < 3 sigma_hat for t, r_a, r_h. + 2. Realizability: rebuild Engine B at theta_hat; min(f_min_j) >= -1e-3. + 3. Plateau: MLE last-10% improvement < 1% of the total decrease. + 4. Occupancy: each component has enough resolved sigma bins (>= MIN_BINS_PER_CMP) + to constrain the fit (STOP, don't weaken, if a component is starved). + +Runtime (measured 2026-06-11, CPU/float64; XLA single-thread; one full run) +-------------------------------------------------------------------------- + R_CUT (cluster edge = max radius) 8.998 (the EFF r_t; all stars used) + R_EDGES 16 quantile bins on [0, R_CUT], FROZEN + per-component sigma-bin occupancy halo 16/16, core 16/16 (n_min=50) + Engine B is QUADRATURE (no ODE), but value_and_grad through a rebuild is + not free (Eddington inversion + OM moments). + compile (cold value_and_grad) 1.44 s + warm jit(value_and_grad) 0.137 s -- PASS (<= 5 s budget) + grad(z1) finite AND nonzero in all 3 of (t, r_a, r_h). + end-to-end wall-time ~3.7 min -- DOMINATED by the MLE stage: + three dispersed-init Adam runs, each a jit-compiled 400-step lax.scan + whose body rebuilds Engine B and backprops; the three scan COMPILES are + the bottleneck (the warm per-eval cost above is tiny by comparison). +Self-consistency at truth: predict(truth) vs sig_hat/beta_hat over all populated +cells -- max |dev/se| = 2.34 (binned-expectation is like-with-like). +Recovery (this run): t=0.602+/-0.008, r_a=3.05+/-0.05, r_h=2.02+/-0.03 (all <1sigma). +""" + +import os +import sys +import time + +import jax +import jax.numpy as jnp + +import progenax # noqa: F401 -- enables float64 at import +from progenax import EFFProfile, MultiComponentCluster, PlummerProfile +from jaxstro.units import STELLAR + +sys.path.insert(0, os.path.dirname(__file__)) +import _demo_inference as di # noqa: E402 + +# --------------------------------------------------------------------------- # +# Truth configuration (module constants) -- verbatim from +# tests/validation/test_engine_b_physics.py:_headline_model + :376. +# --------------------------------------------------------------------------- # +G = STELLAR.G # match test_engine_b_physics.py:41 +T_TRUE = 0.6 # halo (Plummer) mass fraction +R_A_TRUE = 3.0 # halo OM anisotropy radius +R_H_TRUE = 2.0 # Plummer half-mass radius +M_J = [0.5, 1.0] # stellar-mass labels, FIXED +N_STARS = 30_000 + +# Fixed EFF core shape + domain (the EFF extent is the model's r_t). +EFF_A, EFF_GAMMA, EFF_RT = 0.8, 5.0, 9.0 + +# Unconstrained-reparam: t via expit box; r_a, r_h via exp (log-parametrized). +T_BOX = (0.3, 0.9) # truth 0.6 well interior + +# Quadrature resolution of the f_j speed-moment integrals (fixed -> grad-safe). +N_W = 400 # matches _predicted_component_Q n_w +N_BINS = 16 +N_MIN = 50 # binned_sigma_beta occupancy floor +MIN_BINS_PER_CMP = 8 # occupancy gate: resolved sigma bins +_TINY = 1e-30 + + +# --------------------------------------------------------------------------- # +# Box reparametrization theta = (t, r_a, r_h) <- z in R^3 +# --------------------------------------------------------------------------- # +def _theta_of_z(z): + """z in R^3 -> theta = (t in T_BOX via expit, r_a = exp(z1), r_h = exp(z2)).""" + return (di.expit(z[0], *T_BOX), jnp.exp(z[1]), jnp.exp(z[2])) + + +def _dtheta_dz(z): + """Per-component dtheta_i/dz_i at z (shape (3,)): expit' on t, exp' on r_a, r_h.""" + return jnp.array([ + jax.grad(lambda zi: di.expit(zi, *T_BOX))(z[0]), + jnp.exp(z[1]), # d(exp z)/dz = exp z = r_a + jnp.exp(z[2]), # = r_h + ]) + + +# --------------------------------------------------------------------------- # +# Engine-B rebuild (traced) + per-component sigma_1d(r) on the Poisson grid +# --------------------------------------------------------------------------- # +def _engine_b(theta): + """Rebuild the Engine-B state from theta=(t, r_a, r_h) (traceable). + + r_t=EFF_RT passed EXPLICITLY so the (fixed) EFF extent sets the domain + without concretizing a traced r_h (derive_r_t's max-extent float() pick). + """ + t, r_a, r_h = theta + m = MultiComponentCluster.from_density_profiles( + [PlummerProfile(r_h=r_h), + EFFProfile(a=EFF_A, gamma=EFF_GAMMA, r_t=EFF_RT)], + jnp.stack([t, 1.0 - t]), + m_j=jnp.array(M_J), + r_a_j=jnp.stack([r_a, jnp.array(jnp.inf)]), + r_t=EFF_RT, + ) + return m.engine_b + + +def _sigma_sq_on_grid(st, m_fixed): + r"""sigma_j^2(r) = _j(r) / 3 on st.r_poisson for every component j. + + VERBATIM the _predicted_component_Q speed-moment recipe + (test_engine_b_physics.py:146-157): the f_j-row speed moments m0, m2 give + _DF,j = (m2/m0)(1/3 + (2/3)/(1 + r^2/r_a_j^2)); sigma_1d^2 = /3. + Psi clamped to 1e-12 before the sqrt(2 Psi) integration limit (Psi_safe). + + PHYSICAL velocity scale (the EFF kappa = G M / (4 pi mu) pattern, + eddington_engine.py:257-261): the f_j tables are built on mass-normalized + densities (Psi_dimless, M=1), so the sampled speeds carry a factor + v = sqrt(G * M_sampled / (4 pi mu)) -> sigma^2_phys = (G m_fixed / (4 pi mu)) + * sigma^2_dimless, where 4 pi mu == 1 identically (computed explicitly from + rho_j_poisson for honesty; theta-independent). ``m_fixed`` is the measured + total cluster mass (a fixed observed scalar, like B2's M_FIXED), so sigma is + DATA-ANCHORED rather than re-derived from theta inside the loss. + """ + r = st.r_poisson + Psi_safe = jnp.maximum(st.Psi_poisson, 1e-12) + # 4 pi mu = 4 pi sum_j int rho_j r^2 dr == 1 for mass-normalized rho_j; carry + # it explicitly (theta-independent) rather than hard-coding the identity. + mu = jnp.trapezoid(jnp.sum(st.rho_j_poisson, axis=0) * r**2, r) + vel_scale_sq = G * m_fixed / (4.0 * jnp.pi * mu) # = sigma^2_phys / sigma^2_dimless + + def moments(Psi_r, f_row): + w = jnp.linspace(0.0, jnp.sqrt(2.0 * Psi_r), N_W) + f_at = jnp.maximum(jnp.interp(Psi_r - 0.5 * w**2, st.E_grid, f_row), 0.0) + return jnp.trapezoid(w**2 * f_at, w), jnp.trapezoid(w**4 * f_at, w) + + def per_comp(j): + m0, m2 = jax.vmap(lambda P: moments(P, st.f_j_grid[j]))(Psi_safe) + finite = jnp.isfinite(st.r_a_j[j]) + ra_safe = jnp.where(finite, st.r_a_j[j], 1.0) + inv_st2 = jnp.where(finite, 1.0 / (1.0 + (r / ra_safe) ** 2), 1.0) + v2 = (m2 / (m0 + 1e-300)) * (1.0 / 3.0 + (2.0 / 3.0) * inv_st2) + return vel_scale_sq * v2 / 3.0 # sigma_1d^2(r) PHYSICAL + + return jax.vmap(per_comp)(jnp.arange(st.f_j_grid.shape[0])) # (J, n_r) + + +def _bin_average(value_grid, weight_grid, r_grid, r_edges): + r"""Number-weighted bin average of ``value_grid`` over ``r_edges``. + + _k = (int_bin w value dr) / (int_bin w dr), computed by the same + cumulative-trapezoid-at-edges trick as predict_binned: build the cumulative + integrals C_val = int_0^r w value dr and C_w = int_0^r w dr on the grid, then + take the edge differences via jnp.interp. ``value_grid``, ``weight_grid``, + ``r_grid`` are 1-D over the Poisson grid; returns (K,) bin averages. + """ + dr = jnp.diff(r_grid) + wval = weight_grid * value_grid + wval_mid = 0.5 * (wval[1:] + wval[:-1]) + w_mid = 0.5 * (weight_grid[1:] + weight_grid[:-1]) + C_val = jnp.concatenate([jnp.zeros(1), jnp.cumsum(wval_mid * dr)]) + C_w = jnp.concatenate([jnp.zeros(1), jnp.cumsum(w_mid * dr)]) + lo, hi = r_edges[:-1], r_edges[1:] + num = jnp.interp(hi, r_grid, C_val) - jnp.interp(lo, r_grid, C_val) + den = jnp.interp(hi, r_grid, C_w) - jnp.interp(lo, r_grid, C_w) + return num / (den + _TINY) + + +# --------------------------------------------------------------------------- # +# Binned-expectation predictors (sigma both components; beta halo only) +# --------------------------------------------------------------------------- # +def predict_sigma_binned(theta, r_edges, m_fixed): + r"""Per-(component, bin) sqrt(E[sigma_hat^2_jk]) under the model's rho_j(r). + + Number-weighted bin average of sigma_j^2(r) with weight rho_j_poisson[j] + (the prescribed density; 1/m_j cancels in the ratio), then sqrt with the + double-`where` safe pattern (clamp BEFORE sqrt -> empty/zero bins keep a + finite gradient). Returns shape (J, K). ``m_fixed`` is the measured total + mass that sets the physical velocity scale (see _sigma_sq_on_grid). + """ + st = _engine_b(theta) + r = st.r_poisson + sig2 = _sigma_sq_on_grid(st, m_fixed) # (J, n_r) + # NUMBER weight n_j(r) dr = 4 pi r^2 rho_j(r) dr (the data sig_hat pools over + # STARS, so the bin average is number-weighted; the r^2 shell-volume factor + # is essential in the wide outer truncation cell -- amendment 4). The 4 pi + # and 1/m_j constants cancel in the per-bin ratio, so r^2 rho_j suffices. + n_weight = r**2 * st.rho_j_poisson # (J, n_r) + + def per_comp(j): + sig2_bar = _bin_average(sig2[j], n_weight[j], r, r_edges) # (K,) + safe = jnp.where(sig2_bar > 0.0, sig2_bar, 1.0) + sig = jnp.sqrt(safe) + return jnp.where(sig2_bar > 0.0, sig, 0.0) + return jax.vmap(per_comp)(jnp.arange(sig2.shape[0])) # (J, K) + + +def predict_beta_halo_binned(theta, r_edges): + r"""Per-bin number-weighted OM beta(r)=r^2/(r^2+r_a^2) for the HALO (j=0). + + Closed-form in r_a (amendment 3), number-weighted by rho_0(r) the same way + the sigma channel is, for a like-with-like comparison to the binned + beta_hat. Returns shape (K,). + """ + st = _engine_b(theta) + r = st.r_poisson + r_a = st.r_a_j[0] + beta = r**2 / (r**2 + r_a**2) # OM, halo finite r_a + n_weight = r**2 * st.rho_j_poisson[0] # 4 pi const cancels + return _bin_average(beta, n_weight, r, r_edges) # (K,) + + +def sigma_curve(theta, r_eval, m_fixed): + """Pointwise sigma_1d,j(r) at theta on r_eval (FOR PLOTTING only, not loss). + + Interpolates the on-grid sigma_j(r) = sqrt(sigma_j^2(r)) onto r_eval. + Returns (J, len(r_eval)). + """ + st = _engine_b(theta) + sig2 = _sigma_sq_on_grid(st, m_fixed) + sig = jnp.sqrt(jnp.maximum(sig2, 0.0)) + + def per_comp(j): + return jnp.interp(r_eval, st.r_poisson, sig[j]) + return jax.vmap(per_comp)(jnp.arange(sig.shape[0])) + + +# --------------------------------------------------------------------------- # +# Mock data construction (run once; R_CUT, R_EDGES then frozen) +# --------------------------------------------------------------------------- # +def build_truth_data(key=None): + r"""Sample the truth halo+core cluster, bin sigma_1d,j(r) + halo beta_hat(r). + + Returns a dict of constants the loss closes over: R_EDGES (16 quantile bins + on [0, R_CUT]), per-component sig_hat/se/weight/n, the halo beta_hat/se, the + per-(comp, bin) mean radius (for the figure), and R_CUT. + """ + if key is None: + key = jax.random.PRNGKey(0) + + model = MultiComponentCluster.from_density_profiles( + [PlummerProfile(r_h=R_H_TRUE), + EFFProfile(a=EFF_A, gamma=EFF_GAMMA, r_t=EFF_RT)], + jnp.array([T_TRUE, 1.0 - T_TRUE]), + m_j=jnp.array(M_J), + r_a_j=jnp.array([R_A_TRUE, jnp.inf]), + ) + ic = model.sample_cluster(key, n_stars=N_STARS, G=G) + + # Mass-weighted COM subtraction (the validation tests' _com_arrays). + pos = ic.positions - jnp.average(ic.positions, axis=0, weights=ic.masses) + vel = ic.velocities - jnp.average(ic.velocities, axis=0, weights=ic.masses) + cid = ic.component_id # int in [0, J) + + r = jnp.linalg.norm(pos, axis=1) + R_CUT = float(r.max()) # cluster edge ~ r_t + # 16 equal-count radial bins on [0, R_CUT], FROZEN after this one draw. + R_EDGES = jnp.quantile(r, jnp.linspace(0.0, 1.0, N_BINS + 1)) + + res = di.binned_sigma_beta(pos, vel, R_EDGES, component_id=cid, n_min=N_MIN) + sig_hat, se, beta_hat, weight, n = res + + # beta SE: same 3-pooled-component effective sample size as sig_hat. With + # sigma_r, sigma_t each ~ stable, beta_hat = 1 - sigma_t^2/(2 sigma_r^2) + # has Var ~ (1-beta)^2 (1/N_t + ... ) ~ O(1/n); use a delta-method estimate + # beta_se ~ (1 + |beta_hat|) / sqrt(n) -- a conservative honest SE (used only + # as the beta-channel weight; no gate is loosened by it). + beta_se = jnp.where(weight > 0, (1.0 + jnp.abs(beta_hat)) + / jnp.sqrt(jnp.maximum(n, 1.0)), 0.0) + + # Per-(comp, bin) mean radius (for placing the figure's data points). + bin_ids = di._bin_index(r, R_EDGES) + r_sum, r_cnt = di._grouped_bin_sums(r, cid, bin_ids, sig_hat.shape[0], + R_EDGES.shape[0] - 1) + r_mean = jnp.where(r_cnt > 0, r_sum / jnp.where(r_cnt > 0, r_cnt, 1.0), 0.0) + + M_fixed = float(jnp.sum(ic.masses)) # measured total mass (a CONSTANT) + + return dict(r_edges=R_EDGES, r_cut=R_CUT, sig_hat=sig_hat, se=se, + weight=weight, n=n, beta_hat=beta_hat, beta_se=beta_se, + r_mean=r_mean, M_fixed=M_fixed) + + +# --------------------------------------------------------------------------- # +# Residual + joint negloglike (sigma both components + halo beta) +# --------------------------------------------------------------------------- # +def make_residual_fn(data): + r"""Standardized residual vector r(z): sigma (both comps) + halo beta cells. + + r = [ sqrt(w_sig) (sig_hat - sig_model)/se ; sqrt(w_b) (beta_hat - beta_model)/beta_se ] + so -0.5 sum r^2 is the joint Gaussian log-likelihood; jacrev(r) gives the + Gauss-Newton Fisher. Masked cells (weight 0) contribute 0. + """ + r_edges = data["r_edges"] + m_fixed = data["M_fixed"] + sig_hat, se, weight = data["sig_hat"], data["se"], data["weight"] + beta_hat, beta_se = data["beta_hat"], data["beta_se"] + w_halo = weight[0] # halo beta weight + safe_se = jnp.where(se > 0, se, 1.0) + safe_bse = jnp.where(beta_se[0] > 0, beta_se[0], 1.0) + sqrt_w = jnp.sqrt(weight) + sqrt_wb = jnp.sqrt(w_halo) + + def residual_fn(z): + theta = _theta_of_z(z) + sig_model = predict_sigma_binned(theta, r_edges, m_fixed) # (J, K) + beta_model = predict_beta_halo_binned(theta, r_edges) # (K,) + r_sig = (sqrt_w * (sig_hat - sig_model) / safe_se).ravel() + r_beta = sqrt_wb * (beta_hat[0] - beta_model) / safe_bse + return jnp.concatenate([r_sig, r_beta]) + + return residual_fn + + +def make_negloglike(data): + """Joint negloglike(z) = 0.5 sum residual(z)^2 (sigma + halo beta).""" + residual_fn = make_residual_fn(data) + + def negloglike(z): + r = residual_fn(z) + return 0.5 * jnp.sum(r * r) + + return negloglike + + +# --------------------------------------------------------------------------- # +# Self-consistency check (predict(truth) vs data, in units of SE) +# --------------------------------------------------------------------------- # +def self_consistency_check(data): + """predict(truth) vs sig_hat/beta_hat in populated cells (units of se).""" + truth = (T_TRUE, R_A_TRUE, R_H_TRUE) + r_edges = data["r_edges"] + m_fixed = data["M_fixed"] + sig_model = predict_sigma_binned(truth, r_edges, m_fixed) + beta_model = predict_beta_halo_binned(truth, r_edges) + sig_hat, se, weight = data["sig_hat"], data["se"], data["weight"] + beta_hat, beta_se = data["beta_hat"], data["beta_se"] + safe_se = jnp.where(se > 0, se, 1.0) + safe_bse = jnp.where(beta_se[0] > 0, beta_se[0], 1.0) # (K,) halo + dev_sig = jnp.where(weight > 0, (sig_hat - sig_model) / safe_se, 0.0) + dev_beta = jnp.where(weight[0] > 0, (beta_hat[0] - beta_model) / safe_bse, 0.0) + + print("\nSelf-consistency: predict(truth) vs data (deviation in SE)") + print(f"{'chan':>6} {'cmp':>3} {'bin':>3} {'data':>9} {'pred':>9} " + f"{'se':>8} {'dev/se':>8}") + max_dev = 0.0 + for j in range(sig_hat.shape[0]): + for k in range(sig_hat.shape[1]): + if float(weight[j, k]) > 0: + d = float(dev_sig[j, k]) + max_dev = max(max_dev, abs(d)) + print(f"{'sigma':>6} {j:>3} {k:>3} " + f"{float(sig_hat[j, k]):>9.4f} {float(sig_model[j, k]):>9.4f} " + f"{float(se[j, k]):>8.4f} {d:>8.2f}") + for k in range(beta_hat.shape[1]): + if float(weight[0, k]) > 0: + d = float(dev_beta[k]) + max_dev = max(max_dev, abs(d)) + print(f"{'beta':>6} {0:>3} {k:>3} " + f"{float(beta_hat[0, k]):>9.4f} {float(beta_model[k]):>9.4f} " + f"{float(beta_se[0, k]):>8.4f} {d:>8.2f}") + print(f"max |dev/se| over populated cells = {max_dev:.2f}") + return max_dev + + +# --------------------------------------------------------------------------- # +# MLE recovery from dispersed inits + Gauss-Newton Fisher errors +# --------------------------------------------------------------------------- # +N_INITS = 3 +INIT_KEY = 7 +INIT_SCALE = 0.5 +N_ADAM_STEPS = 400 +ADAM_LR = 3e-2 + + +def _z_truth(): + """Unconstrained z at the truth (the z0=truth init starting point).""" + return jnp.array([di.logit(T_TRUE, *T_BOX), jnp.log(R_A_TRUE), jnp.log(R_H_TRUE)]) + + +def dispersed_inits(): + """3 unconstrained inits: z at truth plus 2 draws from N(z_truth, scale^2 I).""" + key = jax.random.PRNGKey(INIT_KEY) + z_t = _z_truth() + draws = z_t[None, :] + jax.random.normal(key, (N_INITS - 1, 3)) * INIT_SCALE + return jnp.concatenate([z_t[None, :], draws], axis=0) + + +def plateau_ok(trace, frac=0.1, rel_tol=0.01): + """Tail-decrease < rel_tol * total-decrease over the loss trace (lagged trace).""" + n = trace.shape[0] + k_tail = int(n * (1.0 - frac)) + total = float(trace[0] - trace[-1]) + tail = float(trace[k_tail] - trace[-1]) + if total <= 0: + return False, total, tail + return (tail < rel_tol * total), total, tail + + +def run_mle(negloglike): + """Adam MLE from N_INITS dispersed inits; return the lowest-loss result.""" + loss_jit = jax.jit(negloglike) + z0s = dispersed_inits() + z_hats, traces, finals = [], [], [] + for i in range(N_INITS): + z_hat, trace = di.mle_adam(negloglike, z0s[i], + n_steps=N_ADAM_STEPS, lr=ADAM_LR) + finals.append(float(loss_jit(z_hat))) + z_hats.append(z_hat) + traces.append(trace) + i_best = int(jnp.argmin(jnp.array(finals))) + return z_hats[i_best], traces[i_best], finals, i_best, z0s + + +def recovery_table(theta_hat, sigma_theta): + """Print param | truth | theta_hat | sigma_hat | pull = (hat-truth)/sigma.""" + truths = (T_TRUE, R_A_TRUE, R_H_TRUE) + names = ("t", "r_a", "r_h") + print(f"\n{'param':>6} {'truth':>8} {'theta_hat':>10} {'sigma_hat':>10} " + f"{'(hat-truth)/sigma':>18}") + pulls = [] + for nm, tr, th, sg in zip(names, truths, theta_hat, sigma_theta): + pull = (float(th) - tr) / float(sg) + pulls.append(pull) + print(f"{nm:>6} {tr:>8.4f} {float(th):>10.4f} {float(sg):>10.4f} " + f"{pull:>18.3f}") + return pulls + + +# --------------------------------------------------------------------------- # +# Figures: data-vs-fit sigma panels + halo beta panel + Fisher ellipses +# --------------------------------------------------------------------------- # +def _ellipse_xy(mean2, cov2, n_sigma=2.0, n_pts=200): + """(x, y) of the n_sigma covariance ellipse for a 2-D Gaussian.""" + t = jnp.linspace(0.0, 2.0 * jnp.pi, n_pts) + circle = jnp.stack([jnp.cos(t), jnp.sin(t)], axis=0) + L = jnp.linalg.cholesky(cov2) + pts = mean2[:, None] + n_sigma * (L @ circle) + return pts[0], pts[1] + + +def make_figure(data, theta_hat, sigma_theta, cov_theta, out_dir): + """4-panel figure: sigma_j(r) (2 comps) + halo beta(r) + 2 Fisher ellipses.""" + import matplotlib.pyplot as plt + import numpy as np + + sys.path.insert(0, os.path.dirname(__file__)) + import _plotstyle as ps # noqa: E402 + + ps.apply_pub_style() + colors = [ps.OI["blue"], ps.OI["vermilion"]] + cmp_names = ["halo (Plummer)", "core (EFF)"] + + sig_hat = np.asarray(data["sig_hat"]) + se = np.asarray(data["se"]) + weight = np.asarray(data["weight"]) + r_mean = np.asarray(data["r_mean"]) + beta_hat = np.asarray(data["beta_hat"]) + beta_se = np.asarray(data["beta_se"]) + m_fixed = data["M_fixed"] + t_hat, ra_hat, rh_hat = (float(x) for x in theta_hat) + + fig, axes = plt.subplots(2, 2, figsize=(8.4, 7.0)) + ax_h, ax_c, ax_b, ax_e = axes[0, 0], axes[0, 1], axes[1, 0], axes[1, 1] + + # Panels (a),(b): per-component sigma(r) data (SE bars) + best-fit curves. + r_grid = jnp.linspace(1e-3, float(data["r_cut"]), 400) + sig_fit = np.asarray(sigma_curve((t_hat, ra_hat, rh_hat), r_grid, m_fixed)) + sig_true = np.asarray(sigma_curve((T_TRUE, R_A_TRUE, R_H_TRUE), r_grid, m_fixed)) + rg = np.asarray(r_grid) + for j, (ax, lab) in enumerate([(ax_h, "(a)"), (ax_c, "(b)")]): + mask = weight[j] > 0 + ax.errorbar(r_mean[j][mask], sig_hat[j][mask], yerr=se[j][mask], + fmt="o", ms=3.5, color=colors[j], capsize=1.5, + elinewidth=0.8, mew=0.0, zorder=3, label="mock data") + ax.plot(rg, sig_fit[j], "-", color=colors[j], lw=1.6, zorder=2, + label="best fit") + ax.plot(rg, sig_true[j], ":", color=ps.OI["black"], lw=0.9, alpha=0.7, + zorder=1, label="truth") + ax.set_xlabel(r"$r$ (model units)") + ax.set_ylabel(rf"$\sigma_{{1\mathrm{{D}}}}(r)$ -- {cmp_names[j]}") + ax.set_xlim(0.0, float(data["r_cut"])) + ax.legend(loc="upper right") + ps.panel_label(ax, lab) + + # Panel (c): halo beta_hat(r) + the closed-form OM beta(r; r_a_hat). + mask = weight[0] > 0 + ax_b.errorbar(r_mean[0][mask], beta_hat[0][mask], yerr=beta_se[0][mask], + fmt="o", ms=3.5, color=ps.OI["blue"], capsize=1.5, + elinewidth=0.8, mew=0.0, zorder=3, label=r"halo $\hat\beta(r)$") + beta_fit = rg**2 / (rg**2 + ra_hat**2) + beta_true = rg**2 / (rg**2 + R_A_TRUE**2) + ax_b.plot(rg, beta_fit, "-", color=ps.OI["blue"], lw=1.6, zorder=2, + label=rf"OM fit $r_a={ra_hat:.2f}$") + ax_b.plot(rg, beta_true, ":", color=ps.OI["black"], lw=0.9, alpha=0.7, + zorder=1, label=rf"truth $r_a={R_A_TRUE}$") + ax_b.axhline(0.0, color=ps.OI["green"], lw=0.8, ls="--", alpha=0.6, + label="core (isotropic)") + ax_b.set_xlabel(r"$r$ (model units)") + ax_b.set_ylabel(r"$\beta(r)$ (Binney anisotropy)") + ax_b.set_xlim(0.0, float(data["r_cut"])) + ax_b.legend(loc="lower right") + ps.panel_label(ax_b, "(c)") + + # Panel (d): 2-sigma Fisher ellipses for (r_a, r_h) and (t, r_a). + cov = np.asarray(cov_theta) + # (r_a, r_h) block = indices (1, 2); (t, r_a) block = indices (0, 1). + mean_rarh = jnp.array([ra_hat, rh_hat]) + cov_rarh = jnp.asarray(cov[1:, 1:]) + mean_tra = jnp.array([t_hat, ra_hat]) + cov_tra = jnp.asarray(cov[jnp.ix_(jnp.array([0, 1]), jnp.array([0, 1]))]) + + # Primary: (r_a, r_h) 2-sigma ellipse + MLE + truth. + x1, y1 = _ellipse_xy(mean_rarh, cov_rarh, n_sigma=2.0) + ax_e.plot(np.asarray(x1), np.asarray(y1), "-", color=ps.OI["purple"], lw=1.8, + label=r"$2\sigma$ $(r_a, r_h)$") + ax_e.fill(np.asarray(x1), np.asarray(y1), color=ps.OI["purple"], alpha=0.10) + ax_e.plot(ra_hat, rh_hat, "o", color=ps.OI["black"], ms=5, zorder=5, + label=r"MLE $\hat\theta$") + ax_e.plot(R_A_TRUE, R_H_TRUE, "*", color=ps.OI["orange"], ms=13, zorder=6, + mec=ps.OI["black"], mew=0.5, label="truth") + ax_e.set_xlabel(r"$r_a$ (halo OM radius)") + ax_e.set_ylabel(r"$r_h$ (Plummer half-mass)") + rho_rarh = float(cov[1, 2] / np.sqrt(cov[1, 1] * cov[2, 2])) + rho_tra = float(cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])) + ax_e.legend(loc="best") + ps.panel_label(ax_e, "(d)") + + # Inset: the second requested pair (t, r_a) 2-sigma ellipse. + axin = ax_e.inset_axes([0.62, 0.10, 0.35, 0.35]) + x2, y2 = _ellipse_xy(mean_tra, cov_tra, n_sigma=2.0) + axin.plot(np.asarray(x2), np.asarray(y2), "-", color=ps.OI["green"], lw=1.4) + axin.fill(np.asarray(x2), np.asarray(y2), color=ps.OI["green"], alpha=0.10) + axin.plot(t_hat, ra_hat, "o", color=ps.OI["black"], ms=3, zorder=5) + axin.plot(T_TRUE, R_A_TRUE, "*", color=ps.OI["orange"], ms=8, zorder=6, + mec=ps.OI["black"], mew=0.4) + axin.set_xlabel(r"$t$", fontsize=7, labelpad=1) + axin.set_ylabel(r"$r_a$", fontsize=7, labelpad=1) + axin.tick_params(labelsize=6) + axin.set_title(r"$2\sigma\,(t, r_a)$", fontsize=7, pad=2) + + cap = (rf"$\hat t={t_hat:.3f}\pm{float(sigma_theta[0]):.3f},\ " + rf"\hat r_a={ra_hat:.3f}\pm{float(sigma_theta[1]):.3f},\ " + rf"\hat r_h={rh_hat:.3f}\pm{float(sigma_theta[2]):.3f}$; " + rf"$\rho(r_a,r_h)={rho_rarh:+.2f},\ \rho(t,r_a)={rho_tra:+.2f}$") + fig.text(0.5, -0.01, cap, ha="center", va="top", fontsize=8.5) + fig.tight_layout() + ps.save_fig(fig, out_dir, "demo_halo_core") + print(f"\nfigure -> {out_dir}/demo_halo_core.png (+ .pdf)") + return rho_rarh, rho_tra + + +# --------------------------------------------------------------------------- # +# Main driver +# --------------------------------------------------------------------------- # +def main(): + print("=" * 72) + print("B3 demo (Task 7): halo+core Engine-B (t, r_a, r_h) MLE + Fisher") + print("=" * 72) + print(f"truth: t={T_TRUE}, r_a={R_A_TRUE}, r_h={R_H_TRUE}; " + f"core EFF(a={EFF_A}, gamma={EFF_GAMMA}, r_t={EFF_RT}), " + f"m_j={M_J}, N_STARS={N_STARS}, G={G:.5g}") + + data = build_truth_data() + sig_hat = data["sig_hat"] + J = sig_hat.shape[0] + weight = data["weight"] + print(f"\nR_CUT (cluster edge = max radius) = {data['r_cut']:.4f}") + print("R_EDGES (16 quantile bins on [0, R_CUT]) =", + [round(float(x), 4) for x in data["r_edges"]]) + print("per-component sigma-bin counts (component x bin):") + n = data["n"] + resolved = [] + for j in range(J): + nb = int(jnp.sum(weight[j] > 0)) + resolved.append(nb) + print(f" component {j}: {[int(x) for x in n[j]]} " + f"({nb}/{sig_hat.shape[1]} resolved)") + + # Self-consistency BEFORE the budget gate. + max_dev = self_consistency_check(data) + sc_ok = max_dev < 4.0 + print(f"self-consistency {'OK' if sc_ok else 'FAIL'} " + f"(max |dev/se| = {max_dev:.2f}, threshold 4.0)") + if not sc_ok: + print("\nSTOP: predict(truth) does not match the data -- miscalibrated " + "oracle or binning. Not proceeding.") + sys.exit(1) + + # Budget gate: one warm jit(value_and_grad). + negloglike = make_negloglike(data) + loss_and_grad = jax.jit(jax.value_and_grad(negloglike)) + z0 = _z_truth() + + t0 = time.perf_counter() + v, g = loss_and_grad(z0) + v.block_until_ready(); g.block_until_ready() + t_compile = time.perf_counter() - t0 + + z1 = z0 + jnp.array([0.2, 0.15, -0.1]) + t0 = time.perf_counter() + v2, g2 = loss_and_grad(z1) + v2.block_until_ready(); g2.block_until_ready() + t_warm = time.perf_counter() - t0 + + grad_finite = bool(jnp.all(jnp.isfinite(g2))) + grad_nonzero = bool(jnp.all(jnp.abs(g2) > 0)) + print("\n" + "=" * 72) + print("RUNTIME BUDGET GATE") + print("=" * 72) + print(f"compile {t_compile:.2f}s, warm {t_warm:.3f}s") + print(f"loss(z0) = {float(v):.4e}, loss(z1) = {float(v2):.4e}") + print(f"grad(z1) = {[float(x) for x in g2]}") + print(f"grad finite: {grad_finite}; grad nonzero (all 3): {grad_nonzero}") + budget_ok = t_warm <= 5.0 + print(f"\nBUDGET {'PASS' if budget_ok else 'STOP'} " + f"(warm {t_warm:.3f}s {'<=' if budget_ok else '>'} 5.0s)") + if not budget_ok: + print("STOP: warm loss+grad exceeds 5 s. NOT degrading the solve. " + "Report measured numbers + options to the orchestrator.") + sys.exit(2) + + # MLE recovery. + print("\n" + "=" * 72) + print("MLE RECOVERY (Adam, %d dispersed inits)" % N_INITS) + print("=" * 72) + z_hat, best_trace, finals, i_best, z0s = run_mle(negloglike) + for i in range(N_INITS): + tag = " <-- WON" if i == i_best else "" + print(f" init {i} (z0={[round(float(x), 2) for x in z0s[i]]}): " + f"final negloglike = {finals[i]:.6e}{tag}") + plat_ok, total_dec, tail_dec = plateau_ok(best_trace) + print(f"\nplateau check (winning trace): total decrease = {total_dec:.4e}, " + f"last-10% decrease = {tail_dec:.4e}") + print(f"plateau {'PASS' if plat_ok else 'FAIL'} " + f"(tail {tail_dec:.3e} {'<' if plat_ok else '>='} " + f"1% of total {0.01 * total_dec:.3e})") + + _, g_hat = loss_and_grad(z_hat) + print(f"grad norm at z_hat = {float(jnp.linalg.norm(g_hat)):.4e} " + "(interior-optimum check)") + + # theta_hat + Gauss-Newton Fisher errors. + theta_hat = _theta_of_z(z_hat) + residual_fn = make_residual_fn(data) + F_z = di.fisher_information_gn(residual_fn, z_hat) + eig = jnp.linalg.eigvalsh(F_z) + F_pd = bool(jnp.all(eig > 0)) + cond = float(eig[-1] / eig[0]) if F_pd else float("inf") + print(f"\nGauss-Newton Fisher F_z PD: {F_pd}; eigenvalues = " + f"{[float(x) for x in eig]}; cond = {cond:.3e}") + if not F_pd: + print("STOP: F_z is not positive definite -- degenerate / saddle fit. " + "Reporting rather than masking.") + sys.exit(3) + + cov_theta = di.constrained_cov(F_z, _dtheta_dz(z_hat)) + sigma_theta = jnp.sqrt(jnp.diag(cov_theta)) + pulls = recovery_table(theta_hat, sigma_theta) + + # Realizability at theta_hat: rebuild Engine B, report f_min_j. + st_hat = _engine_b(theta_hat) + f_min_hat = st_hat.f_min_j + min_fmin = float(jnp.min(f_min_hat)) + print(f"\nRealizability at theta_hat: f_min_j = " + f"{[float(x) for x in f_min_hat]}, min = {min_fmin:.4e}") + + # GATES. + print("\n" + "=" * 72) + print("GATES") + print("=" * 72) + recovery_ok = all(abs(p) < 3.0 for p in pulls) + names = ("t", "r_a", "r_h") + for nm, p in zip(names, pulls): + print(f" 3-sigma {nm}: |pull| = {abs(p):.3f} " + f"{'<' if abs(p) < 3.0 else '>='} 3 " + f"({'PASS' if abs(p) < 3.0 else 'FAIL'})") + realiz_ok = min_fmin >= -1e-3 + print(f" realizability: min(f_min_j) = {min_fmin:.4e} " + f"{'>=' if realiz_ok else '<'} -1e-3 " + f"({'PASS' if realiz_ok else 'FAIL'})") + occ_ok = all(nb >= MIN_BINS_PER_CMP for nb in resolved) + for j, nb in enumerate(resolved): + print(f" occupancy comp {j}: {nb} resolved sigma bins " + f"{'>=' if nb >= MIN_BINS_PER_CMP else '<'} {MIN_BINS_PER_CMP} " + f"({'PASS' if nb >= MIN_BINS_PER_CMP else 'FAIL'})") + print(f" plateau: {'PASS' if plat_ok else 'FAIL'}") + print(f" recovery (3-sigma, all params): {'PASS' if recovery_ok else 'FAIL'}") + + if not occ_ok: + print("\nSTOP: a component has too few resolved sigma bins to constrain " + "the fit. NOT weakening n_min or the bin count -- report.") + sys.exit(4) # stop BEFORE the figure (match the sc/budget early exits) + + # Figure. + out_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), + "validation", "plots") + os.makedirs(out_dir, exist_ok=True) + rho_rarh, rho_tra = make_figure(data, theta_hat, sigma_theta, cov_theta, out_dir) + print(f"Fisher correlations: rho(r_a, r_h) = {rho_rarh:+.3f}, " + f"rho(t, r_a) = {rho_tra:+.3f}") + + all_ok = (sc_ok and grad_finite and grad_nonzero and budget_ok + and plat_ok and occ_ok and recovery_ok and F_pd and realiz_ok) + print("\n" + "=" * 72) + print(f"OVERALL {'ALL PASS' if all_ok else 'FAIL'}") + print("=" * 72) + if not all_ok and not recovery_ok: + print("\nNOTE: the 3-sigma recovery gate is REAL. A >3-sigma miss is a " + "PHYSICS finding -- do NOT widen the gate. Report the table above.") + sys.exit(0 if all_ok else 1) + + +if __name__ == "__main__": + main() diff --git a/src/progenax/__init__.py b/src/progenax/__init__.py index 9f47485c..2946680c 100644 --- a/src/progenax/__init__.py +++ b/src/progenax/__init__.py @@ -25,7 +25,7 @@ ) from .profiles import ( PlummerProfile, KingProfile, MichieProfile, EFFProfile, - LIMEPYProfile, + LIMEPYProfile, UniformSphereProfile, solve_king_profile, solve_michie_profile, solve_limepy_profile, solve_multimass_limepy, find_alpha_for_masses, ) @@ -33,6 +33,13 @@ PlummerVelocityDF, KingVelocityDF, MichieVelocityDF, EFFVelocityDF, LIMEPYVelocityDF, ) +# IMF / binary statistics (audit R7: documented but not previously root-exported) +from .imf import ( + PowerLawIMF, ChabrierIMF, Maschberger, TruncatedIMF, BinaryIMF, + FlatMassRatio, PowerLawMassRatio, TwinPeakedMassRatio, + MoeDiStefano2017, MoeDiStefano2017Full, MoePeriod, MoeJointOrbit, + ConstantBinaryFraction, MassDependentBinaryFraction, +) from .builders import ( ICResult, Systems, @@ -123,12 +130,28 @@ "MichieProfile", "EFFProfile", "LIMEPYProfile", + "UniformSphereProfile", "MultiComponentCluster", "solve_king_profile", "solve_michie_profile", "solve_limepy_profile", "solve_multimass_limepy", "find_alpha_for_masses", + # IMF / binary statistics + "PowerLawIMF", + "ChabrierIMF", + "Maschberger", + "TruncatedIMF", + "BinaryIMF", + "FlatMassRatio", + "PowerLawMassRatio", + "TwinPeakedMassRatio", + "MoeDiStefano2017", + "MoeDiStefano2017Full", + "MoePeriod", + "MoeJointOrbit", + "ConstantBinaryFraction", + "MassDependentBinaryFraction", # Velocity distribution functions "PlummerVelocityDF", "KingVelocityDF", diff --git a/src/progenax/binaries/diagnostics.py b/src/progenax/binaries/diagnostics.py index 0668ea9b..7d5878fa 100644 --- a/src/progenax/binaries/diagnostics.py +++ b/src/progenax/binaries/diagnostics.py @@ -206,13 +206,17 @@ class BinaryEnergyBudget(NamedTuple): Attributes: E_internal: total internal orbital energy of the primordial binaries (Σ ``relative_energy``; < 0 if bound). The separate "reservoir" that the - global virial scaling (``Q``) does NOT touch. + global virial scaling (``Q``) does NOT touch. NB the ``softening`` passed + to the bound-pair finders does NOT soften this internal binding energy — + it only regularizes the inter-system potential (audit S18). T_com, W_com: bulk kinetic / gravitational energy of the *system COMs* — the scale the cluster is virialized on. Q_com: ``T_com / |W_com|`` — the virial ratio the cluster was scaled to (≈ the ``Q`` passed to ``build_binary_cluster``). Q_resolved: ``T / |W|`` on the *resolved* stars — the naive ratio that mixes - the cluster and internal-binary scales (inflated by internal binary KE). + the cluster and internal-binary scales. The deep internal binary binding + dominates |W|, so Q_resolved is DEFLATED below the cluster Q (measured + ≈ 0.31 vs the 0.5 the cluster was scaled to), NOT inflated (audit S10). n_binaries: number of primordial binaries (two positive-mass members). """ diff --git a/src/progenax/builders.py b/src/progenax/builders.py index 90731ccf..6f0338f8 100644 --- a/src/progenax/builders.py +++ b/src/progenax/builders.py @@ -21,7 +21,11 @@ # Single canonical energy implementation lives in dynamics.virial; re-export it # here so the public API (progenax.compute_*_energy) and virial_scale share one # gradient-safe source of truth (Batch 0, F1+F2). -from .dynamics.virial import compute_kinetic_energy, compute_potential_energy +from .dynamics.virial import ( + compute_kinetic_energy, + compute_potential_energy, + rescale_velocities_to_virial, +) from .binaries import resolve_binary_components # Seconds in one (SI) day — exact; used to convert sampled periods (days) into the @@ -154,14 +158,30 @@ class ICResult(eqx.Module): def compute_stellar_radii(masses: Float[Array, "N"]) -> Float[Array, "N"]: """ - Estimate stellar radii from mass (main-sequence + brown dwarfs). - - Returns radii in SOLAR RADII (R☉). - - R/R☉ relations (3 regimes): - - M > 1 M☉: R ∝ M^0.8 (massive stars) - - 0.08 ≤ M ≤ 1 M☉: R ∝ M^0.57 (low-mass main sequence) - - M < 0.08 M☉: R ∝ M^0.08 (brown dwarfs, ~0.1 R☉) + Main-sequence stellar radii from mass, in SOLAR RADII. + + Demircan & Kahraman (1991), Ap&SS 181, 313 (their Eqs. 5-6): + R/Rsun = 1.06 (M/Msun)^0.945 for 0.08 <= M < 1.66 Msun + R/Rsun = 1.33 (M/Msun)^0.555 for M >= 1.66 Msun + Brown dwarfs (M < 0.08): R ~ 0.1 Rsun plateau (electron degeneracy), + R = 0.1 (M/0.08)^0.08. The D&K branches meet at 1.66 Msun with their + own ~3.5% fit discontinuity; the low-mass branch meets the BD plateau + nearly continuously (1.06 * 0.08^0.945 = 0.0975). + + Audit R6: the previous exponents were INVERTED vs MS homology (M^0.8 + above 1 Msun instead of below), giving 10 Msun -> 6.3 Rsun (ZAMS ~4.8) + and 0.2 Msun -> 0.40 Rsun (observed ~0.23), with a factor-2.4 jump at + the hydrogen-burning limit. + + Used for collision radii in downstream N-body; ZAMS values, no evolution. + + .. note:: + TEMPORARY single-relation stand-in. This will be REPLACED by the + ``startrax`` package once it lands: first Tout et al. (1996) ZAMS radii + (metallicity-dependent rational-polynomial fits), then MIST and the + Hurley+2000 SSE tracks (mass- AND age-dependent radii with real stellar + evolution). The D&K91 fit here is a static main-sequence approximation + chosen only to be cited and correct in the ZAMS limit until then. Args: masses: Particle masses (N,) [M☉] @@ -171,17 +191,17 @@ def compute_stellar_radii(masses: Float[Array, "N"]) -> Float[Array, "N"]: """ def radius_high_mass(m): - return jnp.power(m, 0.8) + return 1.33 * jnp.power(m, 0.555) def radius_low_mass(m): - return jnp.power(m, 0.57) + return 1.06 * jnp.power(m, 0.945) def radius_brown_dwarf(m): return 0.1 * jnp.power(m / 0.08, 0.08) - radii = jax.vmap( + return jax.vmap( lambda m: jax.lax.cond( - m > 1.0, + m >= 1.66, radius_high_mass, lambda mv: jax.lax.cond( mv >= 0.08, radius_low_mass, radius_brown_dwarf, mv @@ -190,8 +210,6 @@ def radius_brown_dwarf(m): ) )(masses) - return radii - def to_com_frame( positions: Float[Array, "N 3"], @@ -245,18 +263,18 @@ def virial_scale( Returns: Scaled velocities + Cold input (T=0) raises for concrete inputs / yields NaN under tracing — + see the delegate's docstring. + References: Goodwin & Whitworth (2004) A&A 413, 929 - Sub-virial clusters Baumgardt & Kroupa (2007) MNRAS 380, 1589 - Cluster dissolution """ - T = compute_kinetic_energy(velocities, masses) - V = compute_potential_energy(positions, masses, G=G, softening=softening) - - # Q = T / |V| (NOT 2T / |V|) - Q_current = T / jnp.abs(V) - scale = jnp.sqrt(Q_target / Q_current) - - return velocities * scale + # Single implementation lives in dynamics.virial (audit J5 dedupe); this is + # the builders-level convenience signature (Q_target before G). + return rescale_velocities_to_virial( + positions, velocities, masses, G=G, target_Q=Q_target, softening=softening + ) def build_spatial_ic( diff --git a/src/progenax/cluster/__init__.py b/src/progenax/cluster/__init__.py index c3623e7d..c3ee863e 100644 --- a/src/progenax/cluster/__init__.py +++ b/src/progenax/cluster/__init__.py @@ -43,7 +43,8 @@ ALPHA_LARSON, ) -# Turbulence physics (stays in core; used by EnvironmentIMF) +# Turbulence physics (stays in core; used by the environment-dependent IMF — +# BirthEnvironment / env_to_imf_params, NOT a phantom "EnvironmentIMF" class) from progenax.cluster.turbulence import ( sigma_ln_rho_from_mach, spectral_slope_from_mach, diff --git a/src/progenax/cluster/mass_segregation.py b/src/progenax/cluster/mass_segregation.py index 6e225dec..707cc08c 100644 --- a/src/progenax/cluster/mass_segregation.py +++ b/src/progenax/cluster/mass_segregation.py @@ -57,6 +57,13 @@ def energy_sorted_segregation( equilibrium family (MultiComponentCluster.from_mass_segregation), not a blend of this function's output. + .. warning:: + Reassigning masses by energy rank changes the mass-weighted density, so it + no longer matches the parent profile the orbit pool was drawn from — the + self-consistent potential of the segregated system DIFFERS, and the output is + not exactly in equilibrium (audit S6). Callers should finalize with a global + virial rescale (``virial_scale``), as the validation suite does. + Algorithm (Baumgardt+2008 Appendix, McLuster implementation): 1. Sort masses descending → m_sorted[i] for mass rank i 2. Cumulative-mass coordinate M_cum_norm[i] = sum(m_sorted[0:i+1]) / M_total, diff --git a/src/progenax/diagnostics/mass_segregation.py b/src/progenax/diagnostics/mass_segregation.py index 1352f3f3..1d756bf6 100644 --- a/src/progenax/diagnostics/mass_segregation.py +++ b/src/progenax/diagnostics/mass_segregation.py @@ -13,11 +13,28 @@ Allison et al. (2009), MNRAS 395, 1449 """ +from __future__ import annotations + from typing import Tuple -import numpy as np -from scipy.sparse.csgraph import minimum_spanning_tree -from scipy.spatial.distance import pdist, squareform + +def _diagnostics_deps(): + """Lazy numpy+scipy import for the exact (non-differentiable) mass-segregation + path. progenax core is JAX-native; numpy/scipy are only needed by these + scipy-backed estimators, so they are an OPTIONAL ``[diagnostics]`` extra + (audit R9). ``import progenax.diagnostics`` stays cheap; calling these + functions without the extra raises an actionable ImportError.""" + try: + import numpy as np + from scipy.sparse.csgraph import minimum_spanning_tree + from scipy.spatial.distance import pdist, squareform + except ImportError as e: + raise ImportError( + "progenax.diagnostics.compute_lambda_msr needs numpy+scipy — the " + "exact, non-differentiable diagnostics path. Install them with: " + "uv pip install 'progenax[diagnostics]'" + ) from e + return np, minimum_spanning_tree, pdist, squareform def compute_lambda_msr( @@ -78,6 +95,7 @@ def compute_lambda_msr( Allison et al. (2009), ApJ 700, L99 — application (note: L99 Eq. 1 is the Spitzer t_seg relation, NOT Λ_MSR; verified against the held PDF 2026-06-08). """ + np, minimum_spanning_tree, pdist, squareform = _diagnostics_deps() rng = np.random.default_rng(seed) N = len(masses) @@ -123,6 +141,8 @@ def _compute_mst_length(positions: np.ndarray) -> float: if len(positions) < 2: return 0.0 + _, minimum_spanning_tree, pdist, squareform = _diagnostics_deps() + # Compute pairwise distance matrix dist_matrix = squareform(pdist(positions)) diff --git a/src/progenax/diagnostics/q_approx.py b/src/progenax/diagnostics/q_approx.py index 356b967d..f2db0bd4 100644 --- a/src/progenax/diagnostics/q_approx.py +++ b/src/progenax/diagnostics/q_approx.py @@ -236,7 +236,12 @@ def subsampled(): diffs = xy[idx_i] - xy[idx_j] return jnp.mean(jnp.sqrt(jnp.sum(diffs ** 2, axis=-1) + 1e-12)) / R_cluster - return jax.lax.cond(n_pairs <= max_pairs, lambda _: exact(), lambda _: subsampled(), None) + # n_pairs and max_pairs are STATIC Python ints; a Python if avoids tracing + # and compiling the O(N^2) exact() branch when subsampling (audit J6 — saves + # both compile time and the (N,N) allocation for large N). + if n_pairs <= max_pairs: + return exact() + return subsampled() def q_approx( @@ -266,13 +271,12 @@ def q_approx( elif method == "fast": return q_approx_fast(positions, project_to_2d, calibration=calibration, **kwargs) elif method == "auto": - # Use naive for small N (JIT overhead dominates), fast for large N - return jax.lax.cond( - N > 1000, - lambda _: q_approx_fast(positions, project_to_2d, calibration=calibration, **kwargs), - lambda _: q_approx_naive(positions, project_to_2d, calibration), - operand=None, - ) + # N is a STATIC Python int, so a lax.cond would compile BOTH branches + # every shape (~0.4 s + the fast branch's extra memory). A Python if + # selects exactly one — identical semantics, half the compile (audit J6). + if N > 1000: + return q_approx_fast(positions, project_to_2d, calibration=calibration, **kwargs) + return q_approx_naive(positions, project_to_2d, calibration) else: raise ValueError(f"method must be 'auto', 'naive', or 'fast'") diff --git a/src/progenax/diagnostics/substructure.py b/src/progenax/diagnostics/substructure.py index 1d4117d1..ef7359b8 100644 --- a/src/progenax/diagnostics/substructure.py +++ b/src/progenax/diagnostics/substructure.py @@ -17,9 +17,27 @@ Küpper et al. (2011), MNRAS 417, 2300 """ -import numpy as np -from scipy.sparse.csgraph import minimum_spanning_tree -from scipy.spatial.distance import pdist, squareform +from __future__ import annotations + + +def _diagnostics_deps(): + """Lazy numpy+scipy import for the exact (non-differentiable) substructure + path. progenax core is JAX-native; numpy/scipy are only needed by these + scipy-backed estimators, so they are an OPTIONAL ``[diagnostics]`` extra + (audit R9). ``import progenax.diagnostics`` stays cheap; calling these + functions without the extra raises an actionable ImportError.""" + try: + import numpy as np + from scipy.sparse.csgraph import minimum_spanning_tree + from scipy.spatial.distance import pdist, squareform + except ImportError as e: + raise ImportError( + "progenax.diagnostics substructure estimators (compute_q_parameter, " + "compute_azimuthal_variation) need numpy+scipy — the exact, " + "non-differentiable diagnostics path. Install them with: " + "uv pip install 'progenax[diagnostics]'" + ) from e + return np, minimum_spanning_tree, pdist, squareform def compute_q_parameter(positions: np.ndarray) -> float: @@ -80,6 +98,8 @@ def compute_q_parameter(positions: np.ndarray) -> float: References: Cartwright & Whitworth (2004), MNRAS 348, 589 """ + np, minimum_spanning_tree, pdist, squareform = _diagnostics_deps() + # 1. Project to 2D (CW04 always uses 2D projected positions) if positions.shape[1] == 3: xy = positions[:, :2] @@ -169,6 +189,8 @@ def compute_azimuthal_variation( References: Küpper et al. (2011), MNRAS 417, 2300 """ + np = _diagnostics_deps()[0] + # Project to x-y plane and compute azimuthal angle phi = np.arctan2(positions[:, 1], positions[:, 0]) diff --git a/src/progenax/dynamics/virial.py b/src/progenax/dynamics/virial.py index 394e1384..0188513c 100644 --- a/src/progenax/dynamics/virial.py +++ b/src/progenax/dynamics/virial.py @@ -196,9 +196,16 @@ def per_group_virial_ratio( for ``1/r`` gravity, a single all-ones group reproduces the global virial ratio exactly (``W = V``), and the per-group ``W_j`` partition the global virial. - Positions enter only through differences (origin-independent); velocities are - measured relative to the mass-weighted mean (bulk motion removed) so a moving COM - does not inflate ``T_j``. + Velocities are measured relative to the mass-weighted mean (bulk motion removed) + so a moving COM does not inflate ``T_j``. + + .. warning:: + Per-group ``W_j = sum_{i in j} m_i r_i . a_i`` is ORIGIN-DEPENDENT for a + proper subgroup: under r -> r + d it shifts by ``d . sum_{i in j} m_i a_i``, + which is nonzero because a subgroup feels a net force from the rest of the + cluster (only the GLOBAL sum is origin-independent, ``sum_all m_i a_i = 0``). + Callers must therefore pass positions PRE-CENTERED on the cluster COM (all + in-tree callers do); an off-center origin biases the per-group Q_j (audit S16). Args: positions: Particle positions, shape (N, 3). @@ -226,6 +233,16 @@ def per_group_virial_ratio( return T_j / (jnp.abs(W_j) + 1e-30) +def _is_concrete(x) -> bool: + """True iff x is a concrete value (codebase idiom; cf. king.py auto-sizing).""" + try: + float(x) + return True + except (jax.errors.ConcretizationTypeError, jax.errors.TracerArrayConversionError, + TypeError): + return False + + def rescale_velocities_to_virial( positions: Float[Array, "N 3"], velocities: Float[Array, "N 3"], @@ -236,6 +253,14 @@ def rescale_velocities_to_virial( ) -> Float[Array, "N 3"]: """Rescale velocities to achieve target virial ratio Q = T/|V|. + This is the single implementation of virial velocity rescaling; + ``builders.virial_scale`` delegates here (audit J5 dedupe). + + Cold input (T=0 -> Q_current=0) makes the rescale undefined (0/0). A + CONCRETE T=0 refuses loudly (ValueError); a TRACED T=0 yields NaN — the + honest sentinel, since a traced scalar cannot gate a Python raise (audit J5; + the previous +1e-10 hack silently mapped cold input to ~0 velocities). + Args: positions: Particle positions, shape (N, 3) velocities: Particle velocities, shape (N, 3) @@ -248,5 +273,13 @@ def rescale_velocities_to_virial( Rescaled velocities with Q = target_Q """ Q_current = compute_virial_ratio(positions, velocities, masses, G=G, softening=softening) - scale = jnp.sqrt(target_Q / (Q_current + 1e-10)) + if _is_concrete(Q_current) and float(Q_current) <= 0.0: + raise ValueError( + "cannot rescale from zero kinetic energy (T=0): velocities are all " + "zero, so the virial ratio is undefined. Provide non-zero velocities " + "(e.g. sample from a velocity DF) before virial scaling." + ) + # double-where keeps the untaken (Q<=0) branch finite; traced T=0 -> NaN. + Q_safe = jnp.where(Q_current > 0.0, Q_current, 1.0) + scale = jnp.where(Q_current > 0.0, jnp.sqrt(target_Q / Q_safe), jnp.nan) return velocities * scale diff --git a/src/progenax/imf/base.py b/src/progenax/imf/base.py index 347f90cf..5c79b8a9 100644 --- a/src/progenax/imf/base.py +++ b/src/progenax/imf/base.py @@ -235,13 +235,36 @@ def sample_fixed_n( max_steps: int = 50, ) -> Float[Array, "n"]: """ - Fixed-N mode: exactly n masses summing to m_total via quantile stretching. - - Uses stratified quantile sampling with stretching factor q* solved to - hit target total mass. Differentiable w.r.t. m_total. + Fixed-N mode: n masses whose total APPROXIMATES m_total via quantile + stretching (the realized total carries the residual of the final random + jitter — it is not exact). + + Uses stratified quantile sampling with a stretch factor q* (<= 1) solved + to hit the target total mass. Because the stretch is one-sided (q* <= 1) + AND stratification truncates the heavy tail, the achievable total is + capped at sum(ppf((i+0.5)/n)) — a ceiling that sits a few percent BELOW + n*E[m]. A target ABOVE that ceiling is UNREACHABLE: eager (concrete) + inputs now raise ValueError instead of silently undershooting (audit R5: + target 500 Msun, n=1000 used to return ~349 Msun with no warning). Traced + m_total cannot be checked under jit/grad — the caller owns reachability. + Differentiable w.r.t. m_total. """ u_base = (jnp.arange(n) + 0.5) / n + # Reachability guard (eager inputs only). The one-sided stretch caps the + # total at sum(ppf(u_base)); a target above it silently undershot before. + try: + m_ceiling = float(jnp.sum(self.ppf(u_base))) + if float(m_total) > m_ceiling: + raise ValueError( + f"m_total={float(m_total):.6g} Msun is unreachable for n={n}: " + f"the stratified-quantile ceiling is {m_ceiling:.6g} Msun. " + f"Increase n, lower m_total, or use sample_m_total()." + ) + except (jax.errors.ConcretizationTypeError, + jax.errors.TracerArrayConversionError): + pass # traced m_total / params: caller owns reachability + q_star = self._solve_q_for_m_total(u_base, m_total, max_steps) u_random = jax.random.uniform(key, (n,)) diff --git a/src/progenax/imf/binary/moe_di_stefano.py b/src/progenax/imf/binary/moe_di_stefano.py index 8080546c..49d71496 100644 --- a/src/progenax/imf/binary/moe_di_stefano.py +++ b/src/progenax/imf/binary/moe_di_stefano.py @@ -230,14 +230,20 @@ def _pl_integral(a, b, gamma): class MoeDiStefano2017Full(eqx.Module): """Faithful two-slope, period-dependent Moe & Di Stefano (2017) mass-ratio model. - The mass-ratio pdf is (their §9.1, Eq. 2 + the worked twin example, p.5): + The mass-ratio pdf is a two-slope power law plus a twin excess block, + jointly normalized over [q_min, 1.0]: - p_q(q | M1, P) = (1 - F_twin) * p_2slope(q) + F_twin * Uniform[0.95, 1.0] + p_q(q | M1, P) ∝ p_2slope(q) + [twin block on [0.95, 1.0]] where p_2slope ∝ q^γsmallq on [q_min, 0.3] then q^γlargeq on [0.3, 1.0] - (continuous at q=0.3), normalized over [q_min, 1.0]; F_twin is the EXCESS twin - weight (the fraction of systems that are pure twins, not the total q>0.95 - fraction). γsmallq, γlargeq, F_twin are bilinearly interpolated (clamped) over + (continuous at q=0.3). F_twin is the EXCESS-twin fraction of the q > 0.3 + population (MD17 p.5, Fig. 2 — NOT the fraction of all q > q_min companions, + and NOT the total q > 0.95 fraction). To realize that convention the twin + block carries unnormalized mass ft/(1-ft)·I_B (I_B = power-law mass on + [0.3, 1]), so twin/(twin + I_B) = ft exactly (see `_components`). The pre-fix + code mixed ft against the whole [q_min, 1] population, over-weighting twins + by ~22% at solar logP=1 (realized 0.367 vs Table-13 0.300) — audit R3. + γsmallq, γlargeq, F_twin are bilinearly interpolated (clamped) over Table 13 (verified against the PDF p.52). η(P,M1) for eccentricity is handled by `progenax.binaries.MoeEccentricity` (which reproduces Table 13's η). @@ -276,17 +282,38 @@ def _two_slope(self, periods, masses): I_B = C * _pl_integral(self.q_break, 1.0, gl) return gs, gl, C, I_A, I_B - def pdf(self, q, masses, periods): - """Conditional mass-ratio pdf p(q | M1, P), normalized on [q_min, 1].""" + def _components(self, q, masses, periods): + """Normalized (power-law, twin-excess) mixture components. + + Paper convention (MD17 p.5, Fig. 2): F_twin is the excess-twin fraction + of q > 0.3 companions. The unnormalized twin block therefore carries + mass ft/(1-ft) * I_B (I_B = power-law mass on [q_break, 1]), so that + twin/(twin + I_B) = ft exactly after joint normalization. + """ gs, gl, C, I_A, I_B = self._two_slope(periods, masses) ft = self.f_twin(periods, masses) - Z = I_A + I_B p_lo = jnp.power(q, gs) p_hi = C * jnp.power(q, gl) - p_pl = jnp.where(q < self.q_break, p_lo, p_hi) / Z - twin = jnp.where((q >= 0.95) & (q <= 1.0), 1.0 / 0.05, 0.0) - p = (1.0 - ft) * p_pl + ft * twin - return jnp.where((q >= self.q_min) & (q <= 1.0), p, 0.0) + p_pl_unnorm = jnp.where(q < self.q_break, p_lo, p_hi) + ft_safe = jnp.minimum(ft, 0.95) # Table 13 max is ~0.3; guard 1/(1-ft) + twin_mass = ft_safe / (1.0 - ft_safe) * I_B + twin_unnorm = twin_mass * jnp.where( + (q >= 0.95) & (q <= 1.0), 1.0 / 0.05, 0.0 + ) + Z_tot = I_A + I_B + twin_mass + in_range = (q >= self.q_min) & (q <= 1.0) + p_pl = jnp.where(in_range, p_pl_unnorm / Z_tot, 0.0) + p_twin = jnp.where(in_range, twin_unnorm / Z_tot, 0.0) + return p_pl, p_twin + + def pdf(self, q, masses, periods): + """Conditional mass-ratio pdf p(q | M1, P), normalized on [q_min, 1]. + + F_twin follows the paper convention: excess-twin fraction of the + q > 0.3 population (NOT of all q > q_min companions — audit R3). + """ + p_pl, p_twin = self._components(q, masses, periods) + return p_pl + p_twin def sample(self, key: PRNGKeyArray, masses, periods): """Sample q | (M1 [Msun], P [days]) from the two-slope + twin mixture. diff --git a/src/progenax/imf/power_law.py b/src/progenax/imf/power_law.py index eaa47246..7c54e46e 100644 --- a/src/progenax/imf/power_law.py +++ b/src/progenax/imf/power_law.py @@ -89,10 +89,14 @@ def segment_integral(i): a = alphas[i] lo, hi = bounds[i], bounds[i + 1] e = 1.0 - a + # exp_safe double-where: keep BOTH branches finite so the VJP is + # finite at exactly alpha=1 (bare where backprops 0*NaN — audit R10). + # Same pattern as imf/differentiable.py power_integral. + e_safe = jnp.where(jnp.abs(e) < 1e-12, 1.0, e) return jnp.where( jnp.abs(e) < 1e-12, cont[i] * jnp.log(hi / lo), - cont[i] * (hi**e - lo**e) / e, + cont[i] * (hi**e_safe - lo**e_safe) / e_safe, ) integrals = jax.vmap(segment_integral)(jnp.arange(n_segments)) @@ -197,10 +201,12 @@ def sum_before_idx(idx_val): lo = bounds[idx] a = alphas[idx] e = 1.0 - a + # exp_safe double-where (audit R10): finite VJP at exactly alpha=1. + e_safe = jnp.where(jnp.abs(e) < 1e-12, 1.0, e) within_segment = jnp.where( jnp.abs(e) < 1e-12, self._continuity_factors[idx] * jnp.log(m / lo), - self._continuity_factors[idx] * (m**e - lo**e) / e, + self._continuity_factors[idx] * (m**e_safe - lo**e_safe) / e_safe, ) return cum_before + within_segment @@ -245,11 +251,15 @@ def ppf(self, u: Float[Array, "..."]) -> Float[Array, "..."]: hi = bounds[idx + 1] a = alphas[idx] e = 1.0 - a + # exp_safe double-where (audit R10): e_safe substitutes for e in EVERY + # power/division of the inverse-CDF branch (including the **(1/e) outer + # exponent), so the VJP is finite at exactly alpha=1. + e_safe = jnp.where(jnp.abs(e) < 1e-12, 1.0, e) return jnp.where( jnp.abs(e) < 1e-12, lo * jnp.exp(ur * jnp.log(hi / lo)), - (ur * (hi**e - lo**e) + lo**e) ** (1.0 / e), + (ur * (hi**e_safe - lo**e_safe) + lo**e_safe) ** (1.0 / e_safe), ) def sample(self, key: PRNGKeyArray, n: int) -> Float[Array, "n"]: @@ -273,10 +283,14 @@ def first_moment(i): a = alphas[i] lo, hi = bounds[i], bounds[i + 1] e = 2.0 - a # antiderivative exponent of m·m^(-a) = m^(1-a) + # exp_safe double-where (audit R10): here the removable singularity + # is at alpha=2 (e=0); the alpha=1 NaN in mean_mass flows through Z + # (segment_integral, fixed above), but guard this branch too. + e_safe = jnp.where(jnp.abs(e) < 1e-12, 1.0, e) return jnp.where( jnp.abs(e) < 1e-12, cont[i] * jnp.log(hi / lo), - cont[i] * (hi**e - lo**e) / e, + cont[i] * (hi**e_safe - lo**e_safe) / e_safe, ) moments = jax.vmap(first_moment)(jnp.arange(len(self.exponents))) diff --git a/src/progenax/kinematics/api.py b/src/progenax/kinematics/api.py index f53dc326..3981167e 100644 --- a/src/progenax/kinematics/api.py +++ b/src/progenax/kinematics/api.py @@ -34,7 +34,7 @@ ... ) """ -from typing import Optional, Protocol, runtime_checkable +from typing import Optional import equinox as eqx import jax @@ -48,35 +48,10 @@ apply_differential_rotation, ) from progenax.dynamics.virial import rescale_velocities_to_virial - - -@runtime_checkable -class VelocityDF(Protocol): - """Protocol for velocity distribution functions. - - All velocity DFs must implement this interface for use with the - kinematics API pipeline. - """ - - def sample_velocities( - self, - positions: Float[Array, "N 3"], - masses: Float[Array, "N"], - key: PRNGKeyArray, - G: float | None = None, - ) -> Float[Array, "N 3"]: - """Sample velocities from the distribution function. - - Args: - positions: Particle positions (N, 3) [length units] - masses: Particle masses (N,) [mass units] - key: JAX random key - G: Gravitational constant. If None, uses default. - - Returns: - Cartesian velocities (N, 3) [velocity units] - """ - ... +# Single source of truth for the VelocityDF protocol (audit A1): re-exported +# here so progenax.kinematics.api.VelocityDF stays importable, but defined once +# in progenax.protocols to prevent structural drift. +from progenax.protocols import VelocityDF class RotationParams(eqx.Module): @@ -224,7 +199,12 @@ def sample_velocities_pipeline( # Stage 3: Optional virial rescaling. target_Q=None keeps the DF's native # equilibrium (Plummer/King/EFF are already sampled in detailed equilibrium); - # pass a float only to deliberately force a virial ratio. + # pass a float only to deliberately force an overall virial ratio. + # + # NOTE (audit S3): this is an ISOTROPIC speed rescale. It does NOT restore + # stationarity after a rotation overlay (Stage 2) — it cannot remove the + # injected net angular momentum L_z, so the result stays non-stationary even + # at Q=0.5. Rescaling-to-0.5 is not a fix for the rotation-broken equilibrium. if model.target_Q is not None: v = rescale_velocities_to_virial( positions=positions, diff --git a/src/progenax/kinematics/eff_df.py b/src/progenax/kinematics/eff_df.py index 6d0fd1c6..92206905 100644 --- a/src/progenax/kinematics/eff_df.py +++ b/src/progenax/kinematics/eff_df.py @@ -75,6 +75,16 @@ class EFFVelocityDF(eqx.Module): velocity scale is fixed self-consistently from (G, M_total, a, gamma, r_t), so no external virial rescale is needed. + .. warning:: + The DEFAULT (gamma=3, sharp truncation) is KNOWN ~8% sub-virial by + construction: a truncated rho with rho(r_t)>0 cannot be represented by an + f(E), so the sampled cluster is only approximately stationary (measured + Q = T/|V| ~ 0.458, not 0.5; pinned by test_eff_physics + test_gamma3_default_subvirial_offset_is_pinned). This is a documented + limitation, NOT an inversion error. For mild truncation (gamma>=5) the + offset shrinks to ~1% (Q ~ 0.495); for a strict lowered-DF equilibrium + use the King/Michie models instead. + With ``anisotropy_radius`` (r_a) set, the DF is the Osipkov-Merritt radially anisotropic model for the same EFF density: f = f(Q), Q = E + J^2/2r_a^2, built by Eddington inversion of the augmented density rho_Q = (1 + r^2/r_a^2) rho (Merritt diff --git a/src/progenax/kinematics/king_df.py b/src/progenax/kinematics/king_df.py index c7372d15..96c7483f 100644 --- a/src/progenax/kinematics/king_df.py +++ b/src/progenax/kinematics/king_df.py @@ -87,7 +87,6 @@ class KingVelocityDF(eqx.Module): Attributes: W0: King concentration parameter (dimensionless central potential) r_c: Core radius [length units] (the King core radius) - r_t: Tidal radius [length units] xi_grid, psi_grid: ODE solution of the King model (xi = r/r_c, psi(xi)) speed_method: static, "table" (default) or "quadrature" (exact oracle) @@ -98,7 +97,6 @@ class KingVelocityDF(eqx.Module): W0: Float[Array, ""] r_c: Float[Array, ""] - r_t: Float[Array, ""] xi_grid: Float[Array, "n_ode"] psi_grid: Float[Array, "n_ode"] speed_table: SpeedCDFTable | None @@ -108,7 +106,6 @@ def __init__( self, W0: float = 5.0, r_c: float = 1.0, - r_t: float = 10.0, xi_max: float | None = None, n_ode_points: int | None = None, speed_method: str = "table", @@ -120,7 +117,6 @@ def __init__( self.speed_method = speed_method self.W0 = jnp.asarray(W0) self.r_c = jnp.asarray(r_c) - self.r_t = jnp.asarray(r_t) # Auto-size the ODE domain from W0 (matches KingProfile.from_W0_rc) so the # matched DF stays self-consistent for high-concentration models. auto_xi_max, auto_n_points = _auto_ode_domain(W0) diff --git a/src/progenax/kinematics/rotation.py b/src/progenax/kinematics/rotation.py index bad2485c..1a94923f 100644 --- a/src/progenax/kinematics/rotation.py +++ b/src/progenax/kinematics/rotation.py @@ -2,15 +2,42 @@ Adds streaming rotation to velocity distributions. +.. warning:: + These are ADDITIVE KINEMATIC OVERLAYS (``velocities + v_rotation``): they + inject kinetic energy and net angular momentum L_z, so the output is NOT a + stationary equilibrium — the virial ratio Q = T/|V| rises above 0.5 (audit + S3). They do not solve a rotating-equilibrium DF. An isotropic re-virial + (e.g. ``virial_scale``) does NOT restore stationarity (it rescales speeds + uniformly but cannot remove the rotational L_z). Use only when you want a + deliberately rotating, non-stationary IC. + References: Lynden-Bell (1960) MNRAS 120, 204 - Rotating stellar systems Binney & Tremaine (2008) "Galactic Dynamics" Section 4.8 """ +import jax import jax.numpy as jnp from jaxtyping import Array, Float +def _normalized_rotation_axis(axis: Float[Array, "3"]) -> Float[Array, "3"]: + """Unit rotation axis. A CONCRETE zero axis raises (it has no rotation + direction — the old `/max(mag,1e-30)` silently returned a 0 axis, a no-op, + which the stale 'result will be NaN' comment mis-described — audit S15). A + TRACED zero axis cannot be checked and falls through to the 0-axis no-op.""" + axis_mag = jnp.linalg.norm(axis) + try: + if float(axis_mag) == 0.0: + raise ValueError( + "rotation axis is the zero vector — it has no rotation " + "direction. Pass a nonzero axis (e.g. jnp.array([0., 0., 1.]))." + ) + except (jax.errors.ConcretizationTypeError, jax.errors.TracerArrayConversionError): + pass # traced axis: cannot check; falls through to the no-op below + return axis / jnp.maximum(axis_mag, 1e-30) + + def apply_solid_body_rotation( velocities: Float[Array, "N 3"], positions: Float[Array, "N 3"], @@ -24,6 +51,11 @@ def apply_solid_body_rotation( For solid body rotation, all particles rotate with the same angular velocity omega, giving v_phi = omega * R (cylindrical R). + .. warning:: + Kinematic overlay — injects kinetic energy and L_z. The result is NOT a + stationary equilibrium (Q = T/|V| rises above 0.5); see the module + caveat (audit S3). + Args: velocities: Input velocities (N, 3) positions: Particle positions (N, 3) @@ -44,10 +76,7 @@ def apply_solid_body_rotation( Reference: Binney & Tremaine (2008) Section 4.8 """ - # Normalize axis with epsilon safeguard - axis_mag = jnp.linalg.norm(axis) - # Use safe division - if axis is zero, result will be NaN which is appropriate - axis_norm = axis / jnp.maximum(axis_mag, 1e-30) + axis_norm = _normalized_rotation_axis(axis) # raises on a concrete zero axis # Rotation velocity: v_rot = omega x r = omega * (axis x r) omega_vec = omega * axis_norm # (3,) @@ -74,6 +103,10 @@ def apply_differential_rotation( - v_phi(R_peak) = v_peak (maximum rotation) - v_phi -> 0 as R -> infinity (decreasing at large R) + .. warning:: + Kinematic overlay — injects kinetic energy and L_z. The result is NOT a + stationary equilibrium (Q rises above 0.5); see the module caveat (S3). + Args: velocities: Input velocities (N, 3) positions: Particle positions (N, 3) @@ -91,9 +124,7 @@ def apply_differential_rotation( Lynden-Bell (1960), MNRAS 120, 204 is the classic reference for rotating stellar systems in general, not for this functional form. """ - # Normalize axis with epsilon safeguard - axis_mag = jnp.linalg.norm(axis) - axis_norm = axis / jnp.maximum(axis_mag, 1e-30) + axis_norm = _normalized_rotation_axis(axis) # raises on a concrete zero axis # Compute cylindrical radius R (distance from rotation axis) # R^2 = |r|^2 - (r . axis)^2 diff --git a/src/progenax/profiles/api.py b/src/progenax/profiles/api.py index 350e356e..4491bec8 100644 --- a/src/progenax/profiles/api.py +++ b/src/progenax/profiles/api.py @@ -281,11 +281,21 @@ def compute_profile_potential( rho_t = (1.0 + (rgrid / profile_instance.a) ** 2) ** ( -profile_instance.gamma / 2.0 ) - dr = rgrid[1] - rgrid[0] - - I2 = cumulative_trapezoid(rho_t * rgrid**2, dx=dr) # propto M( <1% + # (stretched). `compute_profile_potential` (profiles/api.py) reads this + # same _r_grid and is kept consistent with a matching non-uniform + # enclosed-mass integral. Smooth in r_t -> differentiable. + u_grid = jnp.linspace(0.0, 1.0, n_grid) + r_grid = r_t_arr * u_grid**2 # Compute density on grid: rho(r) = (1 + r^2/a^2)^(-gamma/2) rho_grid = jnp.power(1.0 + (r_grid / a_arr) ** 2, -gamma_arr / 2.0) @@ -90,13 +96,13 @@ def __init__( # Integrand: 4*pi*r^2*rho(r) integrand = 4.0 * jnp.pi * r_grid**2 * rho_grid - # Cumulative mass via the trapezoid rule (2nd-order). The old - # cumsum(integrand)*dr was a 1st-order left/right-Riemann sum mislabeled - # "trapezoid" and biased the sampled radial distribution (audit M5). - dr = r_grid[1] - r_grid[0] + # Cumulative mass via the NON-UNIFORM trapezoid rule (2nd-order): the + # sqrt-stretched grid has variable spacing, so each trapezoid is weighted + # by its own width diff(r_grid). (The old cumsum(integrand)*dr was a + # 1st-order Riemann sum mislabeled "trapezoid" — audit M5.) M_cum = jnp.concatenate([ jnp.zeros(1, dtype=integrand.dtype), - jnp.cumsum(0.5 * (integrand[1:] + integrand[:-1])) * dr, + jnp.cumsum(0.5 * (integrand[1:] + integrand[:-1]) * jnp.diff(r_grid)), ]) # Normalize to [0, 1] for CDF diff --git a/src/progenax/profiles/king.py b/src/progenax/profiles/king.py index 43318d9a..03ca7c54 100644 --- a/src/progenax/profiles/king.py +++ b/src/progenax/profiles/king.py @@ -26,13 +26,24 @@ class implementing the SpatialProfile protocol. """ import math +import warnings from typing import Tuple import diffrax import equinox as eqx import jax import jax.numpy as jnp -from jaxtyping import Array, Float, PRNGKeyArray +from jaxtyping import Array, Bool, Float, PRNGKeyArray + + +def _is_concrete(x) -> bool: + """True iff x is a concrete value (codebase idiom; cf. _auto_ode_domain).""" + try: + float(x) + return True + except (jax.errors.ConcretizationTypeError, jax.errors.TracerArrayConversionError, + TypeError): + return False # ============================================================================== @@ -328,6 +339,9 @@ class KingProfile(eqx.Module): psi_grid: Float[Array, "n_points"] _r_grid: Float[Array, "n_grid"] _cdf_grid: Float[Array, "n_grid"] + # Diagnostic (audit J4): True iff psi never crossed 0 within the ODE domain, + # so r_t was pinned to the grid boundary (a wrong tidal radius). Traced bool. + r_t_is_pinned: Bool[Array, ""] def __init__( self, @@ -355,8 +369,13 @@ def __init__( xi_grid_arr = jnp.asarray(xi_grid, dtype=jnp.float64) psi_grid_arr = jnp.asarray(psi_grid, dtype=jnp.float64) - # Build radial grid for CDF - r_grid = jnp.linspace(0.0, r_t_arr, n_grid) + # Build radial grid for CDF — sqrt-stretched (r = r_t * u^2): spacing + # dr ∝ sqrt(r) concentrates points in the core. A LINEAR grid leaves + # <10 points inside the core at W0 >= 9 (xi_t grows super-exponentially: + # 131 r_c at W0=9, 548 at W0=12), giving +18%..+270% core-mass errors + # (audit R4, measured). Smooth in r_t -> differentiable in W0/r_c. + u_grid = jnp.linspace(0.0, 1.0, n_grid) + r_grid = r_t_arr * u_grid**2 xi_grid_local = r_grid / r_c_arr # Compute density on grid via interpolation of ODE solution @@ -383,17 +402,41 @@ def __init__( # Integrand: 4*pi*r^2*rho(r) integrand = 4.0 * jnp.pi * r_grid**2 * rho_grid - # Cumulative mass via the trapezoid rule (2nd-order); a 1st-order Riemann - # sum would bias the sampled radial distribution. - dr = r_grid[1] - r_grid[0] + # Cumulative mass via the NON-UNIFORM trapezoid rule (2nd-order): the + # sqrt-stretched grid has variable spacing, so the per-interval width + # diff(r_grid) must weight each trapezoid (a single dr would mis-integrate). M_cum = jnp.concatenate([ jnp.zeros(1, dtype=integrand.dtype), - jnp.cumsum(0.5 * (integrand[1:] + integrand[:-1])) * dr, + jnp.cumsum(0.5 * (integrand[1:] + integrand[:-1]) * jnp.diff(r_grid)), ]) # Normalize to [0, 1] for CDF cdf_grid = M_cum / (M_cum[-1] + 1e-30) + # r_t pinning diagnostic (audit J4): if psi(xi) never reaches 0 within + # the solved domain there is no crossing, so _find_tidal_radius fell back + # to xi_grid[-1] (the boundary) — a wrong tidal radius. Traced bool. + r_t_is_pinned = jnp.logical_not(jnp.any(psi_grid_arr <= 0.0)) + + # r_t consistency (audit S1): the direct constructor accepts an arbitrary + # r_t. For concrete, non-pinned inputs, warn if r_t deviates from the + # c(W0) tidal radius r_c*xi_t by >5% — a non-self-consistent, non- + # equilibrium model. from_W0_rc derives r_t, so it never trips this. + if (_is_concrete(r_t_arr) and _is_concrete(W0_arr) + and not bool(r_t_is_pinned)): + xi_t = _find_tidal_radius(xi_grid_arr, psi_grid_arr) + r_t_consistent = float(r_c_arr * xi_t) + if abs(float(r_t_arr) - r_t_consistent) > 0.05 * r_t_consistent: + warnings.warn( + f"KingProfile r_t={float(r_t_arr):.4g} is inconsistent with the " + f"c(W0={float(W0_arr):.2f}) tidal radius r_c*xi_t=" + f"{r_t_consistent:.4g} (>5% deviation): this builds a NON-self-" + f"consistent, non-equilibrium King model. Use " + f"KingProfile.from_W0_rc(W0, r_c) to derive a consistent r_t.", + UserWarning, + stacklevel=2, + ) + # Store using object.__setattr__ (future-proof Equinox pattern) object.__setattr__(self, "W0", W0_arr) object.__setattr__(self, "r_c", r_c_arr) @@ -402,6 +445,7 @@ def __init__( object.__setattr__(self, "psi_grid", psi_grid_arr) object.__setattr__(self, "_r_grid", r_grid) object.__setattr__(self, "_cdf_grid", cdf_grid) + object.__setattr__(self, "r_t_is_pinned", r_t_is_pinned) @classmethod def from_W0_rc( @@ -444,7 +488,7 @@ def from_W0_rc( xi_grid, psi_grid = solve_king_profile(W0, xi_max=xi_max, n_points=n_ode_points) xi_t = _find_tidal_radius(xi_grid, psi_grid) r_t = r_c * xi_t - return cls( + prof = cls( W0=W0, r_c=r_c, r_t=r_t, @@ -452,6 +496,19 @@ def from_W0_rc( psi_grid=psi_grid, n_grid=n_grid, ) + # Eager refusal (audit J4): for CONCRETE W0 a pinned r_t is a silent + # wrong answer — raise loudly. Under tracing the flag cannot be tested + # (traced bool), so the diagnostic prof.r_t_is_pinned is the only signal + # (mirrors the Engine-B concrete/traced guard, eddington_engine.py). + if _is_concrete(prof.r_t_is_pinned) and bool(prof.r_t_is_pinned): + raise ValueError( + f"King ODE domain too small: psi(xi) never reached 0 within " + f"xi_max={float(xi_max):.1f} for W0={float(W0):.2f}, so r_t is " + f"PINNED to the grid boundary (a wrong tidal radius). Increase " + f"xi_max (and n_ode_points), or omit them to auto-size. For " + f"traced/jit W0 pass an explicit xi_max sized for the largest W0." + ) + return prof def sample_positions( self, diff --git a/src/progenax/profiles/michie.py b/src/progenax/profiles/michie.py index 75ef3459..cc49a8bb 100644 --- a/src/progenax/profiles/michie.py +++ b/src/progenax/profiles/michie.py @@ -174,7 +174,11 @@ def __init__(self, W0, r_c, r_a, r_t, xi_grid, psi_grid, n_grid: int = 1000): xi_a = jnp.asarray(xi_grid, dtype=jnp.float64) psi_a = jnp.asarray(psi_grid, dtype=jnp.float64) - r_grid = jnp.linspace(0.0, r_t_a, n_grid) + # sqrt-stretched grid (r = r_t * u^2): concentrates points in the core, + # which a linear grid under-resolves at high W0 (audit R4 — same fix as + # KingProfile). Smooth in r_t -> differentiable. + u_grid = jnp.linspace(0.0, 1.0, n_grid) + r_grid = r_t_a * u_grid**2 psi_vals = jnp.interp(r_grid / r_c_a, xi_a, psi_a, left=W0_a, right=0.0) s_vals = r_grid / r_a_a # = (r/r_c) / (r_a/r_c) @@ -183,11 +187,11 @@ def __init__(self, W0, r_c, r_a, r_t, xi_grid, psi_grid, n_grid: int = 1000): rho_grid = jnp.where(rho0 > 1e-10, rho_vals / rho0, 0.0) rho_grid = jnp.where(r_grid <= r_t_a, rho_grid, 0.0) + # Non-uniform trapezoid: variable spacing -> weight by diff(r_grid). integrand = 4.0 * jnp.pi * r_grid**2 * rho_grid - dr = r_grid[1] - r_grid[0] M_cum = jnp.concatenate([ jnp.zeros(1, dtype=integrand.dtype), - jnp.cumsum(0.5 * (integrand[1:] + integrand[:-1])) * dr, + jnp.cumsum(0.5 * (integrand[1:] + integrand[:-1]) * jnp.diff(r_grid)), ]) cdf_grid = M_cum / (M_cum[-1] + 1e-30) diff --git a/src/progenax/tidal.py b/src/progenax/tidal.py index eb928d5a..068a0a59 100644 --- a/src/progenax/tidal.py +++ b/src/progenax/tidal.py @@ -133,6 +133,14 @@ def apply_tidal_truncation( differentiable. Unlike boolean-mask indexing, the output keeps a static shape ``N``, so the function is ``jit`` / ``vmap`` / ``grad`` safe. + .. warning:: + The survivors keep the velocities they were drawn with for the UNtruncated + potential, so the truncated set is SUPER-VIRIAL w.r.t. its own (now shallower) + potential — some stars near ``r_t`` are formally unbound and the set is not a + stationary equilibrium (audit S4). Re-virialize the survivors + (``virial_scale`` / ``rescale_velocities_to_virial``) or use an r_t-consistent + equilibrium model (King / LIMEPY) if you need a stationary IC. + Args: positions: Particle positions (N, 3) velocities: Particle velocities (N, 3) diff --git a/tests/README.md b/tests/README.md index 2b839534..47ed16a6 100644 --- a/tests/README.md +++ b/tests/README.md @@ -2,7 +2,8 @@ Comprehensive physics validation for stellar initial conditions generation. -**874 tests** across 3 tiers validating astrophysics from published literature (1911-2018). +**~1243 released-core tests** across 3 tiers validating astrophysics from published +literature (1911–2018). Counts drift as the suite grows — **see CI for the live number**. ## Table of Contents @@ -45,10 +46,10 @@ pytest tests/validation/test_plummer_physics.py -v | Tier | Tests | Purpose | Runtime | |------|-------|---------|---------| -| **Unit** | 742 | Functional correctness (shapes, bounds, roundtrips) | ~30s | -| **Integration** | 24 | JAX transformations (JIT, grad, vmap) | ~10s | -| **Validation** | 108 | Physics accuracy against literature | ~15s | -| **Total** | **874** | Comprehensive coverage | ~55s | +| **Unit** | ~956 | Functional correctness (shapes, bounds, roundtrips) | — | +| **Integration** | ~43 | JAX transformations (JIT, grad, vmap) | — | +| **Validation** | ~244 | Physics accuracy against literature | — | +| **Total** | **~1243** | Comprehensive coverage (see CI for the live count) | — | --- @@ -561,4 +562,4 @@ N_stats = 10000 # Precise statistics --- -*Last updated: 2026-06-03 | 874 tests | progenax v0.1.0* +*Last updated: 2026-06-12 | ~1243 released-core tests (see CI for the live count) | progenax v0.1.0* diff --git a/tests/conftest.py b/tests/conftest.py index 35f6f6fa..ec53ea7b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,6 +8,8 @@ - JAX configuration """ +import os + import pytest import jax import jax.numpy as jnp @@ -30,6 +32,21 @@ def pytest_configure(config): ) +@pytest.fixture(autouse=True, scope="module") +def _bound_jax_compilation_cache(): + """Release the JAX compilation cache after each test MODULE in CI. + + JAX caches every compiled executable; the profile/DF modules compile many + distinct diffrax-ODE / Eddington functions, so on the 2-core/7 GB CI runner + the cumulative cache OOM-kills a long shard (exit 143 — audit R1 follow-up). + Clearing between modules bounds the growth. Gated on CI (GitHub sets + ``CI=true``) so the local inner loop keeps its warm cache; force locally with + ``PROGENAX_CLEAR_JAX_CACHES=1``.""" + yield + if os.environ.get("CI") or os.environ.get("PROGENAX_CLEAR_JAX_CACHES"): + jax.clear_caches() + + # ============================================================================= # Standard Particle Counts # ============================================================================= diff --git a/tests/unit/test_demo_inference.py b/tests/experimental/unit/test_demo_inference.py similarity index 97% rename from tests/unit/test_demo_inference.py rename to tests/experimental/unit/test_demo_inference.py index 476dee8a..20d00e0a 100644 --- a/tests/unit/test_demo_inference.py +++ b/tests/experimental/unit/test_demo_inference.py @@ -17,8 +17,12 @@ import jax.numpy as jnp import numpy as np -# tests/unit/test_demo_inference.py -> parents[0]=unit, [1]=tests, [2]=repo root. -_REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +# tests/experimental/unit/test_demo_inference.py +# -> parents[0]=unit, [1]=experimental, [2]=tests, [3]=repo root. +# Lives in the experimental tier: the helper pulls optax/blackjax (the +# [experimental] extra), so it must NOT gate the released-core unit shard, +# which syncs --extra dev only (audit R1 follow-up: red CI had hidden this). +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] sys.path.insert(0, str(_REPO_ROOT / "scripts")) import _demo_inference as di # noqa: E402 diff --git a/tests/integration/test_end_to_end.py b/tests/integration/test_end_to_end.py index 889cfc3a..368ca4a8 100644 --- a/tests/integration/test_end_to_end.py +++ b/tests/integration/test_end_to_end.py @@ -54,9 +54,8 @@ def test_chabrier_to_king_ic(self): # Create King IC W0 = 6.0 - xi_grid, psi_grid = solve_king_profile(W0=W0) - profile = KingProfile(W0=W0, r_c=1.0, r_t=10.0, xi_grid=xi_grid, psi_grid=psi_grid) - velocity_df = KingVelocityDF(W0=W0, r_c=profile.r_c, r_t=profile.r_t) + profile = KingProfile.from_W0_rc(W0=W0, r_c=1.0) # self-consistent r_t + velocity_df = KingVelocityDF(W0=W0, r_c=profile.r_c) G = 1.0 result = build_spatial_ic( diff --git a/tests/integration/test_readme_examples.py b/tests/integration/test_readme_examples.py new file mode 100644 index 00000000..c9fac2ff --- /dev/null +++ b/tests/integration/test_readme_examples.py @@ -0,0 +1,21 @@ +"""Execute every python code block in README.md (audit R7/D3: would have caught +the deleted-API examples automatically). + +Skips blocks marked ``# doctest: +SKIP`` (e.g. GPU-only or illustrative +fragments). Each block runs in a fresh namespace. +""" +import re +from pathlib import Path + +import pytest + +README = (Path(__file__).resolve().parents[2] / "README.md").read_text() +BLOCKS = [ + b for b in re.findall(r"```python\n(.*?)```", README, re.DOTALL) + if "+SKIP" not in b +] + + +@pytest.mark.parametrize("i", range(len(BLOCKS))) +def test_readme_block_executes(i): + exec(compile(BLOCKS[i], f"README.md:block{i}", "exec"), {}) diff --git a/tests/unit/diagnostics/test_lazy_imports.py b/tests/unit/diagnostics/test_lazy_imports.py new file mode 100644 index 00000000..cb3f098f --- /dev/null +++ b/tests/unit/diagnostics/test_lazy_imports.py @@ -0,0 +1,27 @@ +import subprocess +import sys + + +def test_diagnostics_import_without_scipy_gives_actionable_error(): + """Audit R9: in an env without scipy, importing the subpackage must NOT + crash at import time; the JAX-native q_approx must still work; and calling + an exact (scipy-backed) function must raise an ImportError naming the + [diagnostics] extra.""" + code = ( + "import sys;" + "sys.modules['scipy'] = None; sys.modules['scipy.sparse'] = None;" + "sys.modules['scipy.sparse.csgraph'] = None; sys.modules['scipy.spatial'] = None;" + "sys.modules['scipy.spatial.distance'] = None;" + "import progenax.diagnostics as d;" # must import fine + "from progenax.diagnostics import q_approx;" # JAX-native: must work + "import jax.numpy as jnp, jax;" + "q_approx(jax.random.normal(jax.random.PRNGKey(0), (50, 3)));" + "exc = None\n" + "try:\n" + " d.compute_q_parameter(jax.random.normal(jax.random.PRNGKey(0), (50, 3)))\n" + "except ImportError as e:\n" + " exc = e\n" + "assert exc is not None and 'diagnostics' in str(exc), exc\n" + ) + r = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert r.returncode == 0, r.stderr diff --git a/tests/unit/imf/test_imf_gradients.py b/tests/unit/imf/test_imf_gradients.py index f18ccf7d..6593c4f5 100644 --- a/tests/unit/imf/test_imf_gradients.py +++ b/tests/unit/imf/test_imf_gradients.py @@ -121,3 +121,66 @@ def loss(mu): grad_val = jax.grad(loss)(0.2) assert jnp.isfinite(grad_val), f"Gradient is {grad_val}" assert jnp.abs(grad_val) > 1e-6, f"Gradient is effectively zero: {grad_val}" + + +class TestAlphaOneGradients: + """Audit R10: bare where(|1-a| mid > hi # both stats decrease as alpha steepens + second_diff = abs((lo - mid) - (mid - hi)) + assert second_diff < 0.05 * abs(lo - hi) # near-linear: no kink at a=1 + + def test_sample_statistic_grad_finite_at_alpha_one(self): + def loss(a): + m = self._imf(a).sample(jax.random.PRNGKey(0), 500) + return jnp.mean(jnp.log(m)) + assert jnp.isfinite(jax.grad(loss)(1.0)) diff --git a/tests/unit/imf/test_moe_full.py b/tests/unit/imf/test_moe_full.py index 58af5880..1fd28dfb 100644 --- a/tests/unit/imf/test_moe_full.py +++ b/tests/unit/imf/test_moe_full.py @@ -96,8 +96,16 @@ def test_range_and_twin_fraction(self): m1 = jnp.full(n, 1.0); P = jnp.full(n, 10.0) # solar logP=1, Ftwin=0.30 q = moe.sample(key, m1, P) assert jnp.all((q >= 0.1) & (q <= 1.0)) - # twin fraction >= the excess weight (0.30), plus power-law content in [0.95,1] - assert jnp.mean(q > 0.95) > 0.28 + # P(q>0.95) is the twin block PLUS the power-law content in [0.95,1], + # NOT F_twin itself (F_twin is the EXCESS over q>0.3 — audit R3). Derive + # the expectation by quadrature of the fixed pdf rather than a magic + # constant; the old 0.28 gate encoded the pre-fix +22% twin overweight. + qs = jnp.linspace(0.1, 1.0, 200_001) + p = jax.vmap(lambda qq: moe.pdf(qq, jnp.asarray(1.0), jnp.asarray(10.0)))(qs) + expected = float( + jnp.trapezoid(jnp.where(qs >= 0.95, p, 0.0), qs) / jnp.trapezoid(p, qs) + ) + assert abs(float(jnp.mean(q > 0.95)) - expected) < 0.01 def test_sampled_mean_matches_pdf(self): moe = _full() @@ -138,6 +146,63 @@ def loss(qmin): assert jnp.abs(ad - fd) / (jnp.abs(ad) + 1e-12) < 1e-3, f"ad={ad} fd={fd}" +class TestFTwinPaperConvention: + """MD17 p.5 / Fig.2: F_twin = excess-twin fraction of q > 0.3 companions. + + Audit finding R3: the pre-fix mixture realized F_twin = 0.367 instead of + 0.300 at (M1=1, logP=1). These tests pin the paper convention at four + Table-13 nodes via quadrature of the pdf's mixture components. + """ + + NODES = [(1.0, 1.0), (1.0, 3.0), (3.2, 1.0), (12.0, 3.0)] # (M1 [Msun], logP) + + def test_f_twin_excess_over_q_gt_03(self): + md = _full() + qs = jnp.linspace(md.q_min, 1.0, 200_001) + for m1, logP in self.NODES: + mass = jnp.asarray(m1) + period = jnp.asarray(10.0**logP) + ft_table = float(md.f_twin(period, mass)) + p_pl, p_twin = md._components(qs, mass, period) + mask = qs >= 0.3 + twin_mass = float(jnp.trapezoid(jnp.where(mask, p_twin, 0.0), qs)) + total_gt03 = float( + jnp.trapezoid(jnp.where(mask, p_pl + p_twin, 0.0), qs) + ) + realized = twin_mass / total_gt03 + assert abs(realized - ft_table) < 2e-3, ( + f"(M1={m1}, logP={logP}): realized paper-convention F_twin " + f"{realized:.4f} != Table 13 value {ft_table:.4f}" + ) + + def test_pdf_still_normalized(self): + md = _full() + qs = jnp.linspace(md.q_min, 1.0, 200_001) + for m1, logP in self.NODES: + p = jax.vmap( + lambda q: md.pdf(q, jnp.asarray(m1), jnp.asarray(10.0**logP)) + )(qs) + Z = float(jnp.trapezoid(p, qs)) + assert abs(Z - 1.0) < 1e-3 + + def test_sample_matches_pdf_twin_fraction(self): + """Sampled P(q >= 0.95) must match the quadrature of the FIXED pdf.""" + md = _full() + n = 200_000 + key = jax.random.PRNGKey(0) + m1 = jnp.full((n,), 1.0) + periods = jnp.full((n,), 10.0) + q = md.sample(key, m1, periods) + qs = jnp.linspace(md.q_min, 1.0, 200_001) + p = jax.vmap(lambda qq: md.pdf(qq, jnp.asarray(1.0), jnp.asarray(10.0)))(qs) + expected = float( + jnp.trapezoid(jnp.where(qs >= 0.95, p, 0.0), qs) + / jnp.trapezoid(p, qs) + ) + observed = float(jnp.mean(q >= 0.95)) + assert abs(observed - expected) < 0.01 # shot noise ~0.001 at n=2e5 + + class TestMoePeriod: def test_range(self): from progenax.imf.binary import MoePeriod diff --git a/tests/unit/imf/test_sampling_modes.py b/tests/unit/imf/test_sampling_modes.py new file mode 100644 index 00000000..aa343084 --- /dev/null +++ b/tests/unit/imf/test_sampling_modes.py @@ -0,0 +1,67 @@ +"""Tests for the BaseIMF mass-target sampling modes (audit R5: previously zero tests). + +Semantics (from progenax/imf/base.py): + - sample_m_total(key, m_total, n_max=None) -> (masses_padded, n_live): + prefix-cut; the realized total OVERSHOOTS m_total by at most one star + (the boundary star pushes the cumsum past the target). + - sample_m_total_packed(key, m_total, n_max) -> (masses_packed, n_live_float): + fractional boundary star -> hits m_total (near-)exactly. + - sample_fixed_n(key, n, m_total) -> (n,): + stratified quantile stretch (q <= 1). The achievable total is capped at + sum(ppf(u_base)); a target above that ceiling used to undershoot SILENTLY. +""" +import jax +import jax.numpy as jnp +import pytest + +from progenax.imf import Maschberger + + +class TestSampleFixedN: + def test_reachable_target_is_hit(self): + imf = Maschberger() + key = jax.random.PRNGKey(0) + n = 1000 + # natural total ~ n * E[m]; pick a comfortably reachable target + target = 250.0 + m = imf.sample_fixed_n(key, n, target) + assert m.shape == (n,) + assert abs(float(m.sum()) / target - 1.0) < 0.05 # ~target, not "exactly" + + def test_unreachable_target_raises(self): + """Pre-fix behavior: silently returned 349 Msun for a 500 Msun target.""" + imf = Maschberger() + key = jax.random.PRNGKey(0) + with pytest.raises(ValueError, match="unreachable"): + imf.sample_fixed_n(key, 1000, 500.0) + + def test_masses_within_imf_bounds(self): + imf = Maschberger() + m = imf.sample_fixed_n(jax.random.PRNGKey(1), 500, 120.0) + assert float(m.min()) >= imf.m_min + assert float(m.max()) <= imf.m_max + + +class TestSampleMTotalModes: + @pytest.mark.parametrize("seed", [0, 1]) + def test_sample_m_total_hits_target(self, seed): + imf = Maschberger() + masses, n_live = imf.sample_m_total(jax.random.PRNGKey(seed), m_total=500.0) + total = float(masses.sum()) + # prefix-cut overshoots by AT MOST one (max-mass) star, never undershoots + # (unless n_max is exhausted; the auto n_max is sized 2x so it isn't). + assert 500.0 <= total <= 500.0 + imf.m_max + assert int(n_live) == int((masses > 0).sum()) + + def test_sample_m_total_packed_consistent_with_unpacked(self): + imf = Maschberger() + key = jax.random.PRNGKey(0) + n_max = 4000 # generous: ~1500-2500 stars make 500 Msun + packed, n_float = imf.sample_m_total_packed(key, 500.0, n_max) + # the fractional boundary star makes the packed total hit the target + # (near-)exactly -- well within one stellar mass. + assert abs(float(packed.sum()) - 500.0) < imf.m_max + # same draw, unpacked: the live count brackets the packed fractional count + # (n_live = floor(boundary)+1; n_float = floor(boundary)+fraction). + unpacked, n_live = imf.sample_m_total(key, 500.0, n_max=n_max) + assert abs(float(n_float) - float(n_live)) <= 1.5 diff --git a/tests/unit/kinematics/test_df_gradients.py b/tests/unit/kinematics/test_df_gradients.py index e076341c..f97291ce 100644 --- a/tests/unit/kinematics/test_df_gradients.py +++ b/tests/unit/kinematics/test_df_gradients.py @@ -77,7 +77,7 @@ def test_grad_matches_fd(self, param, base, eps, rtol): key = jax.random.PRNGKey(11) def loss(x): - kw = {"W0": 7.0, "r_c": 1.0, "r_t": 10.0} + kw = {"W0": 7.0, "r_c": 1.0} kw[param] = x v = KingVelocityDF(**kw).sample_velocities(pos, masses, key, G=G) return jnp.mean(jnp.sum(v**2, axis=1)) @@ -92,7 +92,7 @@ def loss(x): def test_jit_compatible(self): pos = _fixed_positions(128, r_scale=3.0, seed=12) masses = jnp.ones(128) - df = KingVelocityDF(W0=7.0, r_c=1.0, r_t=10.0) + df = KingVelocityDF(W0=7.0, r_c=1.0) v = jax.jit(lambda p, m, k: df.sample_velocities(p, m, k, G=G))( pos, masses, jax.random.PRNGKey(2) ) diff --git a/tests/unit/kinematics/test_eddington_invert.py b/tests/unit/kinematics/test_eddington_invert.py index cf3f9233..66462bc6 100644 --- a/tests/unit/kinematics/test_eddington_invert.py +++ b/tests/unit/kinematics/test_eddington_invert.py @@ -4,15 +4,16 @@ import numpy as np import pytest -# Exact-value pins from _eff_eddington_table(1.0, 4.0, 12.0, r_a) -# (full repr precision). Re-capture (never loosen) on JAX upgrades or -# DELIBERATE op-order changes only. -# Re-captured 2026-06-10 (consolidation Task 8): the table's cumulative -# trapezoid moved onto progenax.numerics.cumulative_trapezoid (scalar dx -# OUTSIDE the cumsum; the old inline used the non-uniform jnp.diff(r) -# INSIDE), shifting the raw table integral by ~1 ulp (~9e-16); the -# Eddington inversion (d2rho/dPsi2 + Abel weight) amplifies that to -# <=8e-15 relative (~55 ulp, worst iso f[250]) on the pinned f values. +# Value pins from _eff_eddington_table(1.0, 4.0, 12.0, r_a), captured +# 2026-06-10 (consolidation Task 8). Compared with a TIGHT relative tolerance +# (rel=1e-9), NOT bit-exact == : the table's cumulative-trapezoid op order +# shifts the raw integral by ~1 ulp and the Eddington inversion (d2rho/dPsi2 + +# Abel weight) amplifies that to <=8e-15 relative locally; cross-platform FP +# (macOS-ARM dev vs Linux-x86 CI) pushes a few more ulp, which made the old +# `==` fail in CI (audit T8: bit-exact across platforms is non-portable). +# rel=1e-9 sits ~6 orders above that noise floor, so it still catches any real +# refactoring drift while staying portable. Re-capture on DELIBERATE op-order +# changes only. PINS = { "iso": ( 5.5042467299701805, @@ -41,16 +42,19 @@ class TestExtractionRegression: @pytest.mark.parametrize("tag,ra", [("iso", None), ("om", 2.0)]) - def test_eff_table_bit_identical_after_refactor(self, tag, ra): + def test_eff_table_stable_after_refactor(self, tag, ra): """The refactored _eff_eddington_table (now calling eddington_invert) - reproduces the pre-refactor values EXACTLY (same grids, same ops).""" + reproduces the pre-refactor values to a tight relative tolerance + (rel=1e-9) — a PORTABLE refactoring guard. Not bit-exact: == failed + cross-platform by ~1 ulp (audit T8).""" from progenax.kinematics.eff_df import _eff_eddington_table r, Psi, E, f, mu = _eff_eddington_table(1.0, 4.0, 12.0, ra) Psi0, mu_pin, f_pins = PINS[tag] - assert float(Psi[0]) == Psi0 and float(mu) == mu_pin + assert float(Psi[0]) == pytest.approx(Psi0, rel=1e-9) + assert float(mu) == pytest.approx(mu_pin, rel=1e-9) for idx, pin in zip((1, 250, 500, 750, 998), f_pins): - assert float(f[idx]) == pin, f"f[{idx}] drifted" + assert float(f[idx]) == pytest.approx(pin, rel=1e-9), f"f[{idx}] drifted" def _plummer_grids(a=1.0, rt=100.0, n=20000): diff --git a/tests/unit/kinematics/test_king_df.py b/tests/unit/kinematics/test_king_df.py index b145b7d3..699b2e1a 100644 --- a/tests/unit/kinematics/test_king_df.py +++ b/tests/unit/kinematics/test_king_df.py @@ -17,12 +17,25 @@ G = STELLAR.G +class TestKingVelocityDFApi: + """Audit A3: r_t was a stored-but-never-used field (sigma^2 = GM/(9 r_c mu), + truncation comes from psi(xi)). Removed — no backwards-compat shim.""" + + def test_constructs_without_r_t(self): + df = KingVelocityDF(W0=7.0, r_c=1.0) + assert not hasattr(df, "r_t") + + def test_passing_r_t_is_rejected(self): + with pytest.raises(TypeError): + KingVelocityDF(W0=7.0, r_c=1.0, r_t=10.0) + + class TestKingVelocityDFPhysics: """Test King velocity DF physical properties.""" def test_isotropic_distribution(self): """Velocities are isotropically distributed.""" - df = KingVelocityDF(W0=5.0, r_c=1.0, r_t=10.0) + df = KingVelocityDF(W0=5.0, r_c=1.0) N = 1000 r = 2.0 positions = jnp.array([[r, 0.0, 0.0]] * N) @@ -40,7 +53,7 @@ def test_isotropic_distribution(self): def test_mean_velocity_zero(self): """Mean velocity is zero (no bulk motion).""" - df = KingVelocityDF(W0=5.0, r_c=1.0, r_t=10.0) + df = KingVelocityDF(W0=5.0, r_c=1.0) N = 1000 r = 2.0 positions = jnp.array([[r, 0.0, 0.0]] * N) @@ -57,7 +70,7 @@ def test_mean_velocity_zero(self): def test_dispersion_decreases_outward(self): """Velocity dispersion decreases with radius.""" - df = KingVelocityDF(W0=7.0, r_c=1.0, r_t=10.0) + df = KingVelocityDF(W0=7.0, r_c=1.0) N = 500 r_inner = 1.0 @@ -89,7 +102,7 @@ class TestKingTableRouting: draws are variate-paired (same per-star key splits).""" def _two_dfs(self): - kw = dict(W0=5.0, r_c=1.0, r_t=10.0) + kw = dict(W0=5.0, r_c=1.0) return (KingVelocityDF(**kw), # default: table KingVelocityDF(**kw, speed_method="quadrature")) @@ -117,10 +130,10 @@ def test_speed_distribution_ks(self): assert D < 0.02 def test_table_default_and_quadrature_static(self): - df = KingVelocityDF(W0=5.0, r_c=1.0, r_t=10.0) + df = KingVelocityDF(W0=5.0, r_c=1.0) assert df.speed_method == "table" with pytest.raises(ValueError, match="speed_method"): - KingVelocityDF(W0=5.0, r_c=1.0, r_t=10.0, speed_method="exact") + KingVelocityDF(W0=5.0, r_c=1.0, speed_method="exact") def test_g1_lowered_exponential_is_king_weight(self): """The exact identity the routing relies on: the LIMEPY lowered @@ -144,18 +157,18 @@ class TestKingSamplerOptimization: def test_table_cached_at_construction(self): from progenax.profiles.limepy_tables import SpeedCDFTable - df = KingVelocityDF(W0=5.0, r_c=1.0, r_t=10.0) + df = KingVelocityDF(W0=5.0, r_c=1.0) assert isinstance(df.speed_table, SpeedCDFTable) def test_quadrature_method_has_no_table(self): - df = KingVelocityDF(W0=5.0, r_c=1.0, r_t=10.0, + df = KingVelocityDF(W0=5.0, r_c=1.0, speed_method="quadrature") assert df.speed_table is None def test_cached_table_bit_identical_to_fresh_build(self): from progenax.profiles.limepy_tables import SpeedCDFTable - df = KingVelocityDF(W0=5.0, r_c=1.0, r_t=10.0) + df = KingVelocityDF(W0=5.0, r_c=1.0) fresh = SpeedCDFTable.build(df.W0, jnp.asarray(1.0)) np.testing.assert_array_equal(np.asarray(df.speed_table.cdf), np.asarray(fresh.cdf)) @@ -164,7 +177,7 @@ def test_same_key_same_velocities(self): from progenax.profiles.king import KingProfile prof = KingProfile.from_W0_rc(W0=5.0, r_c=1.0) - df = KingVelocityDF(W0=5.0, r_c=1.0, r_t=10.0) + df = KingVelocityDF(W0=5.0, r_c=1.0) m = jnp.ones(500) pos = prof.sample_positions(m, jax.random.PRNGKey(0)) v1 = df.sample_velocities(pos, m, jax.random.PRNGKey(1), G=1.0) diff --git a/tests/unit/kinematics/test_limepy_df.py b/tests/unit/kinematics/test_limepy_df.py index 644dbe9b..03df5df0 100644 --- a/tests/unit/kinematics/test_limepy_df.py +++ b/tests/unit/kinematics/test_limepy_df.py @@ -97,7 +97,7 @@ def test_g1_velocity_scale_matches_king(self): from progenax.kinematics.king_df import KingVelocityDF from progenax.kinematics.limepy_df import LIMEPYVelocityDF - king = KingVelocityDF(W0=7.0, r_c=1.0, r_t=30.0) + king = KingVelocityDF(W0=7.0, r_c=1.0) lim = LIMEPYVelocityDF(W0=7.0, g=1.0, r_c=1.0) M = jnp.asarray(5000.0) np.testing.assert_allclose( @@ -115,7 +115,7 @@ def test_g1_dispersion_profile_matches_king(self): _, _, m, pos_l, vel_l = _build_ic(W0=7.0, g=1.0, N=N, seed=3) # King prof_k = KingProfile.from_W0_rc(7.0, 1.0) - df_k = KingVelocityDF(W0=7.0, r_c=1.0, r_t=float(prof_k.r_t)) + df_k = KingVelocityDF(W0=7.0, r_c=1.0) kp, kv = jax.random.split(jax.random.PRNGKey(3)) pos_k = prof_k.sample_positions(jnp.ones(N), kp) vel_k = df_k.sample_velocities(pos_k, jnp.ones(N), kv, G=G) diff --git a/tests/unit/kinematics/test_rotation.py b/tests/unit/kinematics/test_rotation.py index 28533443..02c5d3e0 100644 --- a/tests/unit/kinematics/test_rotation.py +++ b/tests/unit/kinematics/test_rotation.py @@ -120,3 +120,21 @@ def test_zero_at_center(self): ) assert jnp.allclose(v_out, 0.0, atol=1e-10) + + +class TestZeroAxisRefused: + """Audit S15: a concrete zero rotation axis has no rotation direction; the + old code silently no-op'd (axis/max(mag,1e-30)=0) under a stale 'NaN' + comment. Concrete zero axis now raises.""" + + def test_solid_body_zero_axis_raises(self): + v = jnp.zeros((5, 3)) + pos = jax.random.normal(jax.random.PRNGKey(0), (5, 3)) + with pytest.raises(ValueError, match="zero vector|rotation direction"): + apply_solid_body_rotation(v, pos, omega=0.1, axis=jnp.zeros(3)) + + def test_differential_zero_axis_raises(self): + v = jnp.zeros((5, 3)) + pos = jax.random.normal(jax.random.PRNGKey(1), (5, 3)) + with pytest.raises(ValueError, match="zero vector|rotation direction"): + apply_differential_rotation(v, pos, v_peak=1.0, R_peak=1.0, axis=jnp.zeros(3)) diff --git a/tests/unit/profiles/test_king.py b/tests/unit/profiles/test_king.py index ce5e1f4b..7cbb63ee 100644 --- a/tests/unit/profiles/test_king.py +++ b/tests/unit/profiles/test_king.py @@ -159,3 +159,48 @@ def sample_and_sum(m, k): result = sample_and_sum(masses, key) assert jnp.isfinite(result) + + +class TestRtBoundaryPinning: + """Audit J4: a too-small ODE domain leaves psi(xi)>0 everywhere, so + _find_tidal_radius silently pins xi_t to the grid boundary (a wrong r_t). + Concrete construction must REFUSE loudly; traced construction can't raise + on a traced bool, so it stores an r_t_is_pinned diagnostic instead (the + Engine-B two-tier concrete/traced pattern).""" + + def test_healthy_solve_not_flagged(self): + prof = KingProfile.from_W0_rc(W0=7.0, r_c=1.0) + assert not bool(prof.r_t_is_pinned) + + def test_concrete_pinned_solve_raises(self): + """Eager construction with a domain that pins xi_t refuses loudly.""" + with pytest.raises(ValueError, match="pinned|xi_max"): + KingProfile.from_W0_rc(W0=12.0, r_c=1.0, xi_max=50.0, n_ode_points=500) + + def test_traced_pinned_solve_is_flagged_not_raised(self): + """Under tracing the raise is impossible (traced bool); the flag is set.""" + prof = jax.jit( + lambda w: KingProfile.from_W0_rc( + W0=w, r_c=1.0, xi_max=50.0, n_ode_points=500 + ) + )(12.0) + assert bool(prof.r_t_is_pinned) + + def test_traced_healthy_solve_constructs(self): + prof = jax.jit(lambda w: KingProfile.from_W0_rc(W0=w, r_c=1.0))(7.0) + assert jnp.isfinite(prof.r_t) and not bool(prof.r_t_is_pinned) + + +class TestRtConsistencyWarning: + """Audit S1: the direct KingProfile constructor accepts an arbitrary r_t + inconsistent with c(W0) — KingProfile(W0=7, r_c=1, r_t=10) silently builds + a non-self-consistent, non-equilibrium model. Concrete inputs now warn.""" + + def test_inconsistent_r_t_warns(self): + xi_grid, psi_grid = solve_king_profile(7.0) # c(7): r_t/r_c ~ 30, not 10 + with pytest.warns(UserWarning, match="inconsistent"): + KingProfile(W0=7.0, r_c=1.0, r_t=10.0, xi_grid=xi_grid, psi_grid=psi_grid) + + def test_consistent_r_t_no_warning(self, recwarn): + KingProfile.from_W0_rc(W0=7.0, r_c=1.0) # derives r_t -> self-consistent + assert not any("inconsistent" in str(w.message) for w in recwarn.list) diff --git a/tests/unit/profiles/test_profile_potential.py b/tests/unit/profiles/test_profile_potential.py index de4e17ac..d612247c 100644 --- a/tests/unit/profiles/test_profile_potential.py +++ b/tests/unit/profiles/test_profile_potential.py @@ -63,7 +63,10 @@ def test_enclosed_mass_matches_profile_cdf(self): # equals the profile's own CDF (both = cumulative trapezoid of rho*r^2). rgrid = prof._r_grid rho_t = (1.0 + (rgrid / prof.a) ** 2) ** (-prof.gamma / 2.0) - dr = rgrid[1] - rgrid[0] + # NON-UNIFORM trapezoid: _r_grid is sqrt-stretched (audit R4), so weight + # each trapezoid by its own width diff(rgrid). The invariant is unchanged + # (potential's enclosed mass == profile CDF, both = cumtrap of rho*r^2). + dr = jnp.diff(rgrid) I2 = jnp.concatenate([jnp.zeros(1), jnp.cumsum(0.5 * (rho_t[1:] * rgrid[1:] ** 2 + rho_t[:-1] * rgrid[:-1] ** 2) * dr)]) cdf_from_pot = I2 / I2[-1] assert jnp.max(jnp.abs(cdf_from_pot - prof._cdf_grid)) < 1e-9 diff --git a/tests/unit/substructure/test_q_approx.py b/tests/unit/substructure/test_q_approx.py index 5cc5ed75..5cfc7a4f 100644 --- a/tests/unit/substructure/test_q_approx.py +++ b/tests/unit/substructure/test_q_approx.py @@ -150,6 +150,18 @@ def test_explicit_method_selection(self): Q_fast = q_approx(positions, method="fast") assert jnp.isfinite(Q_naive) and jnp.isfinite(Q_fast) + def test_auto_dispatch_small_n_equals_naive(self): + """Audit J6: auto picks the naive path for small N (Python-if dispatch).""" + from progenax.diagnostics.q_approx import q_approx + pos = jax.random.normal(jax.random.PRNGKey(0), (200, 3)) + assert jnp.allclose(q_approx(pos, method="auto"), q_approx(pos, method="naive")) + + def test_auto_dispatch_large_n_equals_fast(self): + """Audit J6: auto picks the fast path for large N (Python-if dispatch).""" + from progenax.diagnostics.q_approx import q_approx + pos = jax.random.normal(jax.random.PRNGKey(0), (1500, 3)) + assert jnp.allclose(q_approx(pos, method="auto"), q_approx(pos, method="fast")) + class TestExports: """Test module exports from progenax.diagnostics.""" diff --git a/tests/unit/test_builders.py b/tests/unit/test_builders.py index 9ca53169..e641fdb7 100644 --- a/tests/unit/test_builders.py +++ b/tests/unit/test_builders.py @@ -4,31 +4,39 @@ import jax.numpy as jnp import pytest +from progenax.builders import compute_stellar_radii -class TestStellarRadii: - """Test stellar radii mass-radius relation.""" - def test_stellar_radii_solar_mass(self): - """1 Msun star should have ~1 Rsun radius (R ∝ M^0.8).""" - from progenax.builders import compute_stellar_radii - masses = jnp.array([1.0]) - radii = compute_stellar_radii(masses) - assert jnp.abs(radii[0] - 1.0) < 0.1 +class TestComputeStellarRadii: + """Anchored to Demircan & Kahraman (1991) Ap&SS 181, 313 + observed radii. - def test_stellar_radii_massive_star(self): - """10 Msun star should have R ~ 10^0.8 ~ 6.3 Rsun.""" - from progenax.builders import compute_stellar_radii - masses = jnp.array([10.0]) - radii = compute_stellar_radii(masses) - expected = 10.0**0.8 - assert jnp.abs(radii[0] - expected) < 0.5 + Audit R6: the previous exponents were INVERTED vs MS homology + (M^0.8 above 1 Msun instead of below), giving 10 Msun -> 6.3 Rsun + (ZAMS ~4) and 0.2 Msun -> 0.40 Rsun (observed ~0.22). The old tests + asserted the code's own formula back at itself. + """ + + def test_solar(self): + r = compute_stellar_radii(jnp.array([1.0])) + assert abs(float(r[0]) - 1.06) < 0.01 # D&K91: 1.06 * 1^0.945 + + def test_massive_star_near_zams(self): + r = float(compute_stellar_radii(jnp.array([10.0]))[0]) + assert abs(r - 1.33 * 10.0**0.555) < 1e-6 # = 4.77 Rsun + assert 3.0 < r < 5.5 # within ~25% of ZAMS ~4 Rsun - def test_stellar_radii_brown_dwarf(self): - """Brown dwarf (0.01 Msun) should have R ~ 0.1 Rsun.""" - from progenax.builders import compute_stellar_radii - masses = jnp.array([0.01]) - radii = compute_stellar_radii(masses) - assert 0.05 < radii[0] < 0.15 + def test_m_dwarf(self): + r = float(compute_stellar_radii(jnp.array([0.2]))[0]) + assert abs(r - 1.06 * 0.2**0.945) < 1e-6 # = 0.231 Rsun + assert 0.17 < r < 0.28 # observed ~0.22 Rsun + + def test_brown_dwarf_plateau(self): + r = compute_stellar_radii(jnp.array([0.05, 0.02])) + assert jnp.all(jnp.abs(r - 0.1) < 0.05) + + def test_near_continuous_at_hydrogen_burning_limit(self): + r = compute_stellar_radii(jnp.array([0.079, 0.081])) + assert abs(float(r[1]) / float(r[0]) - 1.0) < 0.1 # no factor-2.4 jump (audit F7) class TestToCOMFrame: @@ -121,6 +129,36 @@ def test_q_less_than_half_collapsing(self): assert Q < 0.5 # Subvirial assert jnp.abs(Q - 0.25) < 0.05 + def test_virial_scale_zero_velocities_raises(self): + """Audit J5: cold input (T=0) used to return all-NaN velocities silently.""" + from progenax.builders import virial_scale + pos = jax.random.normal(jax.random.PRNGKey(0), (50, 3)) + vel = jnp.zeros((50, 3)) + m = jnp.ones(50) + with pytest.raises(ValueError, match="zero kinetic"): + virial_scale(pos, vel, m, Q_target=0.5, G=0.00449) + + def test_rescale_velocities_to_virial_zero_raises(self): + """The deduped sibling shares the same eager T=0 guard (audit J5).""" + from progenax.dynamics.virial import rescale_velocities_to_virial + pos = jax.random.normal(jax.random.PRNGKey(1), (50, 3)) + vel = jnp.zeros((50, 3)) + m = jnp.ones(50) + with pytest.raises(ValueError, match="zero kinetic"): + rescale_velocities_to_virial(pos, vel, m, G=0.00449, target_Q=0.5) + + def test_virial_scale_traced_zero_returns_nan(self): + """Under tracing the guard can't fire; NaN is the honest sentinel.""" + from progenax.builders import virial_scale + pos = jax.random.normal(jax.random.PRNGKey(2), (10, 3)) + m = jnp.ones(10) + + def f(vel): + return virial_scale(pos, vel, m, Q_target=0.5, G=0.00449) + + out = jax.jit(f)(jnp.zeros((10, 3))) + assert jnp.all(jnp.isnan(out)) + class TestBuildSpatialIC: """Test main IC builder.""" diff --git a/tests/unit/test_documented_api.py b/tests/unit/test_documented_api.py new file mode 100644 index 00000000..f3efd783 --- /dev/null +++ b/tests/unit/test_documented_api.py @@ -0,0 +1,27 @@ +"""Every symbol the docs promise must import from the package root (audit R7). + +IGIMF and EnvironmentIMF are deliberately ABSENT: they never existed +(the environment-dependent IMF is the functional BirthEnvironment + +env_to_imf_params API); the docs are being fixed to stop advertising them. +""" +import pytest + +DOCUMENTED_ROOT_SYMBOLS = [ + "PowerLawIMF", "ChabrierIMF", "Maschberger", "TruncatedIMF", "BinaryIMF", + "FlatMassRatio", "PowerLawMassRatio", "TwinPeakedMassRatio", + "MoeDiStefano2017", "MoeDiStefano2017Full", "MoePeriod", "MoeJointOrbit", + "ConstantBinaryFraction", "MassDependentBinaryFraction", + "UniformSphereProfile", # audit L1: only profile missing from the root +] + + +@pytest.mark.parametrize("name", DOCUMENTED_ROOT_SYMBOLS) +def test_symbol_importable_from_root(name): + import progenax + assert hasattr(progenax, name), f"progenax.{name} promised by docs, not exported" + + +def test_phantom_classes_stay_dead(): + import progenax + for phantom in ("IGIMF", "EnvironmentIMF"): + assert not hasattr(progenax, phantom) diff --git a/tests/unit/test_protocols.py b/tests/unit/test_protocols.py index 3efabd6e..4efb664b 100644 --- a/tests/unit/test_protocols.py +++ b/tests/unit/test_protocols.py @@ -30,6 +30,14 @@ def test_velocity_dfs_conform(cls): assert issubclass(cls, VelocityDF) +def test_velocity_df_protocol_single_source(): + """Audit A1: kinematics.api must re-export the canonical protocols.VelocityDF, + not a structural duplicate (which would silently drift).""" + from progenax.protocols import VelocityDF as Canonical + from progenax.kinematics.api import VelocityDF as ApiExport + assert ApiExport is Canonical + + @pytest.mark.parametrize( "imf", [ diff --git a/tests/validation/test_eff_physics.py b/tests/validation/test_eff_physics.py index c02ba7cc..0873a759 100644 --- a/tests/validation/test_eff_physics.py +++ b/tests/validation/test_eff_physics.py @@ -248,6 +248,35 @@ def test_eff_eddington_virial_ratio_mild_truncation(self): ) assert abs(Q - 0.5) < 0.05, f"unscaled Q={Q:.3f} (expected ~0.5 for mild truncation)" + def test_gamma3_default_subvirial_offset_is_pinned(self): + """Audit S2: the gamma=3 DEFAULT + sharp truncation is KNOWN ~8% sub-virial + by construction (a truncated rho with rho(r_t)>0 is not representable by an + f(E)). Pin the band so a regression (to -15%, Q~0.425) OR a vanished offset + (Q~0.5) cannot pass silently. This is a DOCUMENTED LIMITATION, not a target: + use the King model for a strict lowered-DF equilibrium (eff_df.py docstring). + + Measured (N=20000, seeds 0/1/2, G=STELLAR.G): Q = 0.458 +/- 0.0004. The band + 0.45>3sigma) while still excluding the + two regressions (Q<=0.425 and Q>=0.49).""" + from jaxstro.units import STELLAR + from progenax.builders import compute_kinetic_energy, compute_potential_energy + + a, gamma, r_t = 1.0, 3.0, 10.0 # the EFF default config + profile = EFFProfile(a=a, gamma=gamma, r_t=r_t) + df = EFFVelocityDF(a=a, gamma=gamma, r_t=r_t) + masses = jnp.ones(20000) + kp, kv = jax.random.split(jax.random.PRNGKey(0)) + pos = profile.sample_positions(masses, kp) + vel = df.sample_velocities(pos, masses, kv, G=STELLAR.G) + Q = float( + compute_kinetic_energy(vel, masses) + / jnp.abs(compute_potential_energy(pos, masses, G=STELLAR.G)) + ) + assert 0.45 < Q < 0.47, ( + f"gamma=3 Q={Q:.4f} drifted out of the documented sub-virial band " + f"[0.45, 0.47] (measured 0.458 ~8% sub-virial)" + ) + def test_eff_eddington_f_is_physical(self): """The tabulated Eddington DF f(E) is non-negative and increases with energy.""" df = EFFVelocityDF(a=1.0, gamma=3.0, r_t=10.0) @@ -353,5 +382,36 @@ def test_larger_scale_radius_more_core_mass(self, N_validation, key): f"Mass fraction within a should increase with a: a=[0.5,1,2] -> f={fractions}" +@pytest.mark.slow +class TestEFFHighRtOverACoreResolution: + """Audit R4 (EFF extension): the linear CDF grid under-resolves the core at + large r_t/a. Reference = dense quadrature of rho*r^2 (independent of the + profile's internal CDF). Measured pre-fix error at 0.3a (r_t/a=100): +4.2%; + the sqrt-stretched grid (r = r_t*u^2) brings it to <0.1%.""" + + def test_sampled_core_mass_matches_dense_reference_high_rt_over_a(self): + a, gamma, r_t = 1.0, 4.0, 100.0 # r_t/a = 100 + s = jnp.linspace(0.0, r_t, 4_000_000) + rho = (1.0 + (s / a) ** 2) ** (-gamma / 2.0) + integ = s**2 * rho + M = jnp.concatenate( + [jnp.zeros(1), jnp.cumsum(0.5 * (integ[1:] + integ[:-1]) * jnp.diff(s))] + ) + prof = EFFProfile(a=a, gamma=gamma, r_t=r_t) + n = 4_000_000 + r = jnp.linalg.norm( + prof.sample_positions(jnp.ones(n), jax.random.PRNGKey(3)), axis=1 + ) + for r_probe in (0.3, 1.0, 3.0): + m_ref = float(jnp.interp(r_probe, s, M) / M[-1]) + m_samp = float(jnp.mean(r < r_probe)) + shot = 3.0 / (m_ref * n) ** 0.5 + tol = max(0.02, shot) # 2% grid budget or 3-sigma shot, whichever larger + assert abs(m_samp / m_ref - 1.0) < tol, ( + f"r={r_probe}a: sampled {m_samp:.4e} vs reference {m_ref:.4e} " + f"(rel err {(m_samp/m_ref-1)*100:+.2f}%, tol {tol*100:.2f}%)" + ) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/validation/test_king_physics.py b/tests/validation/test_king_physics.py index ecc4d549..393c1ed0 100644 --- a/tests/validation/test_king_physics.py +++ b/tests/validation/test_king_physics.py @@ -155,7 +155,7 @@ def test_velocities_bound_against_king_escape_speed(self, N_validation, key): G = 1.0 profile = KingProfile.from_W0_rc(W0, r_c) - df = KingVelocityDF(W0=W0, r_c=r_c, r_t=float(profile.r_t)) + df = KingVelocityDF(W0=W0, r_c=r_c) masses = jnp.ones(N_validation) key_pos, key_vel = jax.random.split(key) @@ -178,7 +178,7 @@ def test_velocity_isotropy(self, N_stats, key): xi_grid, psi_grid = solve_king_profile(W0) profile = KingProfile(W0=W0, r_c=r_c, r_t=r_t, xi_grid=xi_grid, psi_grid=psi_grid) - df = KingVelocityDF(W0=W0, r_c=r_c, r_t=r_t) + df = KingVelocityDF(W0=W0, r_c=r_c) masses = jnp.ones(N_stats) key_pos, key_vel = jax.random.split(key) @@ -200,7 +200,7 @@ def test_velocity_dispersion_decreases_outward(self, N_validation, key): W0, r_c, r_t = 7.0, 1.0, 10.0 G = 1.0 - df = KingVelocityDF(W0=W0, r_c=r_c, r_t=r_t) + df = KingVelocityDF(W0=W0, r_c=r_c) masses = jnp.ones(N_validation) M_total = float(jnp.sum(masses)) @@ -304,7 +304,7 @@ class TestKingEquilibriumVelocityDF: def _build_ic(self, W0=7.0, r_c=1.0, N=5000, seed=0): from jaxstro.units import STELLAR prof = KingProfile.from_W0_rc(W0, r_c) - df = KingVelocityDF(W0=W0, r_c=r_c, r_t=float(prof.r_t)) + df = KingVelocityDF(W0=W0, r_c=r_c) masses = jnp.ones(N) kp, kv = jax.random.split(jax.random.PRNGKey(seed)) pos = prof.sample_positions(masses, kp) @@ -336,7 +336,7 @@ def test_velocity_sampling_is_differentiable(self): def loss(r_c): prof = KingProfile.from_W0_rc(7.0, 1.0) - df = KingVelocityDF(W0=7.0, r_c=r_c, r_t=float(prof.r_t)) + df = KingVelocityDF(W0=7.0, r_c=r_c) m = jnp.ones(200) kp, kv = jax.random.split(jax.random.PRNGKey(1)) pos = prof.sample_positions(m, kp) @@ -433,7 +433,7 @@ def test_velocity_df_auto_domain_high_W0_equilibrium(self): from progenax.builders import compute_kinetic_energy, compute_potential_energy prof = KingProfile.from_W0_rc(W0=12.0, r_c=1.0) - df = KingVelocityDF(W0=12.0, r_c=1.0, r_t=float(prof.r_t)) # auto domain + df = KingVelocityDF(W0=12.0, r_c=1.0) # auto domain m = jnp.ones(3000) kp, kv = jax.random.split(jax.random.PRNGKey(0)) pos = prof.sample_positions(m, kp) @@ -444,5 +444,69 @@ def test_velocity_df_auto_domain_high_W0_equilibrium(self): assert abs(Q - 0.5) < 0.06, f"W0=12 auto-domain Q={Q:.3f} (expected ~0.5)" +def _dense_king_cumulative_mass(W0): + """Reference M(= 9. + + Reference = direct quadrature of rho_hat(psi) on a dense ODE solve + (xi_max=600, n_points=20000) — independent of the profile's internal CDF. + Measured pre-fix errors at 0.3 r_c: +18% (W0=9), +270% (W0=12). + """ + + @pytest.mark.parametrize("W0", [7.0, 9.0, 12.0]) + def test_sampled_core_mass_matches_dense_reference(self, W0): + xi, M = _dense_king_cumulative_mass(W0) + prof = KingProfile.from_W0_rc(W0=W0, r_c=1.0) + n = 2_000_000 + pos = prof.sample_positions(jnp.ones(n), jax.random.PRNGKey(3)) + r = jnp.linalg.norm(pos, axis=1) + for r_probe in (0.3, 1.0, 3.0): + m_ref = float(jnp.interp(r_probe, xi, M) / M[-1]) + m_samp = float(jnp.mean(r < r_probe)) + shot = 3.0 / (m_ref * n) ** 0.5 # 3 sigma binomial + tol = max(0.03, shot) # 3% grid budget or shot noise, whichever larger + assert abs(m_samp / m_ref - 1.0) < tol, ( + f"W0={W0}, r={r_probe} r_c: sampled M(= '3.14' and sys_platform == 'win32'", - "python_full_version >= '3.14' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'emscripten'", - "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'emscripten'", "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } wheels = [ @@ -715,6 +706,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" }, ] +[[package]] +name = "execnet" +version = "2.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bf/89/780e11f9588d9e7128a3f87788354c7946a9cbb1401ad38a48c4db9a4f07/execnet-2.1.2.tar.gz", hash = "sha256:63d83bfdd9a23e35b9c6a3261412324f964c2ec8dcd8d3c6916ee9373e0befcd", size = 166622, upload-time = "2025-11-12T09:56:37.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ab/84/02fc1827e8cdded4aa65baef11296a9bbe595c474f0d6d758af082d849fd/execnet-2.1.2-py3-none-any.whl", hash = "sha256:67fba928dd5a544b783f6056f449e5e3931a5c378b128bc18501f7ea79e296ec", size = 40708, upload-time = "2025-11-12T09:56:36.333Z" }, +] + [[package]] name = "fastcore" version = "1.13.2" @@ -751,37 +751,57 @@ wheels = [ ] [[package]] -name = "fluxax" -version = "0.1.0" -source = { editable = "../fluxax" } +name = "flowjax" +version = "19.0.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] dependencies = [ - { name = "equinox" }, + { name = "equinox", marker = "python_full_version < '3.11'" }, { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "jax", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "jaxlib", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "jaxstro" }, { name = "jaxtyping", version = "0.3.7", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "jaxtyping", version = "0.3.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "lineax", version = "0.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "optax", marker = "python_full_version < '3.11'" }, + { name = "paramax", marker = "python_full_version < '3.11'" }, + { name = "tqdm", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4b/ab/a780735762adfc5001fc9b18f600f89e87b8a3b28e0c025fdc7a2d8b19eb/flowjax-19.0.0.tar.gz", hash = "sha256:58a1659ccf7d64c4d160df8948afc208cf61d40a819acf19ace61a56778d2432", size = 477201, upload-time = "2026-02-18T20:19:13.726Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/7b/191656d31e13bce736178052c9830b30d9170e7ec256ee2438b4104c6e70/flowjax-19.0.0-py3-none-any.whl", hash = "sha256:633924e38d00732ea885fea4f641cfa7466ab094a99ec66dd2c478ca18209892", size = 61387, upload-time = "2026-02-18T20:19:12.515Z" }, ] -[package.metadata] -requires-dist = [ - { name = "black", marker = "extra == 'dev'", specifier = ">=23.0" }, - { name = "equinox", specifier = ">=0.11.0" }, - { name = "fluxax", extras = ["dev", "viz"], marker = "extra == 'all'" }, - { name = "isort", marker = "extra == 'dev'", specifier = ">=5.0" }, - { name = "jax", specifier = ">=0.4.0" }, - { name = "jaxlib", specifier = ">=0.4.0" }, - { name = "jaxstro", specifier = ">=0.1.0" }, - { name = "jaxtyping", specifier = ">=0.2.0" }, - { name = "matplotlib", marker = "extra == 'viz'", specifier = ">=3.5" }, - { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.0" }, - { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, - { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0" }, - { name = "seaborn", marker = "extra == 'viz'", specifier = ">=0.12" }, +[[package]] +name = "flowjax" +version = "19.1.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +dependencies = [ + { name = "equinox", marker = "python_full_version >= '3.11'" }, + { name = "jax", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "jaxtyping", version = "0.3.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "lineax", version = "0.1.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "optax", marker = "python_full_version >= '3.11'" }, + { name = "paramax", marker = "python_full_version >= '3.11'" }, + { name = "tqdm", marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4c/37/375042415ab43e26cea0cc3138f29767d09cb2506ed1cbb7d17f118f7112/flowjax-19.1.0.tar.gz", hash = "sha256:0c22ace574ec9e22eac9a1abc2ab84264d5d7634a6b648a522575468cee89bcf", size = 477188, upload-time = "2026-02-21T13:47:38.187Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/27/19/112666688ab1b389e2b5b463848a8b1a34846e2e2a15e0284fa17c7c5a8a/flowjax-19.1.0-py3-none-any.whl", hash = "sha256:9ec9e08eb8e4956ae99d3261228df16096cb705dedec553108a0c22b35b5372d", size = 61388, upload-time = "2026-02-21T13:47:36.365Z" }, ] -provides-extras = ["dev", "viz", "all"] [[package]] name = "fonttools" @@ -840,35 +860,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2c/47/c99d5268f354002ce80f8d029cd9d7d872969da1de8b93d32de4dc56d6f4/fonttools-4.63.0-py3-none-any.whl", hash = "sha256:445af2eab030a16b9171ea8bdda7ebf7d96bda2df88ee182a464252f6e05e20d", size = 1164562, upload-time = "2026-05-14T12:04:29.092Z" }, ] -[[package]] -name = "gravax" -version = "0.1.0" -source = { editable = "../gravax" } -dependencies = [ - { name = "equinox" }, - { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "jax", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "jaxlib", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "jaxstro" }, - { name = "jaxtyping", version = "0.3.7", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "jaxtyping", version = "0.3.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "progenax" }, -] - -[package.metadata] -requires-dist = [ - { name = "equinox", specifier = ">=0.11.0" }, - { name = "jax", specifier = ">=0.4.20" }, - { name = "jaxlib", specifier = ">=0.4.20" }, - { name = "jaxstro" }, - { name = "jaxtyping", specifier = ">=0.2.25" }, - { name = "progenax" }, - { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, - { name = "pytest-cov", marker = "extra == 'dev'" }, -] -provides-extras = ["dev"] - [[package]] name = "h11" version = "0.16.0" @@ -1223,44 +1214,6 @@ requires-dist = [ ] provides-extras = ["dev"] -[[package]] -name = "jaxstroviz" -version = "0.1.0" -source = { editable = "../jaxstroviz" } -dependencies = [ - { name = "fluxax" }, - { name = "gravax" }, - { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "jax", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "jaxlib", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "jaxstro" }, - { name = "matplotlib" }, - { name = "progenax" }, -] - -[package.metadata] -requires-dist = [ - { name = "arviz", marker = "extra == 'experimental'", specifier = ">=0.18" }, - { name = "black", marker = "extra == 'dev'", specifier = ">=23.0" }, - { name = "fluxax", editable = "../fluxax" }, - { name = "gravax", editable = "../gravax" }, - { name = "isort", marker = "extra == 'dev'", specifier = ">=5.12" }, - { name = "jax", specifier = ">=0.4.20" }, - { name = "jaxlib", specifier = ">=0.4.20" }, - { name = "jaxstro", editable = "../jaxstro" }, - { name = "jaxstroviz", extras = ["dev", "seaborn"], marker = "extra == 'all'" }, - { name = "matplotlib", specifier = ">=3.7" }, - { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.0" }, - { name = "progenax", editable = "." }, - { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, - { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0" }, - { name = "scipy", marker = "extra == 'experimental'", specifier = ">=1.11" }, - { name = "seaborn", marker = "extra == 'seaborn'", specifier = ">=0.12" }, - { name = "startrax", marker = "extra == 'tracks'", specifier = ">=0.1.0" }, -] -provides-extras = ["dev", "seaborn", "experimental", "tracks", "all"] - [[package]] name = "jaxtyping" version = "0.3.7" @@ -1491,16 +1444,16 @@ version = "3.10.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "contourpy", version = "1.3.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "cycler" }, - { name = "fonttools" }, - { name = "kiwisolver" }, + { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "cycler", marker = "python_full_version < '3.12'" }, + { name = "fonttools", marker = "python_full_version < '3.12'" }, + { name = "kiwisolver", marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "packaging" }, - { name = "pillow" }, - { name = "pyparsing" }, - { name = "python-dateutil" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "packaging", marker = "python_full_version < '3.12'" }, + { name = "pillow", marker = "python_full_version < '3.12'" }, + { name = "pyparsing", marker = "python_full_version < '3.12'" }, + { name = "python-dateutil", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/63/1b/4be5be87d43d327a0cf4de1a56e86f7f84c89312452406cf122efe2839e6/matplotlib-3.10.9.tar.gz", hash = "sha256:fd66508e8c6877d98e586654b608a0456db8d7e8a546eb1e2600efd957302358", size = 34811233, upload-time = "2026-04-24T00:14:13.539Z" } wheels = [ @@ -1996,6 +1949,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0f/54/68a0978d1ef8502b8492099beaa6e7a0c1b32e3b5d4f677f5810cb08711c/pandas-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b2c95f8bfc1ee412bf482605d7bfd30c12d1d26bd59fdd91efeef1d4718decb1", size = 9466464, upload-time = "2026-05-11T18:54:22.754Z" }, ] +[[package]] +name = "paramax" +version = "0.0.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "equinox" }, + { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "jax", version = "0.10.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "jaxtyping", version = "0.3.7", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "jaxtyping", version = "0.3.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6b/ea/291616b6007b274a677baa984c2886638b4f3a64d4c40943687c3b78ef63/paramax-0.0.5.tar.gz", hash = "sha256:b6710faeb6534f00f44c49f7fff80ebe9364d23636d6e62985ef6cb9b4e60b31", size = 11351, upload-time = "2026-01-23T19:00:46.859Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/f9/57a9a2b706e706cdf89c46e2a3875c9a63b4824c7327d980799f11a49aa2/paramax-0.0.5-py3-none-any.whl", hash = "sha256:c1184d6eadb323588deaaef883c78dbbe4a896376517845b87e19c6108dad707", size = 7757, upload-time = "2026-01-23T19:00:45.273Z" }, +] + [[package]] name = "pillow" version = "12.2.0" @@ -2133,13 +2102,21 @@ dependencies = [ dev = [ { name = "pytest" }, { name = "pytest-cov" }, + { name = "pytest-xdist" }, +] +diagnostics = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] experimental = [ { name = "arviz", version = "0.23.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "arviz", version = "1.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "blackjax", version = "1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "blackjax", version = "1.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "jaxstroviz" }, + { name = "flowjax", version = "19.0.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "flowjax", version = "19.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "optax" }, { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -2151,17 +2128,20 @@ requires-dist = [ { name = "blackjax", marker = "extra == 'experimental'", specifier = ">=1.2,<2" }, { name = "diffrax", specifier = ">=0.4.0" }, { name = "equinox", specifier = ">=0.11.0" }, + { name = "flowjax", marker = "extra == 'experimental'", specifier = ">=17" }, { name = "jax", specifier = ">=0.4.20" }, { name = "jaxlib", specifier = ">=0.4.20" }, { name = "jaxstro", editable = "../jaxstro" }, - { name = "jaxstroviz", marker = "extra == 'experimental'", editable = "../jaxstroviz" }, { name = "jaxtyping", specifier = ">=0.2.25" }, + { name = "numpy", marker = "extra == 'diagnostics'", specifier = ">=1.24" }, { name = "optax", marker = "extra == 'experimental'", specifier = ">=0.2" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0" }, { name = "pytest-cov", marker = "extra == 'dev'" }, + { name = "pytest-xdist", marker = "extra == 'dev'", specifier = ">=3.0" }, + { name = "scipy", marker = "extra == 'diagnostics'", specifier = ">=1.11" }, { name = "scipy", marker = "extra == 'experimental'", specifier = ">=1.11" }, ] -provides-extras = ["dev", "experimental"] +provides-extras = ["dev", "diagnostics", "experimental"] [[package]] name = "pygments" @@ -2213,6 +2193,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload-time = "2026-03-21T20:11:14.438Z" }, ] +[[package]] +name = "pytest-xdist" +version = "3.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "execnet" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791eb921115fca8e44e596a13efeda518b9d845a619450/pytest_xdist-3.8.0.tar.gz", hash = "sha256:7e578125ec9bc6050861aa93f2d59f1d8d085595d6551c2c90b6f4fad8d3a9f1", size = 88069, upload-time = "2025-07-01T13:30:59.346Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -2575,6 +2568,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" }, ] +[[package]] +name = "tqdm" +version = "4.68.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/85/05/0d5260f1f1ca784f4a4a0def9cbe6affe587f5b4025328d446c3d67765f4/tqdm-4.68.2.tar.gz", hash = "sha256:89c230e8dbc67c7615c142487111222f878c77427ea09549960f62389e258add", size = 171923, upload-time = "2026-06-09T13:26:42.539Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/eb/75/1a0392bcc21c44dcdf87b3cf2d137e7829be2c083a1e38d44efca3d57a16/tqdm-4.68.2-py3-none-any.whl", hash = "sha256:d4240441fb5353290b87d6a85968c9decc131a99b8c7faa28269d829de669ede", size = 78578, upload-time = "2026-06-09T13:26:40.731Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0"