def test_reverse_inversion(reparam):
"""Assert the reverse inversion works correctly"""
reparam.parameters = ['x']
reparam.prime_parameters = ['x_prime']
reparam.offsets = {'x': 1.0}
reparam._edges = {'x': 'upper'}
reparam.bounds = {'x': [0, 5]}
x_val = np.array([[-0.7], [0.4]])
x = numpy_array_to_live_points(x_val, ['x'])
x_prime = numpy_array_to_live_points(np.array([3, 4]), ['x_prime'])
log_j = np.zeros(2)
# Return the same value to check that the negative values are handled
# correctly
with patch('nessai.reparameterisations.inverse_rescale_zero_to_one',
side_effect=lambda x, *args: (x, np.array([5, 6]))) as f:
x_out, x_prime_out, log_j_out = RescaleToBounds._reverse_inversion(
reparam,
x,
x_prime,
log_j,
'x',
'x_prime',
)
assert f.call_args_list[0][0][1] == 0.0
assert f.call_args_list[0][0][2] == 5.0
# Should be output of rescaling minus offset
np.testing.assert_array_equal(x_out['x'], np.array([1.3, 1.6]))
# x_prime should be the same
assert x_prime_out is x_prime
# Jacobian should just include jacobian from rescaling
np.testing.assert_array_equal(log_j_out, np.array([5, 6]))
def test_apply_inversion_detect_edge(reparam):
"""Assert detect edge is called with the correct arguments"""
reparam.parameters = ['x']
reparam.prime_parameters = ['x_prime']
reparam.offsets = {'x': 1.0}
reparam._edges = {'x': None}
reparam.detect_edges_kwargs = {'allowed_bounds': ['lower']}
reparam.bounds = {'x': [0, 5]}
reparam.rescale_bounds = {'x': [0, 1]}
x = numpy_array_to_live_points(np.array([1, 2]), ['x'])
x_prime = numpy_array_to_live_points(np.array([3, 4]), ['x_prime'])
log_j = np.zeros(2)
with patch('nessai.reparameterisations.detect_edge',
return_value=False) as mock_fn:
_ = RescaleToBounds._apply_inversion(reparam,
x,
x_prime,
log_j,
'x',
'x_prime',
False,
test=True)
reparam.update_prime_prior_bounds.assert_called_once()
mock_fn.assert_called_once_with(
x_prime['x_prime'],
test=True,
allowed_bounds=['lower'],
)
assert reparam._edges == {'x': False}
def test_reverse_inversion_not_applied(reparam):
"""Assert the reverse inversion works correctly"""
reparam.parameters = ['x']
reparam.prime_parameters = ['x_prime']
reparam.offsets = {'x': 1.0}
reparam._edges = {'x': False}
reparam.bounds = {'x': [0, 5]}
x_val = np.array([[1], [2]])
x = numpy_array_to_live_points(x_val, ['x'])
x_prime = numpy_array_to_live_points(np.array([3, 4]), ['x_prime'])
log_j = np.zeros(2)
with patch('nessai.reparameterisations.inverse_rescale_minus_one_to_one',
side_effect=lambda x, *args, **kwargs:
(x, np.array([5, 6]))) as f:
x_out, x_prime_out, log_j_out = RescaleToBounds._reverse_inversion(
reparam,
x,
x_prime,
log_j,
'x',
'x_prime',
)
assert f.call_args_list[0][1] == {'xmin': 0, 'xmax': 5}
# Should be output of rescaling minus offset
np.testing.assert_array_equal(x_out['x'], np.array([2, 3]))
# x_prime should be the same
assert x_prime_out is x_prime
# Jacobian should just include jacobian from rescaling
np.testing.assert_array_equal(log_j_out, np.array([5, 6]))
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_reparameterise_boundary_inversion(reparam):
"""Test the reparameterise function with boundary inversion"""
reparam.has_pre_rescaling = False
reparam.has_post_rescaling = False
reparam.parameters = ['x']
reparam.prime_parameters = ['x_prime']
reparam.offsets = {'x': 1.0}
reparam.boundary_inversion = {'x': 'split'}
x = numpy_array_to_live_points(np.array([(1.0, ), (2.0, )]),
reparam.parameters)
inversion_out = numpy_array_to_live_points(
np.array([(-1.0, ), (-2.0, ), (1.0, ), (2.0, )]),
reparam.prime_parameters)
x_prime_in = np.zeros([
2,
], dtype=get_dtype(reparam.prime_parameters))
log_j = np.zeros(x.size)
x_ex = np.concatenate([x, x])
x_prime_ex = inversion_out
log_j_ex = np.array([0, 0.5, 0, 0.5])
reparam._apply_inversion = MagicMock(return_value=(x_ex, x_prime_ex,
log_j_ex))
x_out, x_prime_out, log_j_out = RescaleToBounds.reparameterise(
reparam,
x,
x_prime_in,
log_j,
compute_radius=True,
test='test',
)
np.testing.assert_array_equal(
reparam._apply_inversion.call_args_list[0][0][0], x)
np.testing.assert_array_equal(
reparam._apply_inversion.call_args_list[0][0][1],
x_prime_in,
)
np.testing.assert_array_equal(
reparam._apply_inversion.call_args_list[0][0][2], log_j)
assert reparam._apply_inversion.call_args_list[0][0][3] == 'x'
assert reparam._apply_inversion.call_args_list[0][0][4] == 'x_prime'
assert reparam._apply_inversion.call_args_list[0][0][5] is True
assert reparam._apply_inversion.call_args_list[0][1] == {'test': 'test'}
np.testing.assert_array_equal(x_out, x_ex)
np.testing.assert_array_equal(x_prime_out, x_prime_ex)
np.testing.assert_array_equal(log_j_out, log_j_ex)
def test_apply_inversion_duplicate(reparam, inv_type, compute_radius):
"""Assert apply inversion with duplicate works as intended.
This test also covers compute_radius=True
"""
reparam.parameters = ['x', 'y']
reparam.prime_parameters = ['x_prime', 'y']
reparam.offsets = {'x': 1.0}
reparam._edges = {'x': 'lower'}
reparam.bounds = {'x': [0, 5]}
reparam.rescale_bounds = {'x': [0, 1]}
reparam.boundary_inversion = {'x': inv_type}
x_val = np.array([[1.2, 1.0], [1.7, 2.0]])
x_prime = numpy_array_to_live_points(x_val, ['x_prime', 'y'])
x = numpy_array_to_live_points(np.array([[1, 2], [3, 4]]), ['x', 'y'])
log_j = np.zeros(2)
with patch('numpy.random.choice') as rnd, \
patch('nessai.reparameterisations.rescale_zero_to_one',
side_effect=lambda x, *args: (x, np.array([5, 6]))) as f:
x_out, x_prime_out, log_j_out = RescaleToBounds._apply_inversion(
reparam,
x,
x_prime,
log_j,
'x',
'x_prime',
compute_radius,
)
rnd.assert_not_called()
assert f.call_args_list[0][0][1] == 0.0
assert f.call_args_list[0][0][2] == 5.0
# Output should be x_val minus offset
# Then duplicated
np.testing.assert_array_almost_equal(
x_prime_out['x_prime'],
np.array([0.2, 0.7, -0.2, -0.7]),
decimal=10,
)
np.testing.assert_array_almost_equal(
x_prime_out['y'],
np.array([1.0, 2.0, 1.0, 2.0]),
decimal=10,
)
# x should be the same but duplicated
np.testing.assert_array_equal(x_out, np.concatenate([x, x]))
# Jacobian should just include jacobian from rescaling but duplicated
np.testing.assert_array_equal(log_j_out, np.array([5, 6, 5, 6]))
def test_constant_volume_mode(tmpdir, model, check_acceptance,
rescale_parameters):
"""Integration test for constant volume mode.
With q=0.8647 should get a radius of ~2.
"""
output = str(tmpdir.mkdir('flowproposal'))
expected_radius = 2.0
fp = FlowProposal(
model,
output=output,
plot=False,
poolsize=10,
constant_volume_mode=True,
volume_fraction=0.8647,
check_acceptance=check_acceptance,
rescale_parameters=rescale_parameters,
)
fp.initialise()
worst = numpy_array_to_live_points(0.5 * np.ones(fp.dims), fp.names)
fp.populate(worst, N=100)
assert fp.x.size == 100
np.testing.assert_approx_equal(fp.r, expected_radius, 4)
np.testing.assert_approx_equal(fp.fixed_radius, expected_radius, 4)
def test_flowproposal_populate(tmpdir, model, latent_prior, expansion_fraction,
check_acceptance, rescale_parameters,
max_radius):
"""
Test the populate method in the FlowProposal class with a range of
parameters
"""
output = str(tmpdir.mkdir('flowproposal'))
fp = FlowProposal(
model,
output=output,
plot=False,
poolsize=100,
latent_prior=latent_prior,
expansion_fraction=expansion_fraction,
check_acceptance=check_acceptance,
rescale_parameters=rescale_parameters,
max_radius=max_radius,
constant_volume_mode=False,
)
fp.initialise()
worst = numpy_array_to_live_points(0.5 * np.ones(fp.dims), fp.names)
fp.populate(worst, N=100)
assert fp.x.size == 100
def test_update_prime_prior_bounds_integration():
"""Assert the prime prior bounds are correctly computed"""
rescaling = (
lambda x: (x / 2, np.zeros_like(x)),
lambda x: (2 * x, np.zeros_like(x)),
)
reparam = RescaleToBounds(
parameters=['x'],
prior_bounds=[1000, 1001],
prior='uniform',
pre_rescaling=rescaling,
offset=True,
)
np.testing.assert_equal(reparam.offsets['x'], 500.25)
np.testing.assert_array_equal(reparam.prior_bounds['x'], [1000, 1001])
np.testing.assert_array_equal(reparam.pre_prior_bounds['x'], [500, 500.5])
np.testing.assert_array_equal(reparam.bounds['x'], [-0.25, 0.25])
np.testing.assert_array_equal(reparam.prime_prior_bounds['x_prime'],
[-1, 1])
x_prime = numpy_array_to_live_points(
np.array([[-2], [-1], [0.5], [1], [10]]), ['x_prime'])
log_prior = reparam.x_prime_log_prior(x_prime)
expected = np.array([-np.inf, 0, 0, 0, -np.inf])
np.testing.assert_equal(log_prior, expected)
def test_rescale_angle_duplicate_or_compute_radius(reparam, x, args):
"""Test method for rescaling the 'angle'.
For ToCartesian this is the rescaling applied to the parameter.
"""
reparam.parameters = ['x']
reparam.prior_bounds = {'x': [0, 2]}
reparam.mode = args['mode']
reparam.scale = np.pi
expected_angle = np.array([0.0, np.pi / 2, np.pi, 0.0, -np.pi / 2, -np.pi])
x_prime = numpy_array_to_live_points(np.zeros([3, 1]), ['x_prime'])
log_j = np.zeros(3)
with patch('numpy.random.choice', return_value=np.array([1])):
angle, x_out, x_prime_out, log_j_out = ToCartesian._rescale_angle(
reparam, x, x_prime, log_j, compute_radius=args['compute_radius']
)
np.testing.assert_array_equal(angle, expected_angle)
np.testing.assert_array_equal(x_out, np.concatenate([x, x]))
np.testing.assert_array_equal(
x_prime_out, np.concatenate([x_prime, x_prime])
)
np.testing.assert_equal(log_j_out, -np.log(2))
def test_reparameterise(reparam):
"""Test the reparameterise function"""
reparam.has_pre_rescaling = False
reparam.has_post_rescaling = False
reparam.parameters = ['x']
reparam.prime_parameters = ['x_prime']
reparam.offsets = {'x': 1.0}
reparam.boundary_inversion = {}
x = numpy_array_to_live_points(np.array([(1.0, ), (2.0, )]),
reparam.parameters)
x_prime_in = np.zeros([
2,
], dtype=get_dtype(reparam.prime_parameters))
x_prime_val = np.array([0.0, 0.5])
log_j = np.zeros(x.size)
reparam._rescale_to_bounds = MagicMock(return_value=(x_prime_val,
np.array([0, 0.5])))
x_out, x_prime_out, log_j_out = \
RescaleToBounds.reparameterise(reparam, x, x_prime_in, log_j)
np.testing.assert_array_equal(
np.array([0.0, 1.0]),
reparam._rescale_to_bounds.call_args_list[0][0][0])
assert reparam._rescale_to_bounds.call_args_list[0][0][1] == 'x'
np.testing.assert_array_equal(x, x_out)
np.testing.assert_array_equal(x_prime_out['x_prime'], x_prime_val)
np.testing.assert_array_equal(log_j_out, np.array([0.0, 0.5]))
def test_inverse_reparameterise(reparam):
"""Test the inverse_reparameterise function"""
reparam.has_pre_rescaling = False
reparam.has_post_rescaling = False
reparam.parameters = ['x']
reparam.prime_parameters = ['x_prime']
reparam.offsets = {'x': 1.0}
reparam.boundary_inversion = {}
x_prime = numpy_array_to_live_points(np.array([(1.0, ), (2.0, )]),
reparam.prime_parameters)
x_in = np.zeros([
2,
], dtype=get_dtype(reparam.parameters))
x_val = np.array([0.0, 0.5])
log_j = np.zeros(x_prime.size)
reparam._inverse_rescale_to_bounds = MagicMock(
return_value=(x_val, np.array([0, 0.5])))
x_out, x_prime_out, log_j_out = \
RescaleToBounds.inverse_reparameterise(reparam, x_in, x_prime, log_j)
# x[p] is updated in place, can't test inputs
reparam._inverse_rescale_to_bounds.assert_called_once()
assert \
reparam._inverse_rescale_to_bounds.call_args_list[0][0][1] == 'x'
np.testing.assert_array_equal(x_prime_out, x_prime)
np.testing.assert_array_equal(x_out['x'], x_val + 1.0)
np.testing.assert_array_equal(log_j_out, np.array([0.0, 0.5]))
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_numpy_array_multiple_to_live_points(live_points):
"""
Test the function the produces an array of live points given numpy array
of shape [# point, # dimensions]
"""
array = np.array([[1., 2., 3.], [4., 5., 6.]])
x = lp.numpy_array_to_live_points(array, names=['x', 'y', 'z'])
np.testing.assert_array_equal(live_points, x)
def test_apply_inversion_split(reparam):
"""Assert apply inversion with split works as intended.
Also tests "upper" setting.
"""
reparam.parameters = ['x']
reparam.prime_parameters = ['x_prime']
reparam.offsets = {'x': 1.0}
reparam._edges = {'x': 'upper'}
reparam.bounds = {'x': [0, 5]}
reparam.rescale_bounds = {'x': [0, 1]}
reparam.boundary_inversion = {'x': 'split'}
x_val = np.array([[1.2], [1.7]])
x_prime = numpy_array_to_live_points(x_val, ['x_prime'])
x = numpy_array_to_live_points(np.array([3, 4]), ['x'])
log_j = np.zeros(2)
with patch('numpy.random.choice', return_value=np.array([1])) as rnd, \
patch('nessai.reparameterisations.rescale_zero_to_one',
side_effect=lambda x, *args: (x, np.array([5, 6]))) as f:
x_out, x_prime_out, log_j_out = RescaleToBounds._apply_inversion(
reparam,
x,
x_prime,
log_j,
'x',
'x_prime',
False,
)
rnd.assert_called_once_with(2, 1, replace=False)
assert f.call_args_list[0][0][1] == 0.0
assert f.call_args_list[0][0][2] == 5.0
# Output should be x_val minus offset
# Then 1 - that for 'upper'
# Then *= -1 with the output of rnd
np.testing.assert_array_almost_equal(
x_prime_out['x_prime'],
np.array([0.8, -0.3]),
decimal=10,
)
# x should be the same
assert x_out is x
# Jacobian should just include jacobian from rescaling
np.testing.assert_array_equal(log_j_out, np.array([5, 6]))
def test_reparameterise(reparam, n):
"""Test the reparameterise method"""
reparam.parameters = ['x', 'y']
reparam.prime_parameters = ['x_prime', 'y_prime']
reparam.scale = {'x': -2.0, 'y': 4.0}
x = numpy_array_to_live_points(np.ones((n, 2)), reparam.parameters)
x_prime = numpy_array_to_live_points(np.zeros((n, 2)),
reparam.prime_parameters)
log_j = np.zeros(n)
x_out, x_prime_out, log_j_out = \
Rescale.reparameterise(reparam, x, x_prime, log_j)
assert np.array_equal(x, x_out)
assert np.array_equal(log_j_out, -np.log(8 * np.ones(n)))
assert (x_prime_out['x_prime'] == -0.5).all()
assert (x_prime_out['y_prime'] == 0.25).all()
def test_empty_numpy_array_to_live_points(empty_live_point):
"""
Test the function the produces an array of live points given an empty
numpy array
"""
np.testing.assert_array_equal(
empty_live_point,
lp.numpy_array_to_live_points(np.array([]), names=['x', 'y', 'z']))
def test_inverse_reparameterise_overflow(reparam, scale):
"""Test the inverse_reparameterise method with very small and large scales.
"""
reparam.parameters = ['x']
reparam.prime_parameters = ['x_prime']
reparam.scale = {'x': scale}
x_array = np.arange(100.0, dtype=float)
x = numpy_array_to_live_points(np.ones((x_array.size, 1)),
reparam.parameters)
x_prime = numpy_array_to_live_points(x_array[:, np.newaxis],
reparam.prime_parameters)
log_j = np.zeros(x.size)
x_out, x_prime_out, log_j_out = \
Rescale.inverse_reparameterise(reparam, x, x_prime, log_j)
np.testing.assert_array_equal(x_array * scale, x_out['x'])
assert (log_j == np.log(scale)).all()
def test_new_point_log_prob(model):
"""Test the log prob for new points.
Should be zero.
"""
x = numpy_array_to_live_points(np.random.randn(2, 1), ['x'])
log_prob = Model.new_point_log_prob(model, x)
assert log_prob.size == 2
assert (log_prob == 0).all()
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_in_bounds(model):
"""Test the `in_bounds` method.
Tests both finite and infinite prior bounds.
"""
x = numpy_array_to_live_points(np.array([[0.5, 1], [2, 1]]), ['x', 'y'])
model.names = ['x', 'y']
model.bounds = {'x': [0, 1], 'y': [-np.inf, np.inf]}
val = Model.in_bounds(model, x)
np.testing.assert_array_equal(val, np.array([True, False]))
def test_pre_rescaling_integration(is_invertible, model):
"""Test the pre-scaling feature"""
def forward(x):
return np.log(x), -np.log(x)
def inv(x):
return np.exp(x), x.copy()
reparam = RescaleToBounds(
parameters='x',
prior_bounds={'x': [1.0, np.e]},
pre_rescaling=(forward, inv),
rescale_bounds=[-1.0, 1.0],
)
x = numpy_array_to_live_points(
np.array([[1.0], [np.e**0.5], [2.0], [np.e]]), ['x'])
x_prime = numpy_array_to_live_points(np.empty([x.size, 1]), ['x_prime'])
log_j = np.zeros(x.size)
x_out, x_prime_out, log_j_out = reparam.reparameterise(x, x_prime, log_j)
np.testing.assert_array_equal(x_out, x)
np.testing.assert_array_equal(x_prime_out['x_prime'],
np.array([-1, 0.0, 2 * np.log(2) - 1, 1]))
np.testing.assert_array_equal(log_j_out, -np.log(x['x']) + np.log(2))
x_in = numpy_array_to_live_points(np.empty([x_prime_out.size, 1]), ['x'])
log_j = np.zeros(x.size)
x_out, x_prime_final, log_j_final = \
reparam.inverse_reparameterise(x_in, x_prime_out, log_j)
np.testing.assert_array_equal(log_j_final, np.log(x_out['x']) - np.log(2))
np.testing.assert_array_equal(x_out['x'], x['x'])
np.testing.assert_array_equal(x_prime_final, x_prime_out)
np.testing.assert_array_equal(log_j_final, -log_j_out)
# Trick to get a 1d model to test only x
model._names.remove('y')
model._bounds = {'x': [1.0, np.e]}
assert is_invertible(reparam, model=model, decimal=12)
def test_inverse_rescale_angle(reparam):
"""Test the inverse method for rescaling the 'angle'."""
reparam.parameters = ['x']
reparam.prior_bounds = {'x': [0, 2]}
x = numpy_array_to_live_points(np.array([[0.], [-0.5], [1.0]]), ['x'])
expected_x = x.copy()
expected_x['x'] = np.array([0., 1.0, 2.0])
x_prime = numpy_array_to_live_points(
np.array([0., 1.0, 3.0]), ['x_prime'])
log_j = np.zeros(3)
x_out, x_prime_out, log_j_out = ToCartesian._inverse_rescale_angle(
reparam, x, x_prime, log_j
)
np.testing.assert_array_equal(x_out, expected_x)
np.testing.assert_array_equal(x_prime_out, x_prime)
np.testing.assert_equal(log_j_out, np.log(2))
def test_check_prior_bounds(proposal):
"""Test the check prior bounds method."""
x = numpy_array_to_live_points(np.arange(10)[:, np.newaxis], ['x'])
proposal.model = Mock()
proposal.model.names = ['x']
proposal.model.bounds = {'x': [0, 5]}
y = np.arange(10)
x_out, y_out = FlowProposal.check_prior_bounds(proposal, x, y)
np.testing.assert_array_equal(x_out, x[:6])
np.testing.assert_array_equal(y_out, y[:6])
def test_convert_to_samples(proposal):
"""Test convert to sample without the prime prior"""
samples = numpy_array_to_live_points(np.random.randn(10, 4), ['x', 'y'])
proposal.use_x_prime_prior = False
proposal.model = MagicMock()
proposal.model.names = ['x']
proposal.model.log_prior = MagicMock(return_value=np.ones(10))
out_samples = FlowProposal.convert_to_samples(proposal, samples, plot=True)
assert out_samples.dtype.names == ('x', 'logP', 'logL')
def test_inverse_reparameterise_boundary_inversion(reparam):
"""Test the inverse_reparameterise function with boundary inversion"""
reparam.has_pre_rescaling = False
reparam.has_post_rescaling = False
reparam.parameters = ['x']
reparam.prime_parameters = ['x']
reparam.offsets = {'x': 1.0}
reparam.boundary_inversion = {'x': 'split'}
x_prime = numpy_array_to_live_points(np.array([(-1.0, ), (2.0, )]),
reparam.prime_parameters)
inversion_out = numpy_array_to_live_points(np.array([(1.0, ), (2.0, )]),
reparam.parameters)
x_in = np.zeros([
2,
], dtype=get_dtype(reparam.parameters))
log_j = np.zeros(x_prime.size)
x_ex = inversion_out
x_prime_ex = x_prime
log_j_ex = np.array([0, 0.5])
reparam._reverse_inversion = MagicMock(return_value=(x_ex, x_prime_ex,
log_j_ex))
x_out, x_prime_out, log_j_out = RescaleToBounds.inverse_reparameterise(
reparam,
x_in,
x_prime,
log_j,
)
np.testing.assert_array_equal(
reparam._reverse_inversion.call_args_list[0][0][0], x_in)
np.testing.assert_array_equal(
reparam._reverse_inversion.call_args_list[0][0][1], x_prime)
np.testing.assert_array_equal(
reparam._reverse_inversion.call_args_list[0][0][2], log_j)
np.testing.assert_array_equal(x_out, x_ex)
np.testing.assert_array_equal(x_prime_out, x_prime_ex)
np.testing.assert_array_equal(log_j_out, log_j_ex)
def test_save_live_points(tmp_path):
"""Test the function for saving live points"""
d = {'x': [1, 2], 'y': [3, 4], 'logP': [0, 0], 'logL': [0, 0]}
live_points = numpy_array_to_live_points(np.array([[1, 3], [2, 4]]),
['x', 'y'])
filename = str(tmp_path) + 'test.json'
save_live_points(live_points, filename)
with open(filename, 'r') as fp:
d_out = json.load(fp)
assert d_out == d
def test_reparameterise_negative_radius(reparam):
"""Assert an error is radius if the radius is negative."""
x = numpy_array_to_live_points(
np.array([[1.0, 1.0, -1.0]]), ['theta', 'phi', 'radius']
)
x_prime = x.copy()
log_j = 0.0
reparam.radial = 'radius'
reparam.chi = None
with pytest.raises(RuntimeError) as excinfo:
AnglePair.reparameterise(reparam, x, x_prime, log_j)
assert 'Radius cannot be negative' in str(excinfo.value)
def test_compute_weights(proposal):
"""Test the compute weights method"""
x = numpy_array_to_live_points(np.array([[1], [2], [3]]), 'x')
proposal.model = Mock()
proposal.model.log_prior = Mock(return_value=np.array([6, 6, 6]))
proposal.log_proposal = Mock(return_value=np.array([3, 4, np.nan]))
log_w = np.array([0, -1, np.nan])
out = RejectionProposal.compute_weights(proposal, x)
proposal.model.log_prior.assert_called_once_with(x)
proposal.log_proposal.assert_called_once_with(x)
np.testing.assert_array_equal(out, log_w)
def test_plot_pool_1d(proposal, tmpdir, alt_dist):
"""Test `plot_pool` when plotting is not 'all'.
Test cases when there is an alternative base distribution is used and
when one is not used.
"""
output = tmpdir.mkdir('test_plot_pool_1d')
proposal.output = output
proposal._plot_pool = True
proposal.populated_count = 0
z = np.random.randn(10, 2)
x = numpy_array_to_live_points(np.random.randn(10, 2), ['x', 'y'])
training_data = \
numpy_array_to_live_points(np.random.randn(10, 2), ['x', 'y'])
proposal.training_data = training_data
log_p = torch.arange(10)
proposal.flow = MagicMock()
proposal.flow.device = 'cpu'
if alt_dist:
proposal.alt_dist = MagicMock()
proposal.alt_dist.log_prob = MagicMock(return_value=log_p)
else:
proposal.flow.model.base_distribution_log_prob = \
MagicMock(return_value=log_p)
proposal.alt_dist = None
with patch('nessai.proposal.flowproposal.plot_1d_comparison') as plot:
FlowProposal.plot_pool(proposal, z, x)
plot.assert_called_once_with(training_data,
x,
labels=['live points', 'pool'],
filename=f'{output}/pool_0.png')
assert os.path.exists(f'{output}/pool_0_log_q.png')
if alt_dist:
proposal.alt_dist.log_prob.assert_called_once()
else:
proposal.flow.model.base_distribution_log_prob.assert_called_once()
