:py:mod:`nessai.proposal.augmented`
===================================

.. py:module:: nessai.proposal.augmented

.. autoapi-nested-parse::

   Augmented version of FlowProposal.

   ..
       !! processed by numpydoc !!


Module Contents
---------------

Classes
~~~~~~~

.. autoapisummary::

   nessai.proposal.augmented.AugmentedFlowProposal




Attributes
~~~~~~~~~~

.. autoapisummary::

   nessai.proposal.augmented.logger


.. py:data:: logger
   

   
















   ..
       !! processed by numpydoc !!

.. py:class:: AugmentedFlowProposal(model, augment_dims=1, generate_augment='gaussian', marginalise_augment=False, n_marg=50, **kwargs)

   Bases: :py:obj:`nessai.proposal.flowproposal.FlowProposal`

   
   Version 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** : :obj:`nessai.model.Model`
           User defined model

       **augment_dims** : int
           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_augment** : bool, optional
           Use the marginalised likelihood when performing rejection sampling.
           Adds significant computation cost.

       **n_marg** : int, optional
           Number of samples to use when approximating the marginalised
           likelihood.














   ..
       !! processed by numpydoc !!
   .. py:method:: set_rescaling(self)

      
      Configure the rescaling.

      Calls the method from the parent class first and then adds the
      auxiliary parameters.















      ..
          !! processed by numpydoc !!

   .. py:method:: update_flow_config(self)

      
      Update the flow configuration dictionary
















      ..
          !! processed by numpydoc !!

   .. py:method:: augmented_prior(self, x)

      
      Log Gaussian for augmented variables.

      If self.marginalise_augment is True, log_prior is 0.















      ..
          !! processed by numpydoc !!

   .. py:method:: log_prior(self, x)

      
      Compute the prior probability in the non-prime space.
















      ..
          !! processed by numpydoc !!

   .. py:method:: x_prime_log_prior(self, x)

      
      Compute prior probability in the prime space.
















      ..
          !! processed by numpydoc !!

   .. py:method:: backward_pass(self, z, rescale=True)

      
      A backwards pass from the model (latent -> real)


      :Parameters:

          **z** : array_like
              Structured array of points in the latent space

          **rescale** : bool, optional (True)
              Apply inverse rescaling function

      :Returns:

          **x** : array_like
              Samples in the latent space

          **log_prob** : array_like
              Log probabilities corresponding to each sample (including the
              Jacobian)













      ..
          !! processed by numpydoc !!