def test_rescaling_generate_unknown(proposal, x): """Test the rescaling method with with an unknown method for generate""" proposal._base_rescale = MagicMock(return_value=[x, np.ones(x.size)]) with pytest.raises(RuntimeError) as excinfo: AugmentedFlowProposal._augmented_rescale(proposal, x, generate_augment='ones') assert 'Unknown method' in str(excinfo.value)
def test_set_rescaling(mock, proposal): """Test the set rescaling method""" proposal.names = ['x', 'y'] proposal.rescaled_names = ['x_prime', 'y_prime'] proposal.rescale_parameters = proposal.names proposal.augment_dims = 2 AugmentedFlowProposal.set_rescaling(proposal) assert proposal.names == ['x', 'y', 'e_0', 'e_1'] assert proposal.rescaled_names == ['x_prime', 'y_prime', 'e_0', 'e_1'] assert proposal.augment_names == ['e_0', 'e_1'] mock.assert_called_once()
def test_update_flow_config(proposal): """Test update flow config""" proposal.rescaled_dims = 4 proposal.augment_dims = 2 proposal.flow_config = dict(model_config={}) with patch('nessai.proposal.augmented.FlowProposal.update_flow_config' ) as mock: AugmentedFlowProposal.update_flow_config(proposal) mock.assert_called_once() mask = np.array([1, 1, -1, -1]) np.testing.assert_array_equal( proposal.flow_config['model_config']['kwargs']['mask'], mask)
def test_backward_pass(proposal, model, log_p, marg): """Test the backward pass method""" n = 2 acc = int(np.isfinite(log_p).sum()) x = np.random.randn(n, model.dims) z = np.random.randn(n, model.dims) proposal._marginalise_augment = MagicMock(return_value=log_p) proposal.inverse_rescale = \ MagicMock(side_effect=lambda a: (a, np.ones(a.size))) proposal.rescaled_names = model.names proposal.alt_dist = None proposal.check_prior_bounds = MagicMock(side_effect=lambda a, b: (a, b)) proposal.flow = MagicMock() proposal.flow.sample_and_log_prob = \ MagicMock(return_value=[x, log_p]) proposal.marginalise_augment = marg x_out, log_p = AugmentedFlowProposal.backward_pass(proposal, z) assert len(x_out) == acc proposal.inverse_rescale.assert_called_once() proposal.flow.sample_and_log_prob.assert_called_once_with(z=z, alt_dist=None) assert proposal._marginalise_augment.called is marg
def test_prime_log_prior(mock_prior, proposal, x): """Test the complete prime log prior""" proposal.augmented_prior = MagicMock(return_value=1) log_p = AugmentedFlowProposal.x_prime_log_prior(proposal, x) mock_prior.assert_called_once_with(x) proposal.augmented_prior.assert_called_once_with(x) assert log_p == 2
def test_rescaling(mock_zeros, mock_randn, proposal, x, generate): """Test the rescaling method""" proposal._base_rescale = MagicMock(return_value=[x, np.ones(x.size)]) proposal.augment_names = ['e_1'] proposal.augment_dims = 1 AugmentedFlowProposal._augmented_rescale(proposal, x, generate_augment=generate, test=True) proposal._base_rescale.assert_called_once_with(x, compute_radius=False, test=True) if generate == 'zeroes': mock_zeros.assert_called_once_with(x.size) else: mock_randn.assert_called_once_with(x.size)
def test_augmented_prior(mock, marg, proposal, x): """Test the augmented prior with and without marginalistion""" proposal.marginalise_augment = marg proposal.augment_names = ['e_1', 'e_2'] log_prior = AugmentedFlowProposal.augmented_prior(proposal, x) if marg: assert log_prior == 0 else: np.testing.assert_array_equal(x['e_1'], mock.call_args_list[0][0][0]) np.testing.assert_array_equal(x['e_2'], mock.call_args_list[1][0][0])
def test_rescaling_generate_none(mock_zeros, mock_randn, proposal, x, compute_radius): """Test the rescaling method with generate_augment=None""" proposal._base_rescale = MagicMock(return_value=[x, np.ones(x.size)]) proposal.augment_names = ['e_1'] proposal.augment_dims = 1 proposal.generate_augment = 'zeros' AugmentedFlowProposal._augmented_rescale(proposal, x, generate_augment=None, test=True, compute_radius=compute_radius) proposal._base_rescale.assert_called_once_with( x, compute_radius=compute_radius, test=True) if not compute_radius: mock_randn.assert_called_once_with(x.size) else: mock_zeros.assert_called_once_with(x.size)
def test_w_default_rescaling(model, tmpdir): """Integration test to make sure augmented proposal works with the default rescaling method. """ x = model.new_point(100) x_out = \ 2 * (x['x'] - model.bounds['x'][0]) / (np.ptp(model.bounds['x'])) - 1 y_out = \ 2 * (x['y'] - model.bounds['y'][0]) / (np.ptp(model.bounds['y'])) - 1 output = tmpdir.mkdir('testdir') proposal = AugmentedFlowProposal(model, output=output, poolsize=100, augment_dims=2, rescale_parameters=True, update_bounds=False) proposal.initialise() assert proposal.rescaled_names == ['x_prime', 'y_prime', 'e_0', 'e_1'] x_prime, log_j = proposal.rescale(x) np.testing.assert_array_equal(x_out, x_prime['x_prime']) np.testing.assert_array_equal(y_out, x_prime['y_prime']) x_inv, log_j_inv = proposal.inverse_rescale(x_prime) np.testing.assert_array_equal(log_j, -log_j_inv) np.testing.assert_array_equal(x[['x', 'y']], x_inv[['x', 'y']])
def test_marginalise_augment(proposal): """Test the marginalise augment function""" proposal.n_marg = 5 proposal.augment_dims = 2 x = np.concatenate([np.random.randn(3, 2), np.zeros((3, 2))], axis=1) z = np.random.randn(15, 4) log_prob = np.random.randn(15) proposal.flow = MagicMock() proposal.flow.forward_and_log_prob = \ MagicMock(return_value=[z, log_prob]) log_prob_out = AugmentedFlowProposal._marginalise_augment(proposal, x) assert len(log_prob_out) == 3
def test_w_reparameterisation(model, tmpdir): """Integration test to make sure augmented proposal works with reparameterisaitons configured using the dictionary method. """ x = model.new_point(100) reparameterisations = { 'x': { 'reparameterisation': 'rescaletobounds', 'update_bounds': False }, 'y': { 'reparameterisation': 'scale', 'scale': 2.0 } } x_out = \ 2 * (x['x'] - model.bounds['x'][0]) / (np.ptp(model.bounds['x'])) - 1 y_out = x['y'] / 2.0 output = tmpdir.mkdir('testdir') proposal = AugmentedFlowProposal(model, output=output, poolsize=100, augment_dims=2, reparameterisations=reparameterisations) proposal.initialise() assert proposal.rescaled_names == ['x_prime', 'y_prime', 'e_0', 'e_1'] x_prime, log_j = proposal.rescale(x) np.testing.assert_array_equal(x_out, x_prime['x_prime']) np.testing.assert_array_equal(y_out, x_prime['y_prime']) x_inv, log_j_inv = proposal.inverse_rescale(x_prime) np.testing.assert_array_equal(log_j, -log_j_inv) np.testing.assert_array_equal(x[['x', 'y']], x_inv[['x', 'y']])
def test_init(model): """Test the init function""" AugmentedFlowProposal(model, poolsize=100)