Sampler Additions and PSR Jump Proposals

[Hazboun6]

This PR combines a number of small changes to the samplers that allow them to be tweaked by users if they wish, as well as adding in 3 new jump proposals and a new function for parameter groupings. Many of these additions were developed for full PTA advanced noise model sampling where pulsars now have upwards of 10+ new parameters pertaining to their noise models, but some of these methods should be more widely useful. The user interface has not changed at all. The changes either need to be implemented by the user (e.g. making pulsar noise parameter groups) or will happen through the course of the usual interface (e.g. chromatic GP prior draw jump proposals will be added if chrom_gp is in the list of signal names for a HyperModel run.)

The small tweaks include:

• Changed one of the checks for signal names to incorporate a name change for common red process GW searches.
• Both the HyperModel and the usual sampler have self.jp=jp added to the code so that the JumpProposal class is available to the user. This allows one to add more jump proposals to the proposal scheme. These can be new proposals or more of the same ones. For instance we have found that weighting the empirical distributions more strongly helps with sampling. Since the sampler setup weights all proposals the same, this allows a user to put more of their preferred jump proposals in the cycle.
• Both the HyperModel and the usual sampler look for a cov.npy file in the outdir before making a generic one. This helps with sampling of resumed chains.
• The usual sampler allows the user to dictate the list of parameter index groups that will be used in the SCAM, AM and DE jumps. If no list of groups is set manually the sampler.get_parameter_groups() method is used. This option was already available in HyperModel.setup_sampler and so homogenizes the code interface.
• A sampler.get_psr_groups() method has been added which allows users to retrieve lists of parameter indices that combine all parameters with a given pulsar's name in the parameter name. This allows for covariance jumps in all of a pulsar's parameters at one.

Below is an example of how one might use some of this new functionality:

#Build my own set of groups, first by using the usual function.
groups = sampler.get_parameter_groups(pta)
#Then by extending the groups with pulsar-centered model groups
groups.extend(sampler.get_psr_groups(pta))

#Use my user-made groups in the call to setup the sampler
Sampler = sampler.setup_sampler(pta, outdir=/PATH/TO/OUTDIR, resume=True, 
empirical_distr = EMP/DIST/PATH, groups=groups)

#Use the Sampler.jp jump proposal class to add more JPs to the cycyle before starting the sampler. 
#These are "new" and are not in any automatic calls within `setup_sampler`.
Sampler.addProposalToCycle(Sampler.jp.draw_from_psr_empirical_distr, 50)
Sampler.addProposalToCycle(Sampler.jp.draw_from_psr_prior, 10)

#These are already called in `setup_sampler` but I would like to add more of them to the cycle.
Sampler.addProposalToCycle(Sampler.jp.draw_from_empirical_distr, 50)
Sampler.addProposalToCycle(Sampler.jp.draw_from_red_prior, 30)
Sampler.addProposalToCycle(Sampler.jp.draw_from_dm_gp_prior, 30)
Sampler.addProposalToCycle(Sampler.jp.draw_from_chrom_gp_prior, 30)
Sampler.addProposalToCycle(Sampler.jp.draw_from_dmexpcusp_prior, 30)

New Jump Proposals:

• A new draw from chromatic GP prior jump proposal has been added. This is similar to the many other prior draw JPs we have added. These all could get reorganized, but for now I would like to keep track of this one as it is important for advanced noise modeling. It is automatically added to hyper model runs that contain chrom_gp type signals.
• A new pulsar prior draw jump proposal has been added. This is a multi-parameter draw that draws new values for all of the parameters in a single pulsar's noise model, i.e. , all parameters with a given pulsar's name in the parameter name. This JP seems to only be useful at the beginning of burn in, but can help decrease burn in time by drawing many parameters at once.
• A new pulsar empirical distribution draw jump proposal has been added. This is a multi-parameter draw that draws new values from a set of empirical distributions that exist for a single pulsar noise model, i.e. , all parameters with a given pulsar's name in the parameter name. This is not an n-dimensional empirical distribution. It just draws new parameters from the 1D and 2D empirical distributions. These jumps are actually fairly often accepted (much more than pulsar prior JP) since they come from the empirical distributions. In fact in some test runs they basically have the same acceptance rate as the normal empirical distribution draws, except that they help much more extensively with mixing when there are many parameters for a given pulsar.

Neither of these last two jump proposals are currently added automatically by the code. They would replicate already existing jump proposals in a "normal" analysis with only achromatic red noise models for the individual pulsars and some common signal. However I've currently left it to the user to add them by hand.

I have not added tests for these proposals, as there are no tests yet for any of our jump proposals. We should discuss the best way to add in tests for all of our jump proposals.

[Hazboun6]

I should point out that this PR does touch on some of the topics addressed in Issue #19, started by @paulthebaker. There is still some work to be done regarding efficient setup of these various proposals though.

[codecov-commenter]

Codecov Report

Merging #127 (e898227) into master (ac88018) will decrease coverage by 0.28%.
The diff coverage is 7.89%.

@@            Coverage Diff             @@
##           master     #127      +/-   ##
==========================================
- Coverage   27.98%   27.69%   -0.29%     
==========================================
  Files          19       19              
  Lines        3277     3416     +139     
==========================================
+ Hits          917      946      +29     
- Misses       2360     2470     +110     

Impacted Files	Coverage Δ
enterprise_extensions/hypermodel.py	0.00% <0.00%> (ø)
enterprise_extensions/sampler.py	20.14% <8.82%> (-1.46%)	⬇️
...rprise_extensions/frequentist/optimal_statistic.py	80.45% <0.00%> (-2.78%)	⬇️
enterprise_extensions/blocks.py	41.72% <0.00%> (ø)
enterprise_extensions/model_orfs.py	32.57% <0.00%> (+2.72%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update ac88018...e898227. Read the comment docs.

[Hazboun6]

Here is a plot of acceptance rates for various jump proposals in a recent solar wind analysis. Note that the pulsar empirical distributions track closely with the regular empirical distribution draws. Here it would be pulling 4 parameters at a time as each pulsar has a DM GP and the red noise parameters. The pulsar prior JP does not do so well here, but it does better in the full PTA advanced noise analysis.

[paulthebaker]

It makes sense that pulsar prior draw is very unlikely to be accepted after the first couple of steps. If you attempt to move many parameters at once, the probability that at least one will end up in a bad spot is pretty high. The code works, but I would guess that it's almost never worth including that particular draw.

[Hazboun6]

@paulthebaker The acceptance rate for that draw is pretty reasonable (18%) in the advanced noise modeling full PTA run. I think it might be because many of the pulsars only have power law red noise so the draw is pulling a sample of red noise parameters only. It is lower than the red noise prior draw, but that makes sense since when it was pulling from a normal pulsar it would be jumping in 2 parameters, and when it was jumping in an advanced noise pulsar it would not be accepted very often for the reasons you give above.

[stevertaylor]

Fantastic modification that we've needed for a while. Thanks Jeff.

[stevertaylor]

Does this need to be even more general, @Hazboun6? This will work if the ORF is None such that the process is a crn, but in model_general the processes get named for their ORF behavior: https://github.com/nanograv/enterprise_extensions/blob/master/enterprise_extensions/models.py#L788. So a H&D process would have gw_hd_log10_A.

[Hazboun6]

You make a good point @stevertaylor . I've changed the jump proposal to cycle through any signals with "gw" and "log10_A" in the parameter names and now sampler.setup_sampler and HyperModel.setup_ampler do searches for the same string pieces in order to add the jump proposal.

See lines 369 and 968 of sampler.py and line 256 of hypermodel.py.