def test_ovewrite_model_coords_dims(self):
"""Check coords and dims from model object can be partially overwrited."""
dim1 = ["a", "b"]
new_dim1 = ["c", "d"]
coords = {"dim1": dim1, "dim2": ["c1", "c2"]}
x_data = np.arange(4).reshape((2, 2))
y = x_data + np.random.normal(size=(2, 2))
with pm.Model(coords=coords):
x = pm.Data("x", x_data, dims=("dim1", "dim2"))
beta = pm.Normal("beta", 0, 1, dims="dim1")
_ = pm.Normal("obs", x * beta, 1, observed=y, dims=("dim1", "dim2"))
trace = pm.sample(100, tune=100, return_inferencedata=False)
idata1 = to_inference_data(trace)
idata2 = to_inference_data(trace, coords={"dim1": new_dim1}, dims={"beta": ["dim2"]})
test_dict = {"posterior": ["beta"], "observed_data": ["obs"], "constant_data": ["x"]}
fails1 = check_multiple_attrs(test_dict, idata1)
assert not fails1
fails2 = check_multiple_attrs(test_dict, idata2)
assert not fails2
assert "dim1" in list(idata1.posterior.beta.dims)
assert "dim2" in list(idata2.posterior.beta.dims)
assert np.all(idata1.constant_data.x.dim1.values == np.array(dim1))
assert np.all(idata1.constant_data.x.dim2.values == np.array(["c1", "c2"]))
assert np.all(idata2.constant_data.x.dim1.values == np.array(new_dim1))
assert np.all(idata2.constant_data.x.dim2.values == np.array(["c1", "c2"]))
def test_no_trace(self):
with pm.Model() as model:
x = pm.Data("x", [1.0, 2.0, 3.0])
y = pm.Data("y", [1.0, 2.0, 3.0])
beta = pm.Normal("beta", 0, 1)
obs = pm.Normal("obs", x * beta, 1, observed=y) # pylint: disable=unused-variable
idata = pm.sample(100, tune=100)
prior = pm.sample_prior_predictive()
posterior_predictive = pm.sample_posterior_predictive(idata)
# Only prior
inference_data = to_inference_data(prior=prior, model=model)
test_dict = {"prior": ["beta"], "prior_predictive": ["obs"]}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
# Only posterior_predictive
inference_data = to_inference_data(posterior_predictive=posterior_predictive, model=model)
test_dict = {"posterior_predictive": ["obs"]}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
# Prior and posterior_predictive but no trace
inference_data = to_inference_data(
prior=prior, posterior_predictive=posterior_predictive, model=model
)
test_dict = {
"prior": ["beta"],
"prior_predictive": ["obs"],
"posterior_predictive": ["obs"],
}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
def test_predictions_constant_data(self):
with pm.Model():
x = pm.Data("x", [1.0, 2.0, 3.0])
y = pm.Data("y", [1.0, 2.0, 3.0])
beta = pm.Normal("beta", 0, 1)
obs = pm.Normal("obs", x * beta, 1, observed=y) # pylint: disable=unused-variable
trace = pm.sample(100, tune=100, return_inferencedata=False)
inference_data = to_inference_data(trace)
test_dict = {"posterior": ["beta"], "observed_data": ["obs"], "constant_data": ["x"]}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
with pm.Model():
x = pm.Data("x", [1.0, 2.0])
y = pm.Data("y", [1.0, 2.0])
beta = pm.Normal("beta", 0, 1)
obs = pm.Normal("obs", x * beta, 1, observed=y) # pylint: disable=unused-variable
predictive_trace = pm.sample_posterior_predictive(inference_data)
assert set(predictive_trace.keys()) == {"obs"}
# this should be four chains of 100 samples
# assert predictive_trace["obs"].shape == (400, 2)
# but the shape seems to vary between pymc3 versions
inference_data = predictions_to_inference_data(predictive_trace, posterior_trace=trace)
test_dict = {"posterior": ["beta"], "~observed_data": ""}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails, "Posterior data not copied over as expected."
test_dict = {"predictions": ["obs"]}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails, "Predictions not instantiated as expected."
test_dict = {"predictions_constant_data": ["x"]}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails, "Predictions constant data not instantiated as expected."
def test_constant_data(self, use_context):
"""Test constant_data group behaviour."""
with pm.Model() as model:
x = pm.Data("x", [1.0, 2.0, 3.0])
y = pm.Data("y", [1.0, 2.0, 3.0])
beta = pm.Normal("beta", 0, 1)
obs = pm.Normal("obs", x * beta, 1, observed=y) # pylint: disable=unused-variable
trace = pm.sample(100, tune=100, return_inferencedata=False)
if use_context:
inference_data = to_inference_data(trace=trace)
if not use_context:
inference_data = to_inference_data(trace=trace, model=model)
test_dict = {"posterior": ["beta"], "observed_data": ["obs"], "constant_data": ["x"]}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
def test_neff(self):
if hasattr(self, "min_n_eff"):
with self.model:
idata = to_inference_data(self.trace[self.burn:])
n_eff = az.ess(idata)
for var in n_eff:
npt.assert_array_less(self.min_n_eff, n_eff[var])
def test_save_warmup_issue_1208_after_3_9(self):
with pm.Model():
pm.Uniform("u1")
pm.Normal("n1")
trace = pm.sample(
tune=100,
draws=200,
chains=2,
cores=1,
step=pm.Metropolis(),
discard_tuned_samples=False,
return_inferencedata=False,
)
assert isinstance(trace, pm.backends.base.MultiTrace)
assert len(trace) == 300
# from original trace, warmup draws should be separated out
idata = to_inference_data(trace, save_warmup=True)
test_dict = {
"posterior": ["u1", "n1"],
"sample_stats": ["~tune", "accept"],
"warmup_posterior": ["u1", "n1"],
"warmup_sample_stats": ["~tune", "accept"],
}
fails = check_multiple_attrs(test_dict, idata)
assert not fails
assert idata.posterior.dims["chain"] == 2
assert idata.posterior.dims["draw"] == 200
# manually sliced trace triggers the same warning as <=3.8
with pytest.warns(UserWarning, match="Warmup samples"):
idata = to_inference_data(trace[-30:], save_warmup=True)
test_dict = {
"posterior": ["u1", "n1"],
"sample_stats": ["~tune", "accept"],
"~warmup_posterior": [],
"~warmup_sample_stats": [],
}
fails = check_multiple_attrs(test_dict, idata)
assert not fails
assert idata.posterior.dims["chain"] == 2
assert idata.posterior.dims["draw"] == 30
def test_autodetect_coords_from_model(self, use_context):
df_data = pd.DataFrame(columns=["date"]).set_index("date")
dates = pd.date_range(start="2020-05-01", end="2020-05-20")
for city, mu in {"Berlin": 15, "San Marino": 18, "Paris": 16}.items():
df_data[city] = np.random.normal(loc=mu, size=len(dates))
df_data.index = dates
df_data.index.name = "date"
coords = {"date": df_data.index, "city": df_data.columns}
with pm.Model(coords=coords) as model:
europe_mean = pm.Normal("europe_mean_temp", mu=15.0, sd=3.0)
city_offset = pm.Normal("city_offset", mu=0.0, sd=3.0, dims="city")
city_temperature = pm.Deterministic(
"city_temperature", europe_mean + city_offset, dims="city"
)
data_dims = ("date", "city")
data = pm.Data("data", df_data, dims=data_dims)
_ = pm.Normal("likelihood", mu=city_temperature, sd=0.5, observed=data, dims=data_dims)
trace = pm.sample(
return_inferencedata=False,
compute_convergence_checks=False,
cores=1,
chains=1,
tune=20,
draws=30,
step=pm.Metropolis(),
)
if use_context:
idata = to_inference_data(trace=trace)
if not use_context:
idata = to_inference_data(trace=trace, model=model)
assert "city" in list(idata.posterior.dims)
assert "city" in list(idata.observed_data.dims)
assert "date" in list(idata.observed_data.dims)
np.testing.assert_array_equal(idata.posterior.coords["city"], coords["city"])
np.testing.assert_array_equal(idata.observed_data.coords["date"], coords["date"])
np.testing.assert_array_equal(idata.observed_data.coords["city"], coords["city"])
def test_priors_separation(self, use_context):
"""Test model is enough to get prior, prior predictive and observed_data."""
with pm.Model() as model:
x = pm.Data("x", [1.0, 2.0, 3.0])
y = pm.Data("y", [1.0, 2.0, 3.0])
beta = pm.Normal("beta", 0, 1)
obs = pm.Normal("obs", x * beta, 1, observed=y) # pylint: disable=unused-variable
prior = pm.sample_prior_predictive()
test_dict = {
"prior": ["beta", "~obs"],
"observed_data": ["obs"],
"prior_predictive": ["obs"],
}
if use_context:
with model:
inference_data = to_inference_data(prior=prior)
else:
inference_data = to_inference_data(prior=prior, model=model)
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
def test_posterior_predictive_keep_size(self, data, chains, draws, eight_schools_params):
with data.model:
posterior_predictive = pm.sample_posterior_predictive(data.obj, keep_size=True)
inference_data = to_inference_data(
trace=data.obj,
posterior_predictive=posterior_predictive,
coords={"school": np.arange(eight_schools_params["J"])},
dims={"theta": ["school"], "eta": ["school"]},
)
shape = inference_data.posterior_predictive.obs.shape
assert np.all(
[obs_s == s for obs_s, s in zip(shape, (chains, draws, eight_schools_params["J"]))]
)
def test_posterior_predictive_warning(self, data, eight_schools_params, caplog):
with data.model:
posterior_predictive = pm.sample_posterior_predictive(data.obj, 370)
inference_data = to_inference_data(
trace=data.obj,
posterior_predictive=posterior_predictive,
coords={"school": np.arange(eight_schools_params["J"])},
dims={"theta": ["school"], "eta": ["school"]},
)
records = caplog.records
shape = inference_data.posterior_predictive.obs.shape
assert np.all([obs_s == s for obs_s, s in zip(shape, (1, 370, eight_schools_params["J"]))])
assert len(records) == 1
assert records[0].levelname == "WARNING"
def get_inference_data(self, data, eight_schools_params):
with data.model:
prior = pm.sample_prior_predictive()
posterior_predictive = pm.sample_posterior_predictive(data.obj)
return (
to_inference_data(
trace=data.obj,
prior=prior,
posterior_predictive=posterior_predictive,
coords={"school": np.arange(eight_schools_params["J"])},
dims={"theta": ["school"], "eta": ["school"]},
model=data.model,
),
posterior_predictive,
)
def get_predictions_inference_data(
self, data, eight_schools_params, inplace
) -> Tuple[InferenceData, Dict[str, np.ndarray]]:
with data.model:
prior = pm.sample_prior_predictive()
posterior_predictive = pm.sample_posterior_predictive(data.obj)
idata = to_inference_data(
trace=data.obj,
prior=prior,
coords={"school": np.arange(eight_schools_params["J"])},
dims={"theta": ["school"], "eta": ["school"]},
)
assert isinstance(idata, InferenceData)
extended = predictions_to_inference_data(
posterior_predictive, idata_orig=idata, inplace=inplace
)
assert isinstance(extended, InferenceData)
assert (id(idata) == id(extended)) == inplace
return (extended, posterior_predictive)
def test_Rhat(self):
with self.model:
idata = to_inference_data(self.trace[self.burn:])
rhat = az.rhat(idata)
for var in rhat:
npt.assert_allclose(rhat[var], 1, rtol=0.01)