def test_multivariate_message(self): p1, p2, p3 = mp.Plate(), mp.Plate(), mp.Plate() x_ = mp.Variable('x', p3, p1) y_ = mp.Variable('y', p1, p2) z_ = mp.Variable('z', p2, p3) n1, n2, n3 = shape = (2, 3, 4) def sumxyz(x, y, z): return (np.moveaxis(x[:, :, None], 0, 2) + y[:, :, None] + z[None]) factor = mp.Factor(sumxyz, x=x_, y=y_, z=z_) x = np.arange(n3 * n1).reshape(n3, n1) * 0.1 y = np.arange(n1 * n2).reshape(n1, n2) * 0.2 z = np.arange(n2 * n3).reshape(n2, n3) * 0.3 sumxyz(x, y, z) variables = {x_: x, y_: y, z_: z} factor(variables) model_dist = mp.MeanField({ x_: mp.NormalMessage(x, 1 * np.ones_like(x)), y_: mp.NormalMessage(y, 1 * np.ones_like(y)), z_: mp.NormalMessage(z, 1 * np.ones_like(z)), }) assert model_dist(variables).log_value.shape == shape
def test_factor_jacobian(): shape = 4, 3 z_ = mp.Variable('z', *(mp.Plate() for _ in shape)) likelihood = mp.NormalMessage(np.random.randn(*shape), np.random.exponential(size=shape)) likelihood_factor = likelihood.as_factor(z_) values = {z_: likelihood.sample()} fval, jval = likelihood_factor.func_jacobian(values, axis=None) ngrad = approx_fprime(values[z_].ravel(), lambda x: likelihood.logpdf(x.reshape(*shape)).sum(), 1e-8).reshape(*shape) assert np.allclose(ngrad, jval[z_])
def test_meanfield_gradients(): n1, n2, n3 = 2, 3, 5 p1, p2, p3 = [graph.Plate() for i in range(3)] v1 = graph.Variable('v1', p1, p2) v2 = graph.Variable('v2', p2, p3) v3 = graph.Variable('v3', p3, p1) mean_field = graph.MeanField({ v1: graph.NormalMessage(np.random.randn(n1, n2), np.random.exponential(size=(n1, n2))), v2: graph.NormalMessage(np.random.randn(n2, n3), np.random.exponential(size=(n2, n3))), v3: graph.NormalMessage(np.random.randn(n3, n1), np.random.exponential(size=(n3, n1))) }) values = mean_field.sample() l0 = mean_field(values, axis=None) logl = mean_field(values, axis=False) assert logl.sum() == pytest.approx(l0, abs=1e-5) logl = mean_field(values, axis=1) assert logl.sum() == pytest.approx(l0, abs=1e-5) logl = mean_field(values, axis=2) assert logl.sum() == pytest.approx(l0, abs=1e-5) logl = mean_field(values, axis=(0, 2)) assert logl.sum() == pytest.approx(l0, abs=1e-5) njac0 = mean_field._numerical_func_jacobian(values, axis=None, _eps=1e-8)[1] njac1 = mean_field._numerical_func_jacobian(values, axis=1, _eps=1e-8)[1] njac2 = mean_field._numerical_func_jacobian(values, axis=(0, 1), _eps=1e-8)[1] njac = mean_field._numerical_func_jacobian_hessian(values, axis=False, _eps=1e-8)[1] grad = mean_field.logpdf_gradient(values, axis=False)[1] for v in grad: norm = np.linalg.norm(grad[v] - njac[v].sum((0, 1, 2))) assert norm == pytest.approx(0, abs=1e-2) norm = np.linalg.norm(grad[v] - njac0[v]) assert norm == pytest.approx(0, abs=1e-2) norm = np.linalg.norm(grad[v] - njac1[v].sum((0, 1))) assert norm == pytest.approx(0, abs=1e-2) norm = np.linalg.norm(grad[v] - njac2[v].sum(0)) assert norm == pytest.approx(0, abs=1e-2)
def test_simple_transform_diagonal(): # testing DiagonalTransform d = 5 scale = np.random.exponential(size=d) A = np.diag(scale**-1) b = np.random.randn(d) x = graph.Variable('x', graph.Plate()) x0 = np.random.randn(d) param_shapes = graph.utils.FlattenArrays({x: (d, )}) def likelihood(x): x = x - b return 0.5 * np.linalg.multi_dot((x, A, x)) factor = graph.Factor(likelihood, x=x, is_scalar=True) func = factor.flatten(param_shapes) res = optimize.minimize(func, x0) assert np.allclose(res.x, b, rtol=1e-2, atol=1e-2) H, iA = res.hess_inv, np.linalg.inv(A) # check R2 score assert 1 - np.square(H - iA).mean() / np.square(iA).mean() > 0.95 scale = np.random.exponential(size=d) A = np.diag(scale**-2) diag = transform.DiagonalTransform(scale) whiten = transform.VariableTransform({x: diag}) white_factor = transform.TransformedNode(factor, whiten) white_func = white_factor.flatten(param_shapes) y0 = diag * x0 res = optimize.minimize(white_func, y0) assert np.allclose(res.x, diag * b) H, iA = res.hess_inv, np.eye(d) # check R2 score assert 1 - np.square(H - iA).mean() / np.square(iA).mean() > 0.95 # testing gradients grad = white_func.jacobian(y0) ngrad = optimize.approx_fprime(y0, white_func, 1e-6) assert np.allclose(grad, ngrad, atol=1e-3, rtol=1e-3)
def test_simple_transform_cholesky(): np.random.seed(0) d = 5 A = stats.wishart(d, np.eye(d)).rvs() b = np.random.rand(d) def likelihood(x): x = x - b return 0.5 * np.linalg.multi_dot((x, A, x)) x = graph.Variable('x', graph.Plate()) x0 = np.random.randn(d) factor = graph.Factor(likelihood, x=x, is_scalar=True) param_shapes = graph.utils.FlattenArrays({x: (d, )}) func = factor.flatten(param_shapes) res = optimize.minimize(func, x0) assert np.allclose(res.x, b, rtol=1e-2) H, iA = res.hess_inv, np.linalg.inv(A) # check R2 score assert 1 - np.square(H - iA).mean() / np.square(iA).mean() > 0.95 # cho = transform.CholeskyTransform(linalg.cho_factor(A)) cho = transform.CholeskyTransform.from_dense(A) whiten = transform.VariableTransform({x: cho}) white_factor = transform.TransformedNode(factor, whiten) white_func = white_factor.flatten(param_shapes) y0 = cho * x0 res = optimize.minimize(white_func, y0) assert np.allclose(res.x, cho * b, atol=1e-3, rtol=1e-3) assert np.allclose(res.hess_inv, np.eye(d), atol=1e-3, rtol=1e-3) # testing gradients grad = white_func.jacobian(y0) ngrad = optimize.approx_fprime(y0, white_func, 1e-6) assert np.allclose(grad, ngrad, atol=1e-3, rtol=1e-3) # testing CovarianceTransform, cho = transform.CovarianceTransform.from_dense(iA) whiten = transform.VariableTransform({x: cho}) white_factor = transform.TransformedNode(factor, whiten) white_func = white_factor.flatten(param_shapes) y0 = cho * x0 res = optimize.minimize(white_func, y0) assert np.allclose(res.x, cho * b, atol=1e-3, rtol=1e-3) assert np.allclose(res.hess_inv, np.eye(d), atol=1e-3, rtol=1e-3) # testing gradients grad = white_func.jacobian(y0) ngrad = optimize.approx_fprime(y0, white_func, 1e-6) assert np.allclose(grad, ngrad, atol=1e-3, rtol=1e-3) # testing FullCholeskyTransform whiten = transform.FullCholeskyTransform(cho, param_shapes) white_factor = transform.TransformedNode(factor, whiten) white_func = white_factor.flatten(param_shapes) y0 = cho * x0 res = optimize.minimize(white_func, y0) assert np.allclose(res.x, cho * b, atol=1e-3, rtol=1e-3) assert np.allclose(res.hess_inv, np.eye(d), atol=1e-3, rtol=1e-3) # testing gradients grad = white_func.jacobian(y0) ngrad = optimize.approx_fprime(y0, white_func, 1e-6) assert np.allclose(grad, ngrad, atol=1e-3, rtol=1e-3)
def test_complex_transform(): n1, n2, n3 = 2, 3, 2 d = n1 + n2 * n3 A = stats.wishart(d, np.eye(d)).rvs() b = np.random.rand(d) p1, p2, p3 = (graph.Plate() for i in range(3)) x1 = graph.Variable('x1', p1) x2 = graph.Variable('x2', p2, p3) mean_field = graph.MeanField({ x1: graph.NormalMessage(np.zeros(n1), 100 * np.ones(n1)), x2: graph.NormalMessage(np.zeros((n2, n3)), 100 * np.ones((n2, n3))), }) values = mean_field.sample() param_shapes = graph.utils.FlattenArrays( {v: x.shape for v, x in values.items()}) def likelihood(x1, x2): x = np.r_[x1, x2.ravel()] - b return 0.5 * np.linalg.multi_dot((x, A, x)) factor = graph.Factor(likelihood, x1=x1, x2=x2, is_scalar=True) cho = transform.CholeskyTransform(linalg.cho_factor(A)) whiten = transform.FullCholeskyTransform(cho, param_shapes) trans_factor = transform.TransformedNode(factor, whiten) values = mean_field.sample() transformed = whiten * values assert np.allclose(factor(values), trans_factor(transformed)) njac = trans_factor._numerical_func_jacobian(transformed)[1] jac = trans_factor.jacobian(transformed) ngrad = param_shapes.flatten(njac) grad = param_shapes.flatten(jac) assert np.allclose(grad, ngrad, atol=1e-3, rtol=1e-3) # test VariableTransform with CholeskyTransform var_cov = { v: (X.reshape((int(X.size**0.5), ) * 2)) for v, X in param_shapes.unflatten(linalg.inv(A)).items() } cho_factors = { v: transform.CholeskyTransform(linalg.cho_factor(linalg.inv(cov))) for v, cov in var_cov.items() } whiten = transform.VariableTransform(cho_factors) trans_factor = transform.TransformedNode(factor, whiten) values = mean_field.sample() transformed = whiten * values assert np.allclose(factor(values), trans_factor(transformed)) njac = trans_factor._numerical_func_jacobian(transformed)[1] jac = trans_factor.jacobian(transformed) ngrad = param_shapes.flatten(njac) grad = param_shapes.flatten(jac) assert np.allclose(grad, ngrad, atol=1e-3, rtol=1e-3) res = optimize.minimize(trans_factor.flatten(param_shapes).func_jacobian, param_shapes.flatten(transformed), method='BFGS', jac=True) assert res.hess_inv.diagonal() == pytest.approx(1., rel=1e-1) # test VariableTransform with CholeskyTransform diag_factors = { v: transform.DiagonalTransform(cov.diagonal()**0.5) for v, cov in var_cov.items() } whiten = transform.VariableTransform(diag_factors) trans_factor = transform.TransformedNode(factor, whiten) values = mean_field.sample() transformed = whiten * values assert np.allclose(factor(values), trans_factor(transformed)) njac = trans_factor._numerical_func_jacobian(transformed)[1] jac = trans_factor.jacobian(transformed) ngrad = param_shapes.flatten(njac) grad = param_shapes.flatten(jac) assert np.allclose(grad, ngrad, atol=1e-3, rtol=1e-3) res = optimize.minimize(trans_factor.flatten(param_shapes).func_jacobian, param_shapes.flatten(transformed), method='BFGS', jac=True) assert res.hess_inv.diagonal() == pytest.approx(1., rel=1e-1)
np.random.seed(1) error_std = 1. prior_std = 10. a = np.array([[-1.3], [0.7]]) b = np.array([-0.5]) n_obs = 100 n_features, n_dims = a.shape x = 5 * np.random.randn(n_obs, n_features) y = x.dot(a) + b + np.random.randn(n_obs, n_dims) obs = graph.Plate(name='obs') features = graph.Plate(name='features') dims = graph.Plate(name='dims') x_ = graph.Variable('x', obs, features) a_ = graph.Variable('a', features, dims) b_ = graph.Variable('b', dims) y_ = graph.Variable('y', obs, dims) z_ = graph.Variable('z', obs, dims) def make_model(): prior_norm = stats.norm(loc=0, scale=prior_std) def prior(x): return prior_norm.logpdf(x).sum()