def test_initialise(tmpdir, proposal, ef, fuzz):
"""Test the initialise method"""
p = tmpdir.mkdir('test')
proposal.output = f'{p}/output/'
proposal.rescaled_dims = 2
proposal.expansion_fraction = ef
proposal.fuzz = 2.0
proposal.flow_config = {'model_config': {}}
proposal.set_rescaling = MagicMock()
proposal.verify_rescaling = MagicMock()
proposal.update_flow_config = MagicMock()
proposal.configure_constant_volume = MagicMock()
fm = MagicMock()
fm.initialise = MagicMock()
with patch('nessai.proposal.flowproposal.FlowModel', new=fm) as mock_fm:
FlowProposal.initialise(proposal)
proposal.set_rescaling.assert_called_once()
proposal.verify_rescaling.assert_called_once()
proposal.update_flow_config.assert_called_once()
proposal.configure_constant_volume.assert_called_once()
mock_fm.assert_called_once_with(config=proposal.flow_config,
output=proposal.output)
proposal.flow.initialise.assert_called_once()
assert proposal.populated is False
assert proposal.initialised
assert proposal.fuzz == fuzz
assert os.path.exists(f'{p}/output')
def test_set_boundary_inversion_no_rescaling(proposal):
"""Check boundary inversion is disabled if paraemters are not rescaled"""
proposal.rescale_parameters = False
proposal.boundary_inversion = ['x']
with pytest.raises(RuntimeError) as excinfo:
FlowProposal.set_boundary_inversion(proposal)
assert 'Boundary inversion requires rescaling' in str(excinfo.value)
def test_training_plots(proposal, tmpdir, plot):
"""Make sure traings plots are correctly produced"""
proposal._plot_training = plot
output = tmpdir.mkdir('test/')
names = ['x', 'y']
prime_names = ['x_prime', 'y_prime']
z = np.random.randn(10, 2)
x = np.random.randn(10, 2)
x_prime = x / 2
proposal.training_data = numpy_array_to_live_points(x, names)
proposal.training_data_prime = \
numpy_array_to_live_points(x_prime, prime_names)
x_gen = numpy_array_to_live_points(x, names)
x_prime_gen = numpy_array_to_live_points(x_prime, prime_names)
proposal.dims = 2
proposal.rescale_parameters = names
proposal.rescaled_names = prime_names
proposal.forward_pass = MagicMock(return_value=(z, None))
proposal.backward_pass = MagicMock(return_value=(x_prime_gen, np.ones(10)))
proposal.inverse_rescale = MagicMock(return_value=(x_gen, np.ones(10)))
proposal.check_prior_bounds = lambda *args: args
proposal.model = MagicMock()
proposal.model.names = names
FlowProposal._plot_training_data(proposal, output)
assert os.path.exists(f'{output}/x_samples.png') is bool(plot)
assert os.path.exists(f'{output}/x_generated.png') is bool(plot)
assert os.path.exists(f'{output}/x_prime_samples.png') is bool(plot)
assert os.path.exists(f'{output}/x_prime_generated.png') is bool(plot)
def test_config_fixed_radius_not_float(proposal):
"""
Test the fixed radius is disabled when the radius cannot be converted to
a float.
"""
FlowProposal.configure_fixed_radius(proposal, 'four')
assert proposal.fixed_radius is False
def test_set_boundary_inversion_incorrect_parameters(proposal):
"""Check an error is raised if a parameter is not recognised"""
proposal.rescale_parameters = True
proposal.boundary_inversion = ['x', 'z']
proposal.names = ['x', 'y']
with pytest.raises(RuntimeError) as excinfo:
FlowProposal.set_boundary_inversion(proposal)
assert 'Boundaries are not in known parameters' in str(excinfo.value)
def test_init_use_default_reparams(model, proposal, value, expected):
"""Assert use_default_reparameterisations is set correctly"""
proposal.use_default_reparameterisations = False
FlowProposal.__init__(proposal,
model,
poolsize=10,
use_default_reparameterisations=value)
assert proposal.use_default_reparameterisations is expected
def test_set_boundary_inversion(proposal, inversion, parameters):
"""Check the correct parameters are inverted"""
proposal.names = ['x', 'y', 'z']
proposal.rescale_parameters = ['x', 'y']
proposal.boundary_inversion = inversion
proposal.inversion_type = 'split'
FlowProposal.set_boundary_inversion(proposal)
assert proposal.boundary_inversion == parameters
def test_set_boundary_inversion_incorrect_inversion_type(proposal):
"""Check an error is raised if the inversion type is unknown"""
proposal.names = ['x']
proposal.rescale_parameters = True
proposal.boundary_inversion = True
proposal.inversion_type = 'half'
with pytest.raises(RuntimeError) as excinfo:
FlowProposal.set_boundary_inversion(proposal)
assert 'Unknown inversion type' in str(excinfo.value)
def test_config_poolsize_none(proposal):
"""
Test the popluation configuration raises an error if poolsize is None.
"""
with pytest.raises(RuntimeError) as excinfo:
FlowProposal.configure_population(proposal, None, None, True, 10, 1.0,
0.0, 'gaussian')
assert 'poolsize' in str(excinfo.value)
def test_plot_pool_all(proposal):
"""Test for the plots that show the pool of samples"""
proposal.output = 'test'
proposal._plot_pool = 'all'
proposal.populated_count = 0
x = numpy_array_to_live_points(np.random.randn(10, 2), ['x', 'y'])
with patch('nessai.proposal.flowproposal.plot_live_points') as plot:
FlowProposal.plot_pool(proposal, None, x)
plot.assert_called_once_with(x, c='logL', filename='test/pool_0.png')
def test_set_boundary_inversion_parameter_not_rescaled(proposal):
"""Check an error is raised if the parameter for inversion is not included
in the rescaled parameters.
"""
proposal.names = ['x', 'y']
proposal.rescale_parameters = ['y']
proposal.boundary_inversion = ['x']
with pytest.raises(RuntimeError) as excinfo:
FlowProposal.set_boundary_inversion(proposal)
assert 'Boundaries are not in rescaled parameters' in str(excinfo.value)
def test_constant_volume_invalid_latent_prior(proposal):
"""Assert an error is raised if the latent prior is not a truncated \
Gaussian
"""
err = "Constant volume requires `latent_prior='truncated_gaussian'`"
proposal.constant_volume_mode = True
proposal.latent_prior = 'gaussian'
with pytest.raises(RuntimeError) as excinfo:
FlowProposal.configure_constant_volume(proposal)
assert str(excinfo.value) == err
def test_resume_w_weights(proposal):
"""Test the resume method with weights"""
proposal.initialise = MagicMock()
proposal.flow = MagicMock()
proposal.mask = None
proposal.update_bounds = False
proposal.weights_file = None
FlowProposal.resume(proposal, None, {}, 'weights.pt')
proposal.initialise.assert_called_once()
proposal.flow.reload_weights.assert_called_once_with('weights.pt')
def test_rejection_sampling_truncate_missing_q(proposal, z, x, log_q):
"""Test rejection sampling method with truncation without without q"""
proposal.use_x_prime_prior = False
proposal.truncate = True
log_q = np.array([0.0, 1.0])
proposal.backward_pass = MagicMock(return_value=(x, log_q))
with pytest.raises(RuntimeError) as excinfo:
FlowProposal.rejection_sampling(proposal, z, worst_q=None)
assert 'Cannot use truncation' in str(excinfo.value)
def test_configure_latent_prior_var(proposal):
"""
Test the latent prior when using a truncated Gaussian with a variance.
"""
proposal.latent_prior = 'truncated_gaussian'
proposal.flow_config = {'model_config': {'kwargs': {'var': 4}}}
proposal.draw_latent_kwargs = {}
FlowProposal.configure_latent_prior(proposal)
assert proposal.draw_latent_prior == \
getattr(utils, 'draw_truncated_gaussian')
assert proposal.draw_latent_kwargs.get('var') == 4
def test_resume(proposal):
"""Test the resume method."""
from numpy import array, array_equal
proposal.initialise = MagicMock()
proposal.mask = [1, 0]
proposal.update_bounds = False
proposal.weights_file = None
FlowProposal.resume(proposal, None, {'model_config': {'kwargs': {}}})
proposal.initialise.assert_called_once()
assert array_equal(proposal.flow_config['model_config']['kwargs']['mask'],
array([1, 0]))
def test_configure_reparameterisations(tmpdir, model, reparameterisation):
"""Test adding one of the default reparameterisations.
Only tests reparameterisations that don't need extra parameters.
"""
proposal = FlowProposal(model,
output=str(tmpdir.mkdir('test')),
poolsize=10,
reparameterisations={'x': reparameterisation})
proposal.set_rescaling()
assert proposal._reparameterisation is not None
def test_set_boundary_inversion_rescale_parameters(proposal,
rescale_parameters,
parameters):
"""Check the correct parameters depending on the type of rescale
parameters.
"""
proposal.names = ['x', 'y']
proposal.rescale_parameters = rescale_parameters
proposal.boundary_inversion = True
proposal.inversion_type = 'split'
FlowProposal.set_boundary_inversion(proposal)
assert proposal.boundary_inversion == parameters
def test_configure_reparameterisation_scale(tmpdir, model, reparameterisation):
"""Test adding the `Rescale` reparameterisation"""
proposal = FlowProposal(model,
output=str(tmpdir.mkdir('test')),
poolsize=10,
reparameterisations={
'x': {
'reparameterisation': reparameterisation,
'scale': 2.0
}
})
proposal.set_rescaling()
assert proposal._reparameterisation is not None
def test_default_reparameterisations(caplog, tmpdir):
"""Assert that by default reparameterisations are not used."""
caplog.set_level('INFO')
model = MagicMock()
model.names = ['x', 'y']
model.bounds = {p: [-1, 1] for p in model.names}
model.reparameterisations = None
proposal = FlowProposal(model,
poolsize=100,
output=str(tmpdir.mkdir('test')))
# Mocked model so can't verify rescaling
proposal.verify_rescaling = MagicMock()
proposal.initialise()
assert proposal._reparameterisation is None
def test_configure_reparameterisation_angle_pair(tmpdir, model):
"""Test adding the `AnglePair` reparameterisation"""
model.names.append('y')
model.bounds = {'x': [0, 2 * np.pi], 'y': [-np.pi / 2, np.pi / 2]}
proposal = FlowProposal(model,
output=str(tmpdir.mkdir('test')),
poolsize=10,
reparameterisations={
'x': {
'reparameterisation': 'angle-pair',
'parameters': ['y']
}
})
proposal.set_rescaling()
assert proposal._reparameterisation is not None
def test_draw_not_popluated(proposal, update):
"""Test the draw method when the proposal is not populated"""
import datetime
proposal.populated = False
proposal.poolsize = 100
proposal.population_time = datetime.timedelta()
proposal.samples = None
proposal.indices = []
proposal.update_poolsize = update
proposal.update_poolsize_scale = MagicMock()
proposal.ns_acceptance = 0.5
def mock_populate(*args, **kwargs):
proposal.populated = True
proposal.samples = np.arange(3)
proposal.indices = list(range(3))
proposal.populate = MagicMock(side_effect=mock_populate)
out = FlowProposal.draw(proposal, 1.)
assert out == 2
assert proposal.populated is True
assert proposal.population_time.total_seconds() > 0.
proposal.populate.assert_called_once_with(1., N=100)
assert proposal.update_poolsize_scale.called == update
def test_reset(proposal):
"""Test reset method"""
proposal.x = 1
proposal.samples = 2
proposal.populated = True
proposal.populated_count = 10
proposal._edges = {'x': 2}
FlowProposal.reset(proposal)
assert proposal.x is None
assert proposal.samples is None
assert proposal.populated is False
assert proposal.populated_count == 0
assert proposal.r is None
assert proposal.alt_dist is None
assert proposal._checked_population
assert proposal._edges['x'] is None
def test_get_state(proposal, populated):
"""Test the get state method used for pickling the proposal.
Tests cases where the proposal is and isn't populated.
"""
proposal.populated = populated
proposal.indices = [1, 2]
proposal._reparameterisation = MagicMock()
proposal.model = MagicMock()
proposal._flow_config = {}
proposal.pool = MagicMock()
proposal.initialised = True
proposal.flow = MagicMock()
proposal.flow.weights_file = 'file'
state = FlowProposal.__getstate__(proposal)
assert state['resume_populated'] is populated
assert state['pool'] is None
assert state['initialised'] is False
assert state['weights_file'] == 'file'
assert '_reparameterisation' not in state
assert 'model' not in state
assert 'flow' not in state
assert '_flow_config' not in state
def test_training(proposal, tmpdir, save, plot, plot_training):
"""Test the training method"""
output = tmpdir.mkdir('test/')
data = np.random.randn(10, 2)
data_prime = data / 2
x = numpy_array_to_live_points(data, ['x', 'y'])
x_prime = numpy_array_to_live_points(data_prime, ['x_prime', 'y_prime'])
log_j = np.ones(data.shape[0])
proposal.training_count = 0
proposal.populated = True
proposal._plot_training = plot_training
proposal.save_training_data = save
proposal.rescale_parameters = ['x']
proposal.rescaled_names = ['x_prime', 'y_prime']
proposal.output = output
proposal.check_state = MagicMock()
proposal.rescale = MagicMock(return_value=(x_prime, log_j))
proposal.flow = MagicMock()
proposal.flow.train = MagicMock()
proposal._plot_training_data = MagicMock()
with patch('nessai.proposal.flowproposal.live_points_to_array',
return_value=data_prime), \
patch('nessai.proposal.flowproposal.save_live_points') as mock_save:
FlowProposal.train(proposal, x, plot=plot)
np.testing.assert_array_equal(x, proposal.training_data)
if save or (plot and plot_training):
output = f'{output}/training/block_0/'
if save:
mock_save.assert_called_once()
if plot and plot_training:
proposal._plot_training_data.assert_called_once_with(output)
elif not plot or not plot_training:
proposal._plot_training_data.assert_not_called()
proposal.check_state.assert_called_once_with(proposal.training_data)
proposal.rescale.assert_called_once_with(x)
proposal.flow.train.assert_called_once_with(
data_prime, output=output, plot=plot and plot_training)
assert proposal.training_count == 1
assert proposal.populated is False
def test_radius_w_log_q(proposal):
"""Test computing the radius with log_q"""
z = np.array([[1, 2, 3], [0, 1, 2]])
log_q = np.array([1, 2])
expected_r = np.sqrt(14)
r, log_q_r = FlowProposal.radius(proposal, z, log_q=log_q)
assert r == expected_r
assert log_q_r == log_q[0]
def test_configure_constant_volume(proposal):
"""Test configuration for constant volume mode."""
proposal.constant_volume_mode = True
proposal.volume_fraction = 0.95
proposal.rescaled_dims = 5
proposal.latent_prior = 'truncated_gaussian'
proposal.max_radius = 3.0
proposal.min_radius = 5.0
proposal.fuzz = 1.5
with patch('nessai.proposal.flowproposal.compute_radius',
return_value=4.0) as mock:
FlowProposal.configure_constant_volume(proposal)
mock.assert_called_once_with(5, 0.95)
assert proposal.fixed_radius == 4.0
assert proposal.min_radius is False
assert proposal.max_radius is False
assert proposal.fuzz == 1.0
def test_draw_populated(proposal):
"""Test the draw method if the proposal is already populated"""
proposal.populated = True
proposal.samples = np.arange(3)
proposal.indices = list(range(3))
out = FlowProposal.draw(proposal, None)
assert out == proposal.samples[2]
assert proposal.indices == [0, 1]
def test_get_alt_distribution_truncated_gaussian(proposal):
"""
Test getting the alternative distribution for the default latent prior, the
truncated Gaussian with var=1. This should return None.
"""
proposal.draw_latent_kwargs = {}
proposal.latent_prior = 'truncated_gaussian'
dist = FlowProposal.get_alt_distribution(proposal)
assert dist is None
def test_resume_w_update_bounds(proposal, data, count):
"""Test the resume method with update bounds"""
proposal.initialise = MagicMock()
proposal.flow = MagicMock()
proposal.mask = None
proposal.update_bounds = True
proposal.weights_file = None
proposal.training_data = data
proposal.training_count = count
proposal.check_state = MagicMock()
if count and data is None:
with pytest.raises(RuntimeError) as excinfo:
FlowProposal.resume(proposal, None, {})
assert 'Could not resume' in str(excinfo.value)
else:
FlowProposal.resume(proposal, None, {})
if data:
proposal.check_state.assert_called_once_with(data)
