docs(research): Ludlow16 JAX concentration feasibility study (#397)

[Jammy2211]

Summary

Phase 1 deliverable for #397 — assess whether the jax.pure_callback wrapping colossus.halo.concentration in autogalaxy/profiles/mass/dark/mcr_util.py can be replaced by a JAX-native implementation.

Verdict: yes, Approach A is viable. A full port of modelLudlow16 — Eisenstein-Hu '98 transfer + Heath '77 growth factor + Einasto gammainc + 200-point Ludlow16 solver — comes to ~330 lines of straight-line JAX in one file. Max relative error in c200c vs colossus over the lensing parameter grid (log M ∈ [10, 14] Msun/h, z ∈ [0.1, 2.5]) is 7.5 × 10⁻⁴. Single-call post-JIT is 0.69 ms vs colossus' 0.83 ms. jax.grad works end-to-end; vmap × 32 is 1.29× faster than colossus serial.

End-to-end propagation through NFWMCRScatterLudlow and cNFWMCRScatterLudlow:

• kappa_s max rel error: 1.07 × 10⁻³
• NFW convergence/deflection per-pixel max rel error: 8.21 × 10⁻⁴
• cNFW deflection per-pixel max rel error: 7.60 × 10⁻⁴

The intrinsic scatter of the Ludlow16 relation itself (σ_log10(c) = 0.13 dex, the dispersion the scatter_sigma parameter marginalises over) is ~350× larger than this offset. Scientifically invisible.

No production code is changed in this PR. All artefacts live under docs/research/nfw_ludlow16_jax/ and serve as the basis for a separate Phase 2 PR that will swap the prototype into mcr_util.py, collapse the JAX/NumPy branches in the two callers, and make colossus an optional dependency.

API Changes

None — internal changes only (research artefact in docs/research/).

Test Plan

• python -m pytest test_autogalaxy/ passes (production code is untouched).
• Reproduce the report numbers per instructions in docs/research/nfw_ludlow16_jax_assessment.md:

  source ~/Code/PyAutoLabs-wt/nfw-jax-port/activate.sh
  cd ~/Code/PyAutoLabs-wt/nfw-jax-port/PyAutoGalaxy
  JAX_ENABLE_X64=1 NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib \
      python docs/research/nfw_ludlow16_jax/validate.py
  

• Confirm docs/research/nfw_ludlow16_jax_assessment.md renders cleanly on GitHub.

Full API Changes (for automation & release notes)

Added

• docs/research/nfw_ludlow16_jax_assessment.md — feasibility report.
• docs/research/nfw_ludlow16_jax/ludlow16_jax.py — Approach A prototype (~330 lines: EH98 transfer, σ(R) via top-hat log-k quadrature, flat-LCDM growth factor via Heath '77 integral, Einasto gammainc mass ratio, 200-point Ludlow16 c-solver).
• docs/research/nfw_ludlow16_jax/validate.py — accuracy comparison vs colossus over the lensing parameter grid.
• docs/research/nfw_ludlow16_jax/bench.py — jax.grad correctness + jax.vmap timing.
• docs/research/nfw_ludlow16_jax/tune.py — (nk, nz) quadrature grid sweep.
• docs/research/nfw_ludlow16_jax/science_check.py — end-to-end NFW / cNFW lensing comparison (helper outputs + convergence + deflection maps).

Removed

None.

Renamed

None.

Changed Signature

None.

Changed Behaviour

None.

Migration

N/A — no production API is touched.

🤖 Generated with Claude Code