nessai.flowmodel.base

Object and functions to handle training the normalising flow.

Module Contents

Classes

FlowModel

Object that contains the normalising flows and handles training and data

class nessai.flowmodel.base.FlowModel(config=None, output=None)

Object that contains the normalising flows and handles training and data pre-processing.

Does NOT use structured arrays for live points, Proposal object should act as the interface between structured used by the sampler and unstructured arrays of live points used for training.

Parameters:

configdict, optional

Configuration used for the normalising flow. If None, default values are used.

outputstr, optional

Path for output, this includes weights files and the loss plot. If not specified, the current working directory is used.

save_input(config, output_file=None)

Save the dictionary used as an inputs as a JSON file in the output directory.

Parameters:

configdict

Dictionary to save.

output_filestr, optional

File to save the config to.

setup_from_input_dict(config)

Setup the trainer from a dictionary, all keys in the dictionary are added as methods to the object. Input is automatically saved.

Parameters:

configdict

Dictionary with parameters that are used to update the defaults.

update_mask()

Method to update the ask upon calling initialise

By default the mask is left unchanged.

get_optimiser(optimiser='adam', **kwargs)

Get the optimiser and ensure it is always correctly initialised.

Returns:

torch.optim.Adam

Instance of the Adam optimiser from torch.optim

initialise()

Initialise the model and optimiser.

This includes:

• Updating the model configuration

• Initialising the normalising flow

• Initialising the optimiser

• Configuring the inference device

move_to(device, update_default=False)

Move the flow to a different device.

Parameters:

devicestr

Device to move flow to.

update_defaultbool, optional

If True, the default device for the flow (and data) will be updated.

static check_batch_size(x, batch_size, min_fraction=0.1)

Check that the batch size is valid.

Tries to ensure that the last batch is at least a minimum fraction of the size of the batch size.

Parameters:

xnumpy.ndarray

Training data.

batch_sizeint

The user-specified batch size

min_fractionfloat

Fraction of user-specified batch size to use a lower limit.

prep_data(samples, val_size, batch_size, weights=None, use_dataloader=False, conditional=None)

Prep data and return dataloaders for training

Parameters:

samplesarray_like

Array of samples to split in to training and validation.

val_sizefloat

Float between 0 and 1 that defines the fraction of data used for validation.

batch_sizeint

Batch size used when constructing dataloaders.

use_dataloaderbool, optional

If True data is returned in a dataloader else a tensor is returned.

weightsarray_like, optional

Array of weights for each sample, weights will used when computing the loss.

Returns:

train_data, val_data

Training and validation data as either a tensor or dataloader

end_iteration()

Calls any functions that should be applied at the end of the iteration.

This functions is called after the flow has been updated on all batches of data but before the validation step.

Calls nessai.flows.base.BaseFlow.end_iteration()

finalise()

Method to finalise the flow before using it for inference.

train(samples, weights=None, conditional=None, max_epochs=None, patience=None, output=None, val_size=None, plot=True)

Train the flow on a set of samples.

Allows for training parameters to specified instead of those given in initial config.

Parameters:

samplesndarray

Unstructured numpy array containing data to train on

weightsarray_like

Array of weights to use with the weight KL when computing the loss.

max_epochsint, optional

Maximum number of epochs that is used instead of value in the configuration.

patienceint, optional

Patience in number of epochs that is used instead of value in the configuration.

val_sizefloat, optional

Fraction of the samples to use for validation

outputstr, optional

Path to output directory that is used instead of the path specified when the object is initialised

plotbool, optional

Boolean to enable or disable plotting the loss function

Returns:

historydict

Dictionary that contains the training and validation losses.

save_weights(weights_file)

Save the weights file. If the file already exists move it to <weights_file>.old and then save the file.

Parameters:

weights_filestr

Path to to file to save weights. Recommended file type is .pt.

load_weights(weights_file)

Load weights for the model and initialises the model if it is not initialised. The weights_file attribute is also updated.

Model is loaded in evaluation mode (model.eval())

Parameters:

weights_filesstr

Path to weights file

reload_weights(weights_file)

Tries to the load the weights file and if not, tries to load the weights file stored internally.

Parameters:

weights_filesstr

Path to weights file

reset_model(weights=True, permutations=False)

Reset the weights of the model and optimiser.

Parameters:

weightsbool, optional

If true the model weights are reset.

permutationsbool, optional

If true any permutations (linear transforms) are reset.

numpy_array_to_tensor(array: numpy.ndarray, /) → torch.Tensor

Convert a numpy array to a tensor and move it to the device

forward_and_log_prob(x: numpy.ndarray, conditional: numpy.ndarray | None = None) → Tuple[numpy.ndarray, numpy.ndarray]

Forward pass through the model and return the samples in the latent space with their log probabilities

Parameters:

xndarray

Array of samples

Returns:

zndarray

Samples in the latent space

log_probndarray

Log probabilities for each samples

log_prob(x: numpy.ndarray, conditional: numpy.ndarray | None = None) → numpy.ndarray

Compute the log-probability of a sample.

Parameters:

xndarray

Array of samples in the X-prime space.

Returns:

ndarray

Array of log-probabilities.

sample(n: int = 1, conditional: numpy.ndarray | None = None) → numpy.ndarray

Sample from the flow.

Parameters:

nint

Number of samples to draw

conditionalnumpy.ndarray

Array of conditional inputs

Returns:

numpy.ndarray

Array of samples

sample_latent_distribution(n: int = 1) → numpy.ndarray

Sample from the latent distribution.

Parameters:

nint

Number of samples to draw

Returns:

numpy.ndarray

Array of samples

sample_and_log_prob(N=1, z=None, alt_dist=None, conditional=None)

Generate samples from samples drawn from the base distribution or and alternative distribution from provided latent samples

Parameters:

Nint, optional

Number of samples to draw if z is not specified

zndarray, optional

Array of latent samples to map the the data space, if alt_dist is not specified they are assumed to be drawn from the base distribution of the flow.

alt_distglasflow.nflows.distribution.Distribution

Distribution object from which the latent samples z were drawn from. Must have a log_prob method that accepts an instance of torch.Tensor

Returns:

samplesndarray

Array containing samples in the latent space.

log_probndarray

Array containing the log probability that corresponds to each sample.