mbank.flow

A module to train and use a normalizing flow model to sample from the probability distribution induced by the metric. The probability distribution is given by:

\[p(\theta) \propto \sqrt{|\text{det}M(\theta)|}\]

A trained flow will learn how to sample from this model.

Although the module can work as a stand-alone, the model is specialized to the bank generation and it is employed extensively in by the tiling handler (mbank.handlers.tiling_handler) to:

interpolate the metric
sample from the tiling

For most standard application, the user can interface with the normalizing flow through the tiling, without using mbank.flow

The module has two submodules:

mbank.flow.flowmodel: defines the normalizing flow model and the functions to train it
mbank.flow.utils: gathers some utitilies to plot and validate the performance of the flow

The module relies on pytorch for a general Machine Learning library and on glasflow for the flow implementation.

The module mbank.flow.utils depends on seaborn, pandas and imageio. To keep the distribution light, they are not among the dependencies of the package. If you want to use the full functionalities, you will need to install them manually.

pip install pandas, seaborn, imageio

Future realease of the code may remove the dependencies on this pacakges.

mbank.flow.flowmodel

Implements the basic normalizing flow model, useful for sampling from the Binary Black Hole parameter space. It requires torch and glasflow.nflows, which are not among the mbank dependencies.

class GW_Flow(transform, distribution, has_constant=True)[source]

A Normalizing flow suitable to reproduce the uniform distribution over the parameter space. It offers an interface to loading and saving the model as well as to the training.

Constructor.

Parameters:

transform (glasflow.nflows.transforms.base.Transform) – A bijection that transforms data into noise (in the glasflow.nflows style)
distribution (glasflow.nflows.distributions.base.Distribution) – The base distribution of the flow that generates the noise (in the glasflow.nflows style)

load_weights(filename)[source]

Load the weights of the flow from file. The weights must match the architecture of the flow. It is equivalent to

self.load_state_dict(torch.load(filename))

Parameters:: filename (str) – File to load the weights from

sample_within_boundaries(num_samples, boundaries=None, seed=None)[source]: Generate a number of samples within the given boundaries

save_weigths(filename)[source]

Saves the weigths to filename. It is equivalent to:

torch.save(self.state_dict(), filename)

Parameters:: filename (str) – Name of the file to save the weights at (a file-like object)

train_flow(loss, N_epochs, train_data, validation_data, optimizer, train_weights=None, validation_weights=None, batch_size=None, validation_step=10, callback=None, lr_scheduler=None, boundaries=None, verbose=False)[source]

Trains the normalizing flow.

It can use several loss functions:

‘forward_KL’: uses the forward KL entropy. In this case, train_data and validation_data must be samples and train_weights and validation_weights will be ignored. See eq. (13) of 1912.02762.
‘ll_mse’: uses the means squared error of the log likelihood between the flow and the target distribution. In this case, train_weights and validation_weights must be provided and have the meaning of log likelihood of each train_data and validation_data respectively. The data doesn’t need to be drawn from any particular distribution

The training supports callbacks and learning rate decay. Early stop can be implemented through callbacks.

Parameters:

N_epochs (int) – Number of training epochs
train_data (ndarray) – Training data. They need to fit the dimensionality of the flow and be convertible into a torch tensor
validation_data (ndarray) – Validation data. They need to fit the dimensionality of the flow and be convertible into a torch tensor

optimizer (torch.optim) –

An torch optimizer object. Typical usage is:

from torch import optim
flow = GW_Flow(**args)
optimizer = optim.Adam(flow.parameters(), lr=0.001)

train_weights (ndarray) – Weights for each of the training data
validation_weights (ndarray) – Weights for each of the validation data
batch_size (int) – Batch size: number of points of the training set to be used at each iteration
validation_step (int) – Number of training steps after which the validation metric is computed
callback (callable) – A callable, called at each validation step of the training. It has to have call signature callback(GW_Flow, epoch): see mbank.flow.utils.plotting_callback() for an example.

lr_scheduler –

A torch learning rate scheduler. Typical example would be:

from torch import optim
flow = GW_Flow(**args)
optimizer = optim.Adam(flow.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99)

Returns:

history – A dictionary keeping the historical values of training & validation loss function + KS metric. It has the following entries:

validation_step: number of epochs between two consecutive evaluation of the validation metrics
train_loss: values of the loss function at each iteration
validation_loss: values of the validation loss function at each validation iteration
valmetric_value: values of the validation metric at each validation iteration
valmetric_mean: expected value of the validation metric for a perfectly trained flow
valmetric_std: standard deviation of the validation metric for a perfectly trained flow
validation_metric: name of the validation metric being used

Return type:

dict

class STD_GW_Flow(D, n_layers, hidden_features, has_constant=True)[source]

An implementation of the standard flow: the flow is composed by one TanhTransform layer and a stack of layers made by NaiveLinear+MaskedAffineAutoregressiveTransform. All the applications within mbank uses of this class.

Initialization of the flow

Parameters:

D (int) – Dimensionality of the flow
n_layers (int) – Number of layers (each made by NaiveLinear+MaskedAffineAutoregressiveTransform)
hidden_features (int) – Number of hidden features of the MaskedAffineAutoregressiveTransform. If a list is given, it is intended to be the number of hidden features for each layer

property boundary_box: Returns the boundaries of the support of the flow (a rectangle)

classmethod load_flow(weigth_file)[source]

Loads the flow from a file. The architecture is inferred from the loaded weights.

new_flow = STD_GW_Flow.load_flow('weight.zip')

Parameters:: weigth_file (str) – File to load the flow from
Returns:: new_flow – Initialized flow
Return type:: STD_GW_Flow

log_volume_element(theta)[source]

Returns an estimation to the (log) volume element \(\log\sqrt{M( heta)}\) givev by the normalizing flow. This is equivalent to the flow log_pdf, scaled by a constant.

Parameters:: theta (ndarray) – Input points to evaluate the at
Returns:: log_sqrt_metric – The natural logarithm of the volume element \(\log\sqrt{M( heta)}\).
Return type:: ndarray

class TanhTransform(D, low=None, high=None)[source]

Implements the Tanh transformation. This maps a Rectangle [low, high] into R^D. It is very recommended to use this as the last layer of every flow you will ever train on GW data.

Initialize the transformation.

Parameters:: D (int) – Dimensionality of the space

forward(inputs, context=None)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class powerlawTransform(D, alphas=None, low=None, high=None)[source]

Implements a trainable power law transformation for each dimension of the flow

Initialize the transformation.

Parameters:: f_low (float) – Lower frequency cut-off for the transformation

forward(inputs, context=None)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class tau0tau3Transform[source]

Implements the tau0tau3 transformation as in https://arxiv.org/pdf/0706.4437.pdf

Initialize the transformation.

Parameters:: f_low (float) – Lower frequency cut-off for the transformation

forward(inputs, context=None)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

mbank.flow.utils

Plotting & validation utilities for the mbank.flow

compare_probability_distribution(data_flow, data_true=None, variable_format=None, title=None, hue_labels=('flow', 'train'), savefile=None, show=False)[source]

Shows the probability distribution learnt by the flow and compares it with the training one. It makes a nice contour plot to visualize the 2D slices of the multidimensional PDF.

Parameters:

data_flow (ndarray) – Samples from the normalizing flow
data_true (ndarray) – Samples from the target (true) distribution (if None, it will not be plotted)
variable_format (str) – Variable format, to place the axes labels properly
title (str) – A title for the plot
hue_labels (list/tuple) – Labels for the two distributions: they will appear in the legend
savefile (str) – File to save the plot at
show (bool) – Whether to show the plot

create_gif(folder, savefile, fps=1)[source]: Given a folder of plots generated by a callback, it creates a gif summarizing the training history

class cross_entropy_metric(data, flow, N_estimation=1000)[source]

Class to compute the validation metric using the Cross Entropy distance

get_dist(data)[source]: Measures the distance between data and the noise distribution (a standard normal)

class early_stopper(patience=10, min_delta=0, temp_file=None, return_best_model=True, verbose=False)[source]: Implements early stopping for the training of the normalizing flow model

class ks_metric(data, flow, N_estimation=1000)[source]

Class to compute the validation metric using the Kolmogorov-Smirnov test

get_dist(data)[source]: Measures the distance between data and the noise distribution (a standard normal)

plot_loss_functions(history, savefolder=None)[source]

Given a history dict, returned by mbank.flow.flowmodel.GW_Flow.train_flow_forward_KL(), it plots the loss function and the validation metric as a function of the epoch

Parameters:

history (dict) – An history dict (as returned by mbank.flow.flowmodel.GW_Flow.train_flow_forward_KL())
savefolder (str) – A folder where to save the plots: they will be saved with the names loss.png and validation_metric.png.

plotting_callback(model, epoch, train_loss, validation_loss, dirname, data_to_plot, variable_format, basefilename=None)[source]: An example callback for plotting the KDE pairplots.

class validation_metric(data, flow, N_estimation=1000)[source]

Base class for validation metric: the method get_dist should be implemented by subclasses

get_dist(data)[source]: Measures the distance between data and the noise distribution (a standard normal)

get_validation_metric()[source]: Check if the data transformed into noise are consistent with the random normal distribution