nessai.flows.base ================= .. py:module:: nessai.flows.base .. autoapi-nested-parse:: Base objects for implementing normalising flows. .. !! processed by numpydoc !! Classes ------- .. autoapisummary:: nessai.flows.base.BaseFlow nessai.flows.base.NFlow Module Contents --------------- .. py:class:: BaseFlow(*args: Any, **kwargs: Any) 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:method:: to(device) Wrapper that stores the device before moving the flow .. !! processed by numpydoc !! .. py:method:: forward(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(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(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(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:: sample_latent_distribution(n, context=None) :abstractmethod: Sample from the latent distribution. .. !! processed by numpydoc !! .. py:method:: base_distribution_log_prob(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(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(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:method:: finalise() Finalise the flow after training. Will be called after training the flow and loading the best weights. For example, can be used to finalise the Monte Carlo estimate of the normalising constant used in a LARS based flow. By default does nothing and should be implemented by the user. .. !! processed by numpydoc !! .. py:method:: end_iteration() Update the model at the end of an iteration. Will be called between training and validation. By default does nothing and should be overridden by an class that inherit from this class. .. !! processed by numpydoc !! .. py:method:: freeze_transform() :abstractmethod: Freeze the transform part of the flow. Must be implemented by the child class. .. !! processed by numpydoc !! .. py:method:: unfreeze_transform() :abstractmethod: Unfreeze the transform part of the flow. Must be implemented by the child class. .. !! processed by numpydoc !! .. py:class:: NFlow(transform, distribution) Bases: :py:obj:`BaseFlow` Base class for flow objects from glasflow.nflows. This replaces `Flow` from glasflow.nflows. It includes additional methods which are called in FlowModel. :Parameters: **transform** : :obj: `glasflow.nflows.transforms.Transform` Object that applies the transformation, must have`forward` and `inverse` methods. See glasflow.nflows for more details. **distribution** : :obj: `glasflow.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 glasflow.nflows for details .. !! processed by numpydoc !! .. py:method:: forward(x, context=None) Apply the forward transformation and return samples in the latent space and log-Jacobian determinant. .. !! processed by numpydoc !! .. py:method:: inverse(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(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(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:: sample_latent_distribution(n, context=None) Sample from the latent distribution. .. !! processed by numpydoc !! .. py:method:: base_distribution_log_prob(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(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(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 !! .. py:method:: finalise() Finalise the flow after training. Checks if the base distribution or transform have finalise methods and calls them. .. !! processed by numpydoc !! .. py:method:: end_iteration() Update the model at the end of an iteration. Will be called between training and validation. .. !! processed by numpydoc !! .. py:method:: freeze_transform() Freeze the transform part of the flow .. !! processed by numpydoc !! .. py:method:: unfreeze_transform() Unfreeze the transform part of the flow .. !! processed by numpydoc !!