nessai.flowmodel

Object and functions to handle training the normalising flow.

Module Contents

Classes

FlowModel

Object that contains the normalsing flows and handles training and data

Functions

update_config(d)

Update the configuration dictionary to include the defaults.

Attributes

logger

nessai.flowmodel.logger
nessai.flowmodel.update_config(d)

Update the configuration dictionary to include the defaults.

Parameters
ddict

Dictionary with configuration

Returns
dict

Dictionary with updated default configuration

Notes

The default configuration is:

lr=0.001
annealing=False
batch_size=100
val_size=0.1
max_epochs=500
patience=20
noise_scale=0.0,
use_dataloader=False,
optimiser='adam',
optimiser_kwargs={}
model_config=default_model

where default model is:

n_neurons=32
n_blocks=4
n_layers=2
ftype='RealNVP'
device_tag='cpu',
flow=None,
inference_device_tag=None,
kwargs={batch_norm_between_layers=True, linear_transform='lu'}

The kwargs can contain any additional keyword arguments that are specific to the type of flow being used.

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

Object that contains the normalsing 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.

model_config
save_input(self, 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(self, 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(self)

Method to update the ask upon calling initialise

By default the mask is left unchanged.

get_optimiser(self, 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(self)

Initialise the model and optimiser.

This includes:

  • Updating the model configuration

  • Initialising the normalising flow

  • Initialising the optimiser

  • Configuring the inference device

move_to(self, 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.

prep_data(self, samples, val_size, batch_size, use_dataloader=False)

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.

Returns
train_data, val_data :

Training and validation data as either a tensor or dataloader

train(self, samples, 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

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

save_weights(self, 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(self, 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(self, 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(self, 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.

forward_and_log_prob(self, x)

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

sample_and_log_prob(self, N=1, z=None, alt_dist=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_distnflows.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.