nessai.proposal.augmented
Augmented version of FlowProposal.
Module Contents
Classes
Version of FlowProposal that uses AugmentedFlows. |
- class nessai.proposal.augmented.AugmentedFlowProposal(model, augment_dims=1, generate_augment='gaussian', marginalise_augment=False, n_marg=50, **kwargs)
Bases:
nessai.proposal.flowproposal.FlowProposalVersion of FlowProposal that uses AugmentedFlows.
Augmented normalising flows were proposed in: https://arxiv.org/abs/2002.07101 and add auxiliary parameters to the inputs of the flow which are drawn from a Gaussian. This improves the flows’ ability to learn multimodal distribution.
- Parameters:
- model
nessai.model.Model User defined model
- augment_dimsint
Number of augment parameters to add to the inputs
- generate_augment{‘gaussian’, ‘zeroes’, ‘zeros’}, optional
Method used when computing the radius of the latent contour.
- marginalise_augmentbool, optional
Use the marginalised likelihood when performing rejection sampling. Adds significant computation cost.
- n_margint, optional
Number of samples to use when approximating the marginalised likelihood.
- model
- set_rescaling()
Configure the rescaling.
Calls the method from the parent class first and then adds the auxiliary parameters.
- update_flow_config()
Update the flow configuration dictionary
- augmented_prior(x)
Log Gaussian for augmented variables.
If self.marginalise_augment is True, log_prior is 0.
- log_prior(x)
Compute the prior probability in the non-prime space.
- x_prime_log_prior(x)
Compute prior probability in the prime space.
- backward_pass(z, rescale=True)
A backwards pass from the model (latent -> real)
- Parameters:
- zarray_like
Structured array of points in the latent space
- rescalebool, optional (True)
Apply inverse rescaling function
- Returns:
- xarray_like
Samples in the latent space
- log_probarray_like
Log probabilities corresponding to each sample (including the Jacobian)