def test_jacobian(self, x, coefficient):
factor = mp.Factor(lambda p: coefficient * p, p=x)
assert factor.jacobian(
{x: 2},
[x],
)[x] == pytest.approx(coefficient)
def make_model_approx(centres, widths):
centres_ = [Variable(f"x_{i}") for i in range(n)]
mu_ = Variable("mu")
logt_ = Variable("logt")
centre_likelihoods = [
NormalMessage(c, w).as_factor(x)
for c, w, x in zip(centres, widths, centres_)
]
normal_likelihoods = [
graph.Factor(
normal_loglike_t,
centre,
mu_,
logt_,
factor_jacobian=normal_loglike_t_jacobian,
) for centre in centres_
]
model = graph.utils.prod(centre_likelihoods + normal_likelihoods)
model_approx = graph.EPMeanField.from_approx_dists(
model,
{
mu_: NormalMessage(0, 10),
logt_: NormalMessage(0, 10),
**{x_: NormalMessage(0, 10)
for x_ in centres_},
},
)
return model_approx
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 make_linear_factor_jac(x_, a_, b_, z_):
return graph.Factor(
linear,
x_,
a_,
b_,
vjp=True,
# factor_jacobian=linear_jacobian,
factor_out=z_)
def make_model():
prior_norm = stats.norm(loc=0, scale=prior_std)
def prior(x):
return prior_norm.logpdf(x).sum()
def linear(x, a, b):
return x.dot(a) + b
linear_factor = graph.Factor(
linear, x_, a_, b_,
factor_out=z_,
vjp=True,
)
likelihood_factor = messages.NormalMessage(y, np.full_like(y, error_std)).as_factor(z_)
prior_a = graph.Factor(prior, a_)
prior_b = graph.Factor(prior, b_)
model = likelihood_factor * linear_factor * prior_a * prior_b
return model
def test_full(data):
samples = {
Variable(f"samples_{i}"): sample
for i, sample in enumerate(data)
}
x_i_ = [Variable(f"x_{i}") for i in range(n)]
logt_i_ = [Variable(f"logt_{i}") for i in range(n)]
mu_x_ = Variable("mu_x")
logt_x_ = Variable("logt_x")
mu_logt_ = Variable("mu_logt")
logt_logt_ = Variable("logt_logt")
hierarchical_params = (mu_x_, logt_x_, mu_logt_, logt_logt_)
# Setting up model
data_loglikes = [
graph.Factor(
normal_loglike_t,
s_,
x_,
logt_,
factor_jacobian=normal_loglike_t_jacobian,
name=f"normal_{i}",
) for i, (s_, x_, logt_) in enumerate(zip(samples, x_i_, logt_i_))
]
centre_loglikes = [
graph.Factor(normal_loglike_t, x_, mu_x_, logt_x_) for x_ in x_i_
]
precision_loglikes = [
graph.Factor(normal_loglike_t, logt_, mu_logt_, logt_logt_)
for logt_ in logt_i_
]
priors = [
messages.NormalMessage(0, 10).as_factor(v, name=f"prior_{v.name}")
for v in hierarchical_params
]
model = graph.utils.prod(data_loglikes + centre_loglikes +
precision_loglikes + priors)
def test_plates(self):
obs = autofit.mapper.variable.Plate(name='obs')
dims = autofit.mapper.variable.Plate(name='dims')
def sub(a, b):
return a - b
a = autofit.mapper.variable.Variable('a', obs, dims)
b = autofit.mapper.variable.Variable('b', dims)
subtract = mp.Factor(sub, a=a, b=b)
x = np.array([[1, 2, 3], [4, 5, 6]])
y = np.array([1, 2, 1])
value = subtract({a: x, b: y}).log_value
assert (value == x - y).all()
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_hierarchical(centres, widths):
centres_ = [Variable(f"x_{i}") for i in range(n)]
mu_ = Variable("mu")
logt_ = Variable("logt")
centre_likelihoods = [
messages.NormalMessage(c, w).as_factor(x)
for c, w, x in zip(centres, widths, centres_)
]
hierarchical_factor = graph.Factor(
hierarchical_loglike_t,
mu_,
logt_,
*centres_,
factor_jacobian=hierarchical_loglike_t_jac,
)
model = graph.utils.prod(centre_likelihoods) * hierarchical_factor
model_approx = graph.EPMeanField.from_approx_dists(
model,
{
mu_: messages.NormalMessage(0.0, 10.0),
logt_: messages.NormalMessage(0.0, 10.0),
**{x_: messages.NormalMessage(0.0, 10.0)
for x_ in centres_},
},
)
laplace = graph.LaplaceOptimiser()
ep_opt = graph.EPOptimiser(model_approx, default_optimiser=laplace)
new_approx = ep_opt.run(model_approx, max_steps=10)
print(new_approx)
mu_ = new_approx.factor_graph.name_variable_dict["mu"]
logt_ = new_approx.factor_graph.name_variable_dict["logt"]
assert new_approx.mean_field[mu_].mean == pytest.approx(np.mean(centres),
rel=0.2)
assert new_approx.mean_field[logt_].mean == pytest.approx(np.log(
np.std(centres)**-2),
rel=0.2)
def make_flat_compound(plus, y, sigmoid):
phi = graph.Factor(log_phi, y)
return phi * plus * sigmoid
def make_likelihood_factor(likelihood, z_, y_):
return mp.Factor(likelihood, z=z_, y=y_)
def make_prior_b(prior, b_):
return mp.Factor(prior, x=b_)
def make_prior_a(prior, a_):
return mp.Factor(prior, x=a_)
def make_linear_factor(
x_, a_, b_, z_
):
return mp.Factor(linear, x=x_, a=a_, b=b_) == z_
def make_flat_compound(plus, y, sigmoid):
g = plus == y
phi = mp.Factor(log_phi, x=y)
return phi * g * sigmoid
def make_linear_factor(x_, a_, b_, z_):
return graph.Factor(linear, x_, a_, b_, factor_out=z_)
def make_likelihood_factor(likelihood, z_, y_):
return graph.Factor(likelihood, z_, y_)
def make_probit_factor(x):
return graph.Factor(stats.norm(loc=0.0, scale=1.0).logcdf, x)
def make_likelihood_factor_jac(z_, y_, obs, dims):
factor = graph.Factor(likelihood, z_, y_, factor_jacobian=likelihood_jacobian)
return factor
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)
def make_plus(x, y):
return graph.Factor(plus_two, x, factor_out=y)
def make_phi(x):
return graph.Factor(log_phi, x)
def make_vectorised_sigmoid(x):
return graph.Factor(log_sigmoid, x)
def make_sigmoid(x):
return graph.Factor(log_sigmoid, x)
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_jacobian(self, x, coefficient):
factor = graph.Factor(lambda p: coefficient * p, x)
assert factor.jacobian({x: 2}).grad()[x] == pytest.approx(coefficient)
def make_likelihood_factor(z_, y_, obs, dims):
factor = graph.Factor(likelihood, z_, y_)
return factor
def make_prior_a(prior, a_):
return graph.Factor(prior, a_)
def _test():
## define parameters of model
np.random.seed(1)
alpha, beta, gamma, delta = 2 / 3, 4 / 3, 1, 1
r = np.array([alpha, -gamma])
A = np.array([[0., beta / alpha], [delta / gamma, 0.]])
K = 1
noise = 0.1
# starting composition
y0 = np.array([1., 1.])
n_species = len(y0)
n_obs = 30
t_space = 1.
t_obs = np.r_[0, (np.arange(n_obs - 1) * t_space +
np.random.rand(n_obs - 1)) * t_space]
def lotka_volterra(t, z, r=r, A=A, K=K):
return z * r * (1 - A.dot(z) / K)
def calc_lotka_volterra(y0, r, A, K, t_obs):
res = integrate.solve_ivp(lotka_volterra, (t_obs[0], t_obs[-1]),
y0,
t_eval=t_obs,
args=(r, A, K),
method='BDF')
y_ = res.y
n = y_.shape[1]
n_obs = len(t_obs)
# make sure output is correct dimension
if n != n_obs:
y_ = np.c_[y_, np.repeat(y_[:,
[-1]], n_obs - n, axis=1)][:, :n_obs]
if y_.shape[1] != n_obs:
raise Exception
return y_
y_true = calc_lotka_volterra(y0, r, A, K, t_obs)
y = y_true + noise * np.random.randn(n_species, n_obs)
## Specifying dimensions of problem
obs = autofit.mapper.variable.Plate(name='obs')
species = autofit.mapper.variable.Plate(name='species')
# Need to specify a second plate for species because
# A is (species, species) and we need a second plate
# to unique specify the second dimension
speciesA = autofit.mapper.variable.Plate(name='species')
dims = autofit.mapper.variable.Plate(name='dims')
## Specifying variables
r_ = autofit.mapper.variable.Variable('r', species)
A_ = autofit.mapper.variable.Variable('A', species, speciesA)
K_ = autofit.mapper.variable.Variable('K')
y0_ = autofit.mapper.variable.Variable('y0', species)
y_ = autofit.mapper.variable.Variable('y', species, obs)
y_obs_ = autofit.mapper.variable.Variable('y_obs', species, obs)
t_obs_ = autofit.mapper.variable.Variable('t_obs', obs)
_norm = stats.norm(loc=0, scale=noise)
_prior = stats.norm(loc=0, scale=10)
_prior_exp = stats.expon(loc=0, scale=1)
def _likelihood(y_obs, y):
return _norm.logpdf(y_obs - y)
## Specifying factors
likelihood = mp.Factor(_likelihood, y_obs=y_obs_, y=y_)
prior_A = mp.Factor(_prior.logpdf, 'prior_A', x=A_)
prior_r = mp.Factor(_prior.logpdf, 'prior_r', x=r_)
prior_y0 = mp.Factor(_prior_exp.logpdf, 'prior_y0', x=y0_)
# calc_lotka_volterra does not vectorise over
# multiple inputs, see `FactorNode._py_vec_call`
LV = mp.Factor(calc_lotka_volterra,
'LV',
vectorised=False,
y0=y0_,
r=r_,
A=A_,
K=K_,
t_obs=t_obs_) == y_
## Defining model
priors = prior_A * prior_r * prior_y0
LV_model = (likelihood * LV) * priors
LV_model._name = 'LV_model'
model_approx = mp.EPMeanField.from_kws(
LV_model,
{
A_:
autofit.graphical.messages.normal.NormalMessage.from_mode(A, 100.),
r_:
autofit.graphical.messages.normal.NormalMessage.from_mode(r, 100.),
y0_:
autofit.graphical.messages.gamma.GammaMessage.from_mode(
np.ones_like(y0), 1),
y_:
autofit.graphical.messages.normal.NormalMessage.from_mode(y, 1),
K_:
autofit.graphical.messages.fixed.FixedMessage(1),
y_obs_:
autofit.graphical.messages.fixed.FixedMessage(y),
t_obs_:
autofit.graphical.messages.fixed.FixedMessage(t_obs)
},
)
history = {}
n_iter = 1
factors = [f for f in LV_model.factors if f not in (LV, )]
np.random.seed(1)
opt = mp.optimise.LaplaceOptimiser(n_iter=n_iter)
for i in range(n_iter):
# perform least squares fit for LV model
model_approx, status = mp.lstsq_laplace_factor_approx(model_approx, LV)
# perform laplace non linear fit for other factors
for factor in factors:
model_approx, status = mp.optimise.laplace_factor_approx(
model_approx,
factor,
status=status,
)
history[i, factor] = model_approx
# model_mean = {v: d.mean for v, d in model_approx.mean_field.items()}
# y_pred = LV_model(model_mean).deterministic_values[y_]
y_pred = model_approx.mean_field[y_].mean
assert np.square(y_pred - y).mean()**0.5 < 2
def make_prior_b(prior, b_):
return graph.Factor(prior, b_)
