def test_bivariate_shapes(bivariate_gaussian):
advi = pm.fit(bivariate_gaussian(), num_steps=5000)
assert advi.losses.numpy().shape == (5000, 2)
samples = advi.approximation.sample(5000)
assert samples.posterior["bivariate_gaussian/density"].values.shape == (
1, 5000, 2)
def test_fit(approximation, conjugate_normal_model):
model = conjugate_normal_model["model"]()
approx = _test_kwargs[approximation]
advi = pm.fit(method=approx["method"](model), **approx["fit_kwargs"])
assert advi is not None
assert advi.losses.numpy().shape == (approx["fit_kwargs"].get(
"num_steps", 10000), )
q_samples = advi.approximation.sample(
**approx.get("sample_kwargs", {"n": 1000}))
# Calculating mean from all draws and comparing to the actual one
calculated_mean = q_samples.posterior["model/mu"].mean(dim=("chain",
"draw"))
np.testing.assert_allclose(calculated_mean,
conjugate_normal_model["estimated_mean"],
rtol=0.05)
if "sample_kwargs" in approx and approx["sample_kwargs"].get(
"include_log_likelihood"):
sample_mean = q_samples.posterior["model/mu"].sel(
chain=0, draw=0) # Single draw
ll_from_scipy = norm.logpdf(conjugate_normal_model["data"],
sample_mean,
conjugate_normal_model["known_sigma"])
ll_from_pymc4 = q_samples.log_likelihood["model/ll"].sel(chain=0,
draw=0)
assert ll_from_scipy.shape == ll_from_pymc4.shape
np.testing.assert_allclose(ll_from_scipy, ll_from_pymc4, rtol=1e-4)
def test_advi_with_deterministics(simple_model_with_deterministic):
advi = pm.fit(simple_model_with_deterministic(), num_steps=1000)
samples = advi.approximation.sample(100)
norm = "simple_model_with_deterministic/simple_model/norm"
determ = "simple_model_with_deterministic/determ"
np.testing.assert_allclose(samples.posterior[determ],
samples.posterior[norm] * 2)
def test_advi_with_deterministics_in_nested_models(
deterministics_in_nested_models):
(
model,
*_,
deterministic_mapping,
) = deterministics_in_nested_models
advi = pm.fit(model(), num_steps=1000)
samples = advi.approximation.sample(100)
for deterministic, (inputs, op) in deterministic_mapping.items():
np.testing.assert_allclose(samples.posterior[deterministic],
op(*[samples.posterior[i] for i in inputs]),
rtol=1e-6)
def test_fit(approximation, simple_model):
model = simple_model["model"]()
approx = _test_kwargs[approximation]
advi = pm.fit(method=approx["method"](model), **approx["fit_kwargs"])
assert advi is not None
assert advi.losses.numpy().shape == (approx["fit_kwargs"].get("num_steps")
or 10000, )
q_samples = advi.approximation.sample(10000)
estimated_mean = simple_model["estimated_mean"]
np.testing.assert_allclose(
np.mean(np.squeeze(q_samples.posterior["model/mu"].values, axis=0)),
estimated_mean,
rtol=0.05,
)
