nessai.proposal.flowproposal

Main proposal object that includes normalising flows.

Module Contents

Classes

FlowProposal

Object that handles training and proposal points

class nessai.proposal.flowproposal.FlowProposal(model, flow_config=None, output='./', poolsize=None, rescale_parameters=True, latent_prior='truncated_gaussian', constant_volume_mode=True, volume_fraction=0.95, fuzz=1.0, plot='min', fixed_radius=False, drawsize=None, check_acceptance=False, truncate_log_q=False, rescale_bounds=[-1, 1], expansion_fraction=4.0, boundary_inversion=False, inversion_type='split', update_bounds=True, min_radius=False, max_radius=50.0, max_poolsize_scale=10, update_poolsize=True, accumulate_weights=False, save_training_data=False, compute_radius_with_all=False, detect_edges=False, detect_edges_kwargs=None, reparameterisations=None, fallback_reparameterisation=None, use_default_reparameterisations=None, reverse_reparameterisations=False)

Bases: nessai.proposal.rejection.RejectionProposal

Object that handles training and proposal points

Parameters:

modelnessai.model.Model

User defined model.

flow_configdict, optional

Configuration for training the normalising flow. If None, uses default settings. Defaults to None.

outputstr, optional

Path to output directory.

plot{True, False, ‘all’, ‘min’}, optional

Controls the plotting level: True - all plots; False - no plots; 'all' - all plots and 'min' - 1d plots and loss.

latent_prior{‘truncated_gaussian’, ‘gaussian’, ‘uniform_nsphere’, ‘gaussian’}, optional

Prior distribution in the latent space. Defaults to ‘truncated_gaussian’.

poolsizeint, optional

Size of the proposal pool. Defaults to 10000.

update_poolsizebool, optional

If True the poolsize is updated using the current acceptance of the nested sampler.

max_poolsize_scaleint, optional

Maximum scale for increasing the poolsize. E.g. if this value is 10 and the poolsize is 1000 the maximum number of points in the pool is 10,000.

drawsizeint, optional

Number of points to simultaneously draw when populating the proposal Defaults to 10000

check_acceptancebool, optional

If True the acceptance is computed after populating the pool. This includes computing the likelihood for every point. Default False.

min_radiusfloat, optional

Minimum radius used for population. If not specified not minimum is used.

max_radiusfloat, optional

If a float then this value is used as an upper limit for the computed radius when populating the proposal. If unspecified no upper limit is used.

fixed_radiusfloat, optional

If specified and the chosen latent distribution is compatible, this radius will be used to draw new samples instead of the value computed with the flow.

constant_volume_modebool

If True, then a constant volume is used for the latent contour used to draw new samples. The exact volume can be set using volume_fraction

volume_fractionfloat

Fraction of the total probability to contain with the latent contour when using a constant volume.

compute_radius_with_allbool, optional

If True all the radius of the latent contour is computed using the maximum radius of all the samples used to train the flow.

fuzzfloat, optional

Fuzz-factor applied to the radius. If unspecified no fuzz-factor is applied.

expansion_fractionfloat, optional

Similar to fuzz but instead a scaling factor applied to the radius this specifies a rescaling for volume of the n-ball used to draw samples. This is translated to a value for fuzz.

truncate_log_qbool, optional

Truncate proposals using minimum log-probability of the training data.

rescale_parameterslist or bool, optional

If True live points are rescaled to rescale_bounds before training. If an instance of list then must contain names of parameters to rescale. If False no rescaling is applied.

rescale_boundslist, optional

Lower and upper bound to use for rescaling. Defaults to [-1, 1]. See rescale_parameters.

update_boundsbool, optional

If True bounds used for rescaling are updated at the starting of training. If False prior bounds are used.

boundary_inversionbool or list, optional

If True boundary inversion is applied to all bounds. If If an instance of list of parameters names, then inversion only applied to these parameters. If False (default )no inversion is used.

inversion_type{‘split’, ‘duplicate’}

Type of inversion to use. 'split' keeps the number of samples the sample but mirrors half around the bound. 'duplicate' mirrors all the samples at the bound.

detect_edgesbool, optional

If True, when applying the version the option of no inversion is allowed.

detect_edges_kwargsdict, optional

Dictionary of keyword arguments passed to nessai.utils.detect_edge().

reparameterisationsdict, optional

Dictionary for configure more flexible reparameterisations. This ignores any of the other settings related to rescaling. For more details see the documentation.

fallback_reparameterisationNone or str

Name of the reparameterisation to be used for parameters that have not been specified in the reparameterisations dictionary. If None, the NullReparameterisation is used. Reparameterisation should support multiple parameters.

use_default_reparameterisationsbool, optional

If True then reparameterisations will be used even if reparameterisations is None. The exact reparameterisations used will depend on add_default_reparameterisations() which may be overloaded by child classes. If not specified then the value of the attribute use_default_reparameterisations is used.

reverse_reparameterisationsbool

Passed to reverse_order in CombinedReparameterisation. Reverses the order of the reparameterisations.

property poolsize

Return the poolsize based of the base value and the current value of the scaling

property flow_config

Return the configuration for the flow

property dims

Return the number of dimensions

property rescaled_dims

Return the number of rescaled dimensions

property x_dtype

Return the dtype for the x space

property x_prime_dtype

Return the dtype for the x prime space

property population_dtype

dtype used for populating the proposal, depends on if the prior is defined in the x space or x-prime space

use_default_reparameterisations = False

Indicates whether reparameterisations will be used be default in this class. Child classes can change this value a force the default behaviour to change without changing the keyword arguments.

configure_population(poolsize, drawsize, update_poolsize, max_poolsize_scale, fuzz, expansion_fraction, latent_prior)

Configure settings related to population

configure_plotting(plot)

Configure plotting.

Plotting is split into training and pool. Training refers to plots produced during training and pool refers to plots produces during the population stage.

Parameters:

plot{True, False, ‘all’, ‘train’, ‘pool’, ‘min’, ‘minimal’}

Level of plotting. all, train and pool enable corner style plots. All other values that evaluate to True enable 1d histogram plots. False disables all plotting.

configure_latent_prior()

Configure the latent prior

configure_fixed_radius(fixed_radius)

Configure the fixed radius

configure_min_max_radius(min_radius, max_radius)

Configure the minimum and maximum radius

configure_constant_volume()

Configure using constant volume latent contour.

update_flow_config()

Update the flow configuration dictionary.

initialise()

Initialise the proposal class.

This includes:

• Setting up the rescaling

• Verifying the rescaling is invertible

• Initialising the FlowModel

update_poolsize_scale(acceptance)

Update poolsize given the current acceptance.

Parameters:

acceptancefloat

Current acceptance.

add_default_reparameterisations()

Add any reparameterisations which are assumed by default

get_reparameterisation(name)

Get the reparameterisation from the name

configure_reparameterisations(reparameterisations)

Configure the reparameterisations.

Parameters:

reparameterisations{dict, None}

Dictionary of reparameterisations. If None, then the defaults from :py:func`get_default_reparameterisations` are used.

set_rescaling()

Set function and parameter names for rescaling

verify_rescaling()

Verify the rescaling functions are invertible

rescale(x, compute_radius=False, **kwargs)

Rescale from the physical space to the primed physical space

Parameters:

x: array_like

Array of live points to rescale

compute_radius: bool (False)

Used to indicate when rescaling is being used for computing the radius for population. This is important for rescaling that uses inversions.

Returns:

array

Array of rescaled values

array

Array of log det|J|

inverse_rescale(x_prime, **kwargs)

Rescale from the primed physical space to the original physical space.

Parameters:

x_primearray_like

Array of live points to rescale.

Returns:

array

Array of rescaled values in the data space.

array

Array of log-Jacobian determinants.

check_state(x)

Update the state of the proposal given some training data.

Includes updating the reparameterisations.

Parameters:

x: array_like

Array of training live points which can be used to set parameters

train(x, plot=True)

Train the normalising flow given some of the live points.

Parameters:

xstructured_array

Array of live points

plot{True, False, ‘all’}

Enable or disable plots for during training. By default the plots are only one-dimensional histograms, ‘all’ includes corner plots with samples, these are often a few MB in size so proceed with caution!

reset_model_weights(**kwargs)

Reset the flow weights.

Parameters:

kwargs

Keyword arguments passed to nessai.flowmodel.FlowModel.reset_model().

check_prior_bounds(x, *args)

Return only values that are within the prior bounds

Parameters:

x: array_like

Array of live points which will compared against prior bounds

*args:

Additional arrays which correspond to the array of live points. Only those corresponding to points within the prior bounds are returned

Returns:

out: tuple of arrays

Array containing the subset of the original arrays which fall within the prior bounds

forward_pass(x, rescale=True, compute_radius=True)

Pass a vector of points through the model

Parameters:

xarray_like

Live points to map to the latent space

rescalebool, optional (True)

Apply rescaling function

compute_radiusbool, optional (True)

Flag parsed to rescaling for rescaling specific to radius computation

Returns:

xarray_like

Samples in the latent space

log_probarray_like

Log probabilities corresponding to each sample (including the jacobian)

backward_pass(z, rescale=True, discard_nans=True, return_z=False)

A backwards pass from the model (latent -> real)

Parameters:

zarray_like

Structured array of points in the latent space

rescalebool, optional (True)

Apply inverse rescaling function

discard_nan: bool

If True, samples with NaNs or Infs in log_q are removed.

return_zbool

If True, return the array of latent samples, this may differ from the input since samples can be discarded.

Returns:

xarray_like

Samples in the data space

log_probarray_like

Log probabilities corresponding to each sample (including the Jacobian)

zarray_like

Samples in the latent space, only returned if return_z=True

radius(z, *arrays)

Calculate the radius of a latent point or set of latent points. If multiple points are parsed the maximum radius is returned.

Parameters:

znp.ndarray

Array of points in the latent space

*arrays

Additional arrays to return the corresponding value

Returns:

tuple of arrays

Tuple of array with the maximum radius and corresponding values from any additional arrays that were passed.

log_prior(x)

Compute the prior probability using the user-defined model

Parameters:

xstructured_array

Array of samples

Returns:

array_like

Array of log prior probabilities

x_prime_log_prior(x)

Compute the prior in the prime space

Parameters:

xarray

Samples in the X-prime space.

compute_weights(x, log_q, return_log_prior=False)

Compute weights for the samples.

Does NOT normalise the weights

Parameters:

xstructured_arrays

Array of points

log_qarray_like

Array of log proposal probabilities.

return_log_prior: bool

If true, the log-prior probability is also returned.

Returns:

array_like

Log-weights for rejection sampling.

rejection_sampling(z, min_log_q=None)

Perform rejection sampling.

Converts samples from the latent space and computes the corresponding weights. Then returns samples using standard rejection sampling.

Parameters:

zndarray

Samples from the latent space

min_log_qfloat, optional

Lower bound on the log-probability computed using the flow that is used to truncate new samples.

Returns:

array_like

Array of accepted latent samples.

array_like

Array of accepted samples in the X space.

convert_to_samples(x, plot=True)

Convert the array to samples ready to be used.

This removes are auxiliary parameters, (e.g. auxiliary radial parameters) and ensures the prior is computed. These samples can be directly used in the nested sampler.

Parameters:

xarray_like

Array of samples

plotbool, optional

If true a 1d histogram for each parameter of the pool is plotted. This includes a comparison the live points used to train the current realisation of the flow.

Returns:

array

Structured array of samples.

prep_latent_prior()

Prepare the latent prior.

draw_latent_prior(n)

Draw n samples from the latent prior.

populate(worst_point, N=10000, plot=True, r=None, max_samples=1000000)

Populate a pool of latent points given the current worst point.

Parameters:

worst_pointstructured_array

The current worst point used to compute the radius of the contour in the latent space.

Nint, optional (10000)

The total number of points to populate in the pool

plot{True, False, ‘all’}

Enable or disable plots for during training. By default the plots are only one-dimensional histograms, ‘all’ includes corner plots with samples, these are often a few MB in size so proceed with caution!

get_alt_distribution()

Get a distribution for the latent prior used to draw samples.

compute_acceptance(logL)

Compute how many of the current pool have log-likelihoods greater than the specified log-likelihood using the current value in the logL field.

Parameters:

floatlogL

Log-likelihood to use as the lower bound

Returns:

float

Acceptance defined on [0, 1]

draw(worst_point)

Draw a replacement point. The new point is independent of the worst point. The worst point is only used during population.

Parameters:

worst_pointstructured_array

The current worst point used to compute the radius of the contour in the latent space.

Returns:

structured_array

New live point

plot_pool(x)

Plot the pool of points.

Parameters:

xarray_like

Corresponding samples to plot in the physical space.

resume(model, flow_config, weights_file=None)

Resume the proposal.

The model and config are not stored so these must be provided.

Parameters:

modelnessai.model.Model

User-defined model used.

flow_configdict

Configuration dictionary for the flow.

weights_filesstr, optional

Weights file to try and load. If not provided the proposal tries to load the last weights file.

reset()

Reset the proposal