:py:mod:`nessai.flows.base`
===========================

.. py:module:: nessai.flows.base

.. autoapi-nested-parse::

   Base objects for implementing normalising flows.

   ..
       !! processed by numpydoc !!


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

Classes
~~~~~~~

.. autoapisummary::

   nessai.flows.base.BaseFlow
   nessai.flows.base.NFlow




.. py:class:: BaseFlow

   Bases: :py:obj:`torch.nn.Module`, :py:obj:`abc.ABC`

   
   Base class for all normalising flows.

   If implementing flows using distributions and transforms see NFlow.















   ..
       !! processed by numpydoc !!
   .. py:attribute:: device
      

      
















      ..
          !! processed by numpydoc !!

   .. py:method:: to(self, device)

      
      Wrapper that stores the device before moving the flow
















      ..
          !! processed by numpydoc !!

   .. py:method:: forward(self, x, context=None)
      :abstractmethod:

      
      Apply the forward transformation and return samples in the
      latent space and the log-Jacobian determinant.



      :Returns:

          :obj:`torch.Tensor`
              Tensor of samples in the latent space

          :obj:`torch.Tensor`
              Tensor of log determinants of the Jacobian of the forward
              transformation













      ..
          !! processed by numpydoc !!

   .. py:method:: inverse(self, z, context=None)
      :abstractmethod:

      
      Apply the inverse transformation and return samples in the
      data space and the log-Jacobian determinant.



      :Returns:

          :obj:`torch.Tensor`
              Tensor of samples in the data space

          :obj:`torch.Tensor`
              Tensor of log determinants of the Jacobian of the forward
              transformation













      ..
          !! processed by numpydoc !!

   .. py:method:: sample(self, n, context=None)
      :abstractmethod:

      
      Generate n samples in the data space



      :Returns:

          :obj:`torch.Tensor`
              Tensor of samples in the data space













      ..
          !! processed by numpydoc !!

   .. py:method:: log_prob(self, x, context=None)
      :abstractmethod:

      
      Compute the log probability for a set of samples in the data space



      :Returns:

          :obj:`torch.Tensor`
              Tensor of log probabilities of the samples













      ..
          !! processed by numpydoc !!

   .. py:method:: base_distribution_log_prob(self, z, context=None)
      :abstractmethod:

      
      Computes the log probability of samples in the latent for
      the base distribution in the flow.



      :Returns:

          :obj:`torch.Tensor`
              Tensor of log probabilities of the latent samples













      ..
          !! processed by numpydoc !!

   .. py:method:: forward_and_log_prob(self, x, context=None)
      :abstractmethod:

      
      Apply the forward transformation and compute the log probability
      of each sample



      :Returns:

          :obj:`torch.Tensor`
              Tensor of samples in the latent space

          :obj:`torch.Tensor`
              Tensor of log probabilities of the samples













      ..
          !! processed by numpydoc !!

   .. py:method:: sample_and_log_prob(self, n, context=None)
      :abstractmethod:

      
      Generates samples from the flow, together with their log probabilities
      in the data space ``log p(x) = log p(z) + log|J|``.

      For flows, this is more efficient that calling ``sample`` and
      ``log_prob`` separately.


      :Returns:

          :obj:`torch.Tensor`
              Tensor of samples in the data space

          :obj:`torch.Tensor`
              Tensor of log probabilities of the samples













      ..
          !! processed by numpydoc !!


.. py:class:: NFlow(transform, distribution)

   Bases: :py:obj:`BaseFlow`

   
   Base class for flow objects implemented according to outline in nflows.
   These take an instance of a transform and a distribution from nflows.

   This replaces `Flow` from nflows. It removes the context and includes
   additional methods which are called in FlowModel.

   :Parameters:

       **transform** : :obj: `nflows.transforms.Transform`
           Object that applies the transformation, must have`forward` and
           `inverse` methods. See nflows for more details.

       **distribution** : :obj: `nflows.distributions.Distribution`
           Object the serves as the base distribution used when sampling
           and computing the log probability. Must have `log_prob` and
           `sample` methods. See nflows for details














   ..
       !! processed by numpydoc !!
   .. py:method:: forward(self, x, context=None)

      
      Apply the forward transformation and return samples in the latent
      space and log-Jacobian determinant.
















      ..
          !! processed by numpydoc !!

   .. py:method:: inverse(self, z, context=None)

      
      Apply the inverse transformation and return samples in the
      data space and log-Jacobian determinant (not log probability).
















      ..
          !! processed by numpydoc !!

   .. py:method:: sample(self, num_samples, context=None)

      
      Produces N samples in the data space by drawing from the base
      distribution and the applying the inverse transform.

      Does NOT need to be specified by the user















      ..
          !! processed by numpydoc !!

   .. py:method:: log_prob(self, inputs, context=None)

      
      Computes the log probability of the inputs samples by apply the
      transform.

      Does NOT need to specified by the user















      ..
          !! processed by numpydoc !!

   .. py:method:: base_distribution_log_prob(self, z, context=None)

      
      Computes the log probability of samples in the latent for
      the base distribution in the flow.
















      ..
          !! processed by numpydoc !!

   .. py:method:: forward_and_log_prob(self, x, context=None)

      
      Apply the forward transformation and compute the log probability
      of each sample



      :Returns:

          :obj:`torch.Tensor`
              Tensor of samples in the latent space

          :obj:`torch.Tensor`
              Tensor of log probabilities of the samples













      ..
          !! processed by numpydoc !!

   .. py:method:: sample_and_log_prob(self, N, context=None)

      
      Generates samples from the flow, together with their log probabilities
      in the data space ``log p(x) = log p(z) + log|J|``.

      For flows, this is more efficient that calling ``sample`` and
      ``log_prob`` separately.















      ..
          !! processed by numpydoc !!