panco3: Differentiable Bayesian Inference in JAX

Github: fkeruzore/panco2

panco3 is a from-scratch JAX port of panco2, rebuilding the same forward-modeling approach to Bayesian pressure profile fitting as a fully differentiable, GPU-capable pipeline.

panco2 runs on an ensemble MCMC sampler (emcee), which is CPU-bound and treats the forward model as a black box. Most of the cost of panco2 is in its forward model, which consists in straightforward tensor operations which can be accelerated using compilation and GPUs; and in sampling efficiency, which can be improved using gradient-based workflows. panco3 reimplements the forward model - profile parametrization, beam convolution, filtering, likelihood evaluation - end-to-end in JAX, making every step JIT-compilable, GPU-accelerated, and differentiable. Sampling moves from ensemble MCMC to gradient-based MCMC/HMC via BlackJAX, which uses those gradients to propose more efficient moves through parameter space.

Most of the panco3 codebase was generated with Claude Code as an agent-assisted development exercise, with the physics, validation, and API design directed and reviewed throughout. panco3 is in active development.