WIP: Noise updates

[jeremy-baier]

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]

moving to-do list higher up

[jeremy-baier]

@rossjjennings rebased to y’all’s changes and switched the base branch from NG20 to main.
let me know if i have bungled the git history somehow in the process.

[rossjjennings]

OK, great! It doesn't look like the history is bungled to me. The tests are still failing, which seems to be because you've added import statements for arviz and la_forge without having declared them as dependencies. You might want to try adding those two packages to the list of dependencies in pyproject.toml, and see if the tests pass then.

[rossjjennings]

You're importing some other new packages too (jax, numpyro, discovery), but those shouldn't be a problem for running the tests because they're imported inside a try/except block. For Discovery, there's not a great way to deal with it since the code is not yet public, but JAX and NumPyro could be added as dependencies in the normal way.

[tcromartie]

I forgot to comment, sorry! This works for me for enterprise-only noise runs with properly updated config files. I am stalled in testing the discovery implementation.

[JPGlaser]

Just a quick note.

[JPGlaser]

Something to note is to check capitalization for "Enterprise"

~ Joe G.

[jeremy-baier]

thanks for pointing this out!
every instance of enterprise is now lower case.

[tcromartie]

This goes in the notebook:

if run_noise_analysis or use_existing_noise_dir:
    mo_new = copy.deepcopy(mo)
    lu.remove_noise(mo_new)
    nu.model_noise(mo_new, to,
                   using_wideband = using_wideband,
                   run_noise_analysis = True,
                   model_kwargs=tc.config['noise_run']['model'],
                   sampler_kwargs=tc.config['noise_run']['inference'],
    )

[tcromartie]

config file:
for discovery:

noise_run:
  model:
      inc_rn: true
      inc_dmgp: false
      inc_chromgp: false
  inference:
     likelihood: discovery
     sampler: HMC
     num_warmup: 100
     num_samples: 1000
     num_chain: 3
     chain_method: parallel

for enterprise only:

     likelihood: Enterprise
     sampler: PTMCMCSampler

[JPGlaser]

Side note @jeremy-baier we also need to check the hard-coded paths that are inside noise_utils.py and a few other places, as brought up by: #43

~ Joe G.

[jeremy-baier]

config file: for discovery:

noise_run:
  model:
      inc_rn: true
      inc_dmgp: false
      inc_chromgp: false
  inference:
     likelihood: discovery
     sampler: HMC
     num_warmup: 100
     num_samples: 1000
     num_chain: 3
     chain_method: parallel

for enterprise only:

     likelihood: Enterprise
     sampler: PTMCMCSampler

See Joe’s comment. ‘Enterprise’ —> ‘enterprise’ for consistency

[jeremy-baier]

hmmm… looks like the tests are still failing. complaining about not being able to load noise chains. i don’t see any test noise chains in the pint pal directory. does anyone know where it’s looking to grab these chains from ?

[jeremy-baier]

Side note @jeremy-baier we also need to check the hard-coded paths that are inside noise_utils.py and a few other places, as brought up by: #43

~ Joe G.

@JPGlaser looking at ross’s comment on this issue. i switched things over so that pint_pal is now using laforge to read in and analyze noise chains. for laforge, you just need to point to the directory so hopefully there is less hardcoding there.
not sure about some of the other places though.

[jeremy-baier]

hmmm… looks like the tests are still failing. complaining about not being able to load noise chains. i don’t see any test noise chains in the pint pal directory. does anyone know where it’s looking to grab these chains from ?

@rossjjennings if you have a chance tomorrow, could you help me figure this one out ?

[rossjjennings]

@jeremy-baier To quote what I said in Slack:

I think what happens with the tests currently is that nu.add_noise_to_model() raises an OSError when it can't find the chains, which is caught in the noise modeling cell, so everything else in that cell gets skipped. In your version it raises a ValueError instead, which doesn't get caught and makes the test fail.

So a simple thing to do would be to change except OSError as e: in the notebook to except ValueError as e:, or except (OSError, ValueError) as e: if you want to catch both.