nessai.flows.base

Base objects for implementing normalising flows.

Module Contents

Classes

BaseFlow	Base class for all normalising flows.
NFlow	Base class for flow objects implemented according to outline in nflows.

class nessai.flows.base.BaseFlow

Bases: torch.nn.Module, abc.ABC

Base class for all normalising flows.

If implementing flows using distributions and transforms see NFlow.

device

to(self, device)

Wrapper that stores the device before moving the flow

abstract forward(self, x, context=None)

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

Returns

torch.Tensor

Tensor of samples in the latent space

torch.Tensor

Tensor of log determinants of the Jacobian of the forward transformation

abstract inverse(self, z, context=None)

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

Returns

torch.Tensor

Tensor of samples in the data space

torch.Tensor

Tensor of log determinants of the Jacobian of the forward transformation

abstract sample(self, n, context=None)

Generate n samples in the data space

Returns

torch.Tensor

Tensor of samples in the data space

abstract log_prob(self, x, context=None)

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

Returns

torch.Tensor

Tensor of log probabilities of the samples

abstract base_distribution_log_prob(self, z, context=None)

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

Returns

torch.Tensor

Tensor of log probabilities of the latent samples

abstract forward_and_log_prob(self, x, context=None)

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

Returns

torch.Tensor

Tensor of samples in the latent space

torch.Tensor

Tensor of log probabilities of the samples

abstract 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.

Returns

torch.Tensor

Tensor of samples in the data space

torch.Tensor

Tensor of log probabilities of the samples

class nessai.flows.base.NFlow(transform, distribution)

Bases: 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

forward(self, x, context=None)

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

inverse(self, z, context=None)

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

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

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

base_distribution_log_prob(self, z, context=None)

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

forward_and_log_prob(self, x, context=None)

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

Returns

torch.Tensor

Tensor of samples in the latent space

torch.Tensor

Tensor of log probabilities of the samples

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.