def test_custom_optimiser(make_model_factor):
other_optimiser = ep.LaplaceOptimiser()
factor_1 = make_model_factor(centre=40, sigma=10, optimiser=other_optimiser)
factor_2 = make_model_factor(centre=60, sigma=15)
factor_model = ep.FactorGraphModel(factor_1, factor_2)
default_optimiser = ep.LaplaceOptimiser()
ep_optimiser = factor_model._make_ep_optimiser(default_optimiser)
factor_optimisers = ep_optimiser.factor_optimisers
assert factor_optimisers[factor_1] is other_optimiser
assert factor_optimisers[factor_2] is default_optimiser
def test_stochastic_linear_regression():
params = [
(50, 5, False),
(20, 60, True),
]
for n_batch, n_iters, inplace in params:
model_approx = make_model_approx()
ep_opt = graph.StochasticEPOptimiser(
model_approx.factor_graph,
graph.LaplaceOptimiser()
)
batches = graph.utils.gen_dict({
obs: graph.utils.gen_subsets(n_batch, n_obs, n_iters=n_iters)
})
new_approx = ep_opt.run(model_approx, batches, inplace=inplace)
mean_field = new_approx.mean_field
X = np.c_[x, np.ones(n_obs)]
XTX = X.T.dot(X) + np.eye(3) / prior_std
cov = np.linalg.inv(XTX)
cov_a = cov[:2, :]
cov_b = cov[2, :]
mean_a = cov_a.dot(X.T.dot(y))
mean_b = cov_b.dot(X.T.dot(y))
a_std = cov_a.diagonal()[:, None]**0.5
b_std = cov_b[[-1]]**0.5
assert mean_field[a_].mean == pytest.approx(mean_a, rel=5e-1), n_batch
assert mean_field[b_].mean == pytest.approx(mean_b, rel=5e-1), n_batch
assert mean_field[a_].sigma == pytest.approx(a_std, rel=2.), n_batch
assert mean_field[b_].sigma == pytest.approx(b_std, rel=2.), n_batch
def _test_optimise_factor_model(factor_model):
laplace = ep.LaplaceOptimiser()
collection = factor_model.optimise(laplace)
assert 25.0 == pytest.approx(collection[0].normalization.mean, rel=0.1)
assert collection[0].normalization is collection[1].normalization
def test_hierarchical_factor(centre_model):
centre_model.add_drawn_variable(af.GaussianPrior(100, 10))
factor = centre_model.factors[0]
assert len(factor.priors) == 3
laplace = g.LaplaceOptimiser()
gaussian = factor.optimise(laplace, max_steps=10)
assert gaussian.instance_from_prior_medians().drawn_prior.mean(
) == pytest.approx(100, abs=1)
def test_model_factor(data, centres):
y = data[0]
centre_argument = af.GaussianPrior(mean=50, sigma=20)
prior_model = af.PriorModel(af.Gaussian,
centre=centre_argument,
normalization=20,
sigma=5)
factor = g.AnalysisFactor(prior_model, analysis=Analysis(x=x, y=y))
laplace = g.LaplaceOptimiser()
gaussian = factor.optimise(laplace, max_steps=10)
assert gaussian.centre.mean == pytest.approx(centres[0], abs=0.1)
def _test_optimise_factor_model(factor_model):
"""
We optimise the model
"""
laplace = ep.LaplaceOptimiser()
collection = factor_model.optimise(laplace)
"""
And what we get back is actually a PriorModelCollection
"""
assert 25.0 == pytest.approx(collection[0].normalization.mean, rel=0.1)
assert collection[0].normalization is collection[1].normalization
def test_approximations(probit_approx, model_approx, x, message):
opt = graph.LaplaceOptimiser()
probit_model_dist, status = opt.optimise_approx(probit_approx)
# get updated factor approximation
probit_project, status = probit_approx.project(probit_model_dist,
delta=1.0)
assert probit_project.model_dist[x].mean == pytest.approx(0.506, rel=0.1)
assert probit_project.model_dist[x].sigma == pytest.approx(0.814, rel=0.1)
assert probit_project.factor_dist[x].mean == pytest.approx(1.499, rel=0.1)
assert probit_project.factor_dist[x].sigma == pytest.approx(1.401, rel=0.1)
def test_simple(model_approx, centres):
laplace = graph.LaplaceOptimiser()
ep_opt = graph.EPOptimiser(model_approx, default_optimiser=laplace)
new_approx = ep_opt.run(model_approx, max_steps=20)
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=1.0)
assert new_approx.mean_field[logt_].mean == pytest.approx(np.log(
np.std(centres)**-2),
rel=1.0)
def test_laplace_method(probit_factor, q_cavity, x):
mf = graph.MeanField({x: q_cavity})
probit_approx = graph.FactorApproximation(
factor=probit_factor,
cavity_dist=mf,
factor_dist=mf,
model_dist=mf,
)
opt = graph.LaplaceOptimiser()
new_dist, s = opt.optimise_approx(probit_approx)
q_probit_laplace = new_dist[x]
assert q_probit_laplace.mean == pytest.approx(-0.258, rel=0.01)
assert q_probit_laplace.sigma == pytest.approx(0.462, rel=0.01)
def test_full_fit(centre_model, data, centres):
graph = g.FactorGraphModel()
for i, y in enumerate(data):
prior_model = af.PriorModel(
af.Gaussian,
centre=af.GaussianPrior(mean=100, sigma=20),
normalization=20,
sigma=5,
)
graph.add(g.AnalysisFactor(prior_model, analysis=Analysis(x=x, y=y)))
centre_model.add_drawn_variable(prior_model.centre)
graph.add(centre_model)
optimiser = g.LaplaceOptimiser()
collection = graph.optimise(optimiser, max_steps=10).model
def test_trivial():
prior = af.UniformPrior(lower_limit=10, upper_limit=20)
prior_model = af.Collection(value=prior)
class TrivialAnalysis(af.Analysis):
def log_likelihood_function(self, instance):
result = -((instance.value - 14)**2)
return result
factor_model = ep.AnalysisFactor(prior_model, analysis=TrivialAnalysis())
optimiser = ep.LaplaceOptimiser()
# optimiser = af.DynestyStatic()
model = factor_model.optimise(optimiser)
assert model.value.mean == pytest.approx(14, rel=0.1)
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 test_gaussian():
n_observations = 100
x = np.arange(n_observations)
y = make_data(Gaussian(centre=50.0, normalization=25.0, sigma=10.0), x)
prior_model = af.PriorModel(
Gaussian,
centre=af.GaussianPrior(mean=50, sigma=20),
normalization=af.GaussianPrior(mean=25, sigma=10),
sigma=af.GaussianPrior(mean=10, sigma=10),
)
factor_model = ep.AnalysisFactor(prior_model, analysis=Analysis(x=x, y=y))
laplace = ep.LaplaceOptimiser()
model = factor_model.optimise(laplace)
assert model.centre.mean == pytest.approx(50, rel=0.1)
assert model.normalization.mean == pytest.approx(25, rel=0.1)
assert model.sigma.mean == pytest.approx(10, rel=0.1)
def test_laplace(
model,
start_approx,
y_,
z_,
):
model_approx = graph.EPMeanField.from_approx_dists(model, start_approx)
laplace = graph.LaplaceOptimiser()
opt = graph.EPOptimiser(model_approx.factor_graph, default_optimiser=laplace)
new_approx = opt.run(model_approx, max_steps=10)
y = new_approx.mean_field[y_].mean
z_pred = new_approx(new_approx.mean_field.mean)[z_]
y_pred = z_pred > 0
(tpr, fpr), (fnr, tnr) = np.dot(
np.array([y, 1 - y]).reshape(2, -1),
np.array([y_pred, 1 - y_pred]).reshape(2, -1).T,
)
accuracy = (tpr + tnr) / (tpr + fpr + fnr + tnr)
assert 0.95 > accuracy > 0.75
def make_laplace():
return g.LaplaceOptimiser()
def test_full_hierachical(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_loglike = graph.Factor(
hierarchical_loglike_t,
mu_x_,
logt_x_,
*x_i_,
name="mean_loglike",
factor_jacobian=hierarchical_loglike_t_jac,
)
logt_loglike = graph.Factor(
hierarchical_loglike_t,
mu_logt_,
logt_logt_,
*logt_i_,
name="logt_loglike",
factor_jacobian=hierarchical_loglike_t_jac,
)
priors = [
messages.NormalMessage(0.0, 10.0).as_factor(v, name=f"prior_{v.name}")
for v in hierarchical_params
]
model = graph.utils.prod(data_loglikes + priors +
[centre_loglike, logt_loglike])
model_approx = graph.EPMeanField.from_approx_dists(
model,
{
**{v: messages.NormalMessage(0.0, 10.0)
for v in model.variables},
**{
v: messages.FixedMessage(sample)
for v, sample in samples.items()
},
},
)
# Mean field approximation
model_approx = graph.EPMeanField.from_approx_dists(
model,
{
**{v: messages.NormalMessage(0.0, 10.0)
for v in model.variables},
**{
v: messages.FixedMessage(sample)
for v, sample in samples.items()
},
},
)
laplace = graph.LaplaceOptimiser()
ep_opt = graph.EPOptimiser(model, default_optimiser=laplace)
new_approx = ep_opt.run(model_approx, max_steps=20)
new_approx.mean_field.subset(hierarchical_params)
m = np.mean([np.mean(sample) for sample in data])
logt = np.mean([np.log(np.std(sample)**-2) for sample in data])
assert new_approx.mean_field[mu_x_].mean == pytest.approx(m, rel=1.0)
assert new_approx.mean_field[mu_logt_].mean == pytest.approx(logt, rel=1.0)
