nanograv / pint_pal

A long-lived repository for NANOGrav Pulsar Timing workflows and analysis.

MIT License

6 stars 19 forks source link

WIP: Noise updates #77

Open jeremy-baier opened 2 months ago

jeremy-baier commented 2 months ago

Drafting a PR for 1) chromatic noise model functionality 2) discovery likelihood 3) enterprise extensions Gibbs sampler. Also taps into la_forge functionality for post processing.

Chromatic noise models

Adds dmgp, and chrom gp, and solar wind to the par files post noise.
Enables user to run dmgp/chromgp in the model_noise function.

Discovery Likelihood

adds functionality to interface with the discovery likelihood
options 3 different types of noise samplers in numpyro (HMC, NUTS, HMC-GIBBS)

Gibbs Sampler

Still a WIP, but using this PR in enterprise extensions.
Modifies Nima Lal’s Gibbs sampler, based on van Haasteren & Vallisneri 2014, enabling DM and Chromatic models to be sampled using a free spec psd.
future functionality will add Fourier coefficients into the timing model via the WaveX and DMWaveX and CMWaveX

Requirements

La_forge for noise post-processing
discovery need access
numpyro: pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html (for GPU)
jax: pip install -U "jax[cuda12]” (for GPU)
enterprise_extensions gibbs sampler: PR here pip install git+https://github.com/jeremy-baier/enterprise_extensions@jgb-chromatic-gibbs

TO DO

Discovery ————

[ ] solar wind !!!
[ ] fix hmc-gibbs
[ ] speed test everything
[ ] figure out optimal sampler parameters for HMC, NUTs

Enterprise Extensions Gibbs ——————————————

[ ] un-hard-code simultaneous use of dm, chrom, red noises in Gibbs
[ ] put free specs into model as fourier coefficients via DMwaveX CMwaveX
[ ] fix issues with CMWaveX ( i think this is just an import thing based off version of pint)
[ ] add option to free up the chromatic index with a PTMCMC step

Misc ——

[ ] interface with the auto-creation of diagnostic plots — add free spec and/or PL plots
[ ] add empirical dists for WN + chrom idx
[ ] test run PTMCMC vs e_e.gibbs_sampling vs discovery
[ ] add Nitu+ solar wind
[ ] test solar wind gp?

jeremy-baier commented 2 months ago

moving to-do list higher up