Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

update posterior predictive sampling and notebook #5268

Merged
merged 9 commits into from
Dec 20, 2021
Merged

update posterior predictive sampling and notebook #5268

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
48fa066
some updates to posterior predictive sampling and conversion
OriolAbril Dec 16, 2021
cd46ada
update posterior predictive notebook
OriolAbril Dec 16, 2021
271633e
improve idata-sample_posterior_predictive integration
OriolAbril Dec 18, 2021
ac9dcaa
try fixing tests
OriolAbril Dec 18, 2021
4d7367a
fix tests
OriolAbril Dec 18, 2021
2097e74
remove plot_posterior_predictive_glm
OriolAbril Dec 18, 2021
ae13f91
add report fallback again
OriolAbril Dec 18, 2021
8c1930c
Update pymc/sampling.py
OriolAbril Dec 19, 2021
676808f
add example and one extra test with thinning
OriolAbril Dec 19, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4,591 changes: 4,378 additions & 213 deletions docs/source/learn/examples/posterior_predictive.ipynb
View file
Open in desktop

Large diffs are not rendered by default.

48 changes: 18 additions & 30 deletions pymc/backends/arviz.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,12 @@ def __init__(
)
self.ntune = len(self.trace) - self.ndraws
self.posterior_trace, self.warmup_trace = self.split_trace()
elif posterior_predictive is not None:
aelem = next(iter(posterior_predictive.values()))
self.nchains, self.ndraws = aelem.shape[:2]
elif predictions is not None:
aelem = next(iter(predictions.values()))
self.nchains, self.ndraws = aelem.shape[:2]
else:
self.nchains = self.ndraws = 0

Expand All @@ -175,29 +181,11 @@ def __init__(
self.log_likelihood = log_likelihood
self.predictions = predictions

def arbitrary_element(dct: Dict[Any, np.ndarray]) -> np.ndarray:
return next(iter(dct.values()))

if trace is None:
# if you have a posterior_predictive built with keep_dims,
# you'll lose here, but there's nothing I can do about that.
self.nchains = 1
get_from = None
if predictions is not None:
get_from = predictions
elif posterior_predictive is not None:
get_from = posterior_predictive
elif prior is not None:
get_from = prior
if get_from is None:
# pylint: disable=line-too-long
raise ValueError(
"When constructing InferenceData must have at least"
" one of trace, prior, posterior_predictive or predictions."
)

aelem = arbitrary_element(get_from)
self.ndraws = aelem.shape[0]
if all(elem is None for elem in (trace, predictions, posterior_predictive, prior)):
raise ValueError(
"When constructing InferenceData must have at least"
" one of trace, prior, posterior_predictive or predictions."
)

self.coords = {**self.model.coords, **(coords or {})}
self.coords = {
Expand Down Expand Up @@ -400,14 +388,10 @@ def translate_posterior_predictive_dict_to_xarray(self, dct) -> xr.Dataset:
"""Take Dict of variables to numpy ndarrays (samples) and translate into dataset."""
data = {}
for k, ary in dct.items():
shape = ary.shape
if shape[0] == self.nchains and shape[1] == self.ndraws:
if (ary.shape[0] == self.nchains) and (ary.shape[1] == self.ndraws):
data[k] = ary
elif shape[0] == self.nchains * self.ndraws:
data[k] = ary.reshape((self.nchains, self.ndraws, *shape[1:]))
else:
data[k] = np.expand_dims(ary, 0)
# pylint: disable=line-too-long
_log.warning(
"posterior predictive variable %s's shape not compatible with number of chains and draws. "
"This can mean that some draws or even whole chains are not represented.",
Expand Down Expand Up @@ -648,14 +632,18 @@ def predictions_to_inference_data(
raise ValueError(
"Do not pass True for inplace unless passing" "an existing InferenceData as idata_orig"
)
new_idata = InferenceDataConverter(
converter = InferenceDataConverter(
trace=posterior_trace,
predictions=predictions,
model=model,
coords=coords,
dims=dims,
log_likelihood=False,
).to_inference_data()
)
if hasattr(idata_orig, "posterior"):
converter.nchains = idata_orig.posterior.dims["chain"]
converter.ndraws = idata_orig.posterior.dims["draw"]
new_idata = converter.to_inference_data()
if idata_orig is None:
return new_idata
elif inplace:
Expand Down
7 changes: 3 additions & 4 deletions pymc/sampling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1541,7 +1541,7 @@ def sample_posterior_predictive(
model: Optional[Model] = None,
var_names: Optional[List[str]] = None,
size: Optional[int] = None,
keep_size: Optional[bool] = False,
keep_size: Optional[bool] = True,
random_seed=None,
progressbar: bool = True,
mode: Optional[Union[str, Mode]] = None,
Expand Down Expand Up @@ -1571,7 +1571,7 @@ def sample_posterior_predictive(
The number of random draws from the distribution specified by the parameters in each
sample of the trace. Not recommended unless more than ndraws times nchains posterior
predictive samples are needed.
keep_size : bool, optional
keep_size : bool, default True
Force posterior predictive sample to have the same shape as posterior and sample stats
data: ``(nchains, ndraws, ...)``. Overrides samples and size parameters.
random_seed : int
Expand All @@ -1582,9 +1582,8 @@ def sample_posterior_predictive(
time until completion ("expected time of arrival"; ETA).
mode:
The mode used by ``aesara.function`` to compile the graph.
return_inferencedata : bool
return_inferencedata : bool, default True
Whether to return an :class:`arviz:arviz.InferenceData` (True) object or a dictionary (False).
Defaults to True.
idata_kwargs : dict, optional
Keyword arguments for :func:`pymc.to_inference_data`

Expand Down
62 changes: 19 additions & 43 deletions pymc/tests/test_idata_conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,16 +72,13 @@ def get_inference_data(self, data, eight_schools_params):
data.obj, return_inferencedata=False
)

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,
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,
)

def get_predictions_inference_data(
Expand Down Expand Up @@ -124,7 +121,7 @@ def make_predictions_inference_data(
return idata, posterior_predictive

def test_to_idata(self, data, eight_schools_params, chains, draws):
inference_data, posterior_predictive = self.get_inference_data(data, eight_schools_params)
inference_data = self.get_inference_data(data, eight_schools_params)
test_dict = {
"posterior": ["mu", "tau", "eta", "theta"],
"sample_stats": ["diverging", "lp", "~log_likelihood"],
Expand All @@ -136,13 +133,6 @@ def test_to_idata(self, data, eight_schools_params, chains, draws):
}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
for key, values in posterior_predictive.items():
ivalues = inference_data.posterior_predictive[key]
for chain in range(chains):
assert np.all(
np.isclose(ivalues[chain], values[chain * draws : (chain + 1) * draws])
)

chains = inference_data.posterior.dims["chain"]
draws = inference_data.posterior.dims["draw"]
obs = inference_data.observed_data["obs"]
Expand All @@ -160,26 +150,24 @@ def test_predictions_to_idata(self, data, eight_schools_params):
}

# check adding non-destructively
inference_data, posterior_predictive = self.get_predictions_inference_data(
data, eight_schools_params, False
)
inference_data, _ = self.get_predictions_inference_data(data, eight_schools_params, False)
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
for key, values in posterior_predictive.items():
ivalues = inference_data.predictions[key]
assert ivalues.shape[0] == 1 # one chain in predictions
assert np.all(np.isclose(ivalues[0], values))
for key, ivalues in inference_data.predictions.items():
assert (
len(ivalues["chain"]) == inference_data.posterior.dims["chain"]
) # same chains as in posterior

# check adding in place
inference_data, posterior_predictive = self.get_predictions_inference_data(
data, eight_schools_params, True
)
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
for key, values in posterior_predictive.items():
ivalues = inference_data.predictions[key]
assert ivalues.shape[0] == 1 # one chain in predictions
assert np.all(np.isclose(ivalues[0], values))
for key, ivalues in inference_data.predictions.items():
assert (
len(ivalues["chain"]) == inference_data.posterior.dims["chain"]
) # same chains as in posterior

def test_predictions_to_idata_new(self, data, eight_schools_params):
# check creating new
Expand All @@ -195,19 +183,7 @@ def test_predictions_to_idata_new(self, data, eight_schools_params):
assert not fails
for key, values in posterior_predictive.items():
ivalues = inference_data.predictions[key]
# could the following better be done by simply flattening both the ivalues
# and the values?
if len(ivalues.shape) == 3:
ivalues_arr = np.reshape(
ivalues.values, (ivalues.shape[0] * ivalues.shape[1], ivalues.shape[2])
)
elif len(ivalues.shape) == 2:
ivalues_arr = np.reshape(ivalues.values, (ivalues.shape[0] * ivalues.shape[1]))
else:
raise ValueError(f"Unexpected values shape for variable {key}")
assert (ivalues.shape[0] == 2) and (ivalues.shape[1] == 500)
assert values.shape[0] == 1000
assert np.all(np.isclose(ivalues_arr, values))
assert (len(ivalues["chain"]) == 2) and (len(ivalues["draw"]) == 500)

def test_posterior_predictive_keep_size(self, data, chains, draws, eight_schools_params):
with data.model:
Expand All @@ -229,7 +205,7 @@ def test_posterior_predictive_keep_size(self, data, chains, draws, eight_schools
def test_posterior_predictive_warning(self, data, eight_schools_params, caplog):
with data.model:
posterior_predictive = pm.sample_posterior_predictive(
data.obj, 370, return_inferencedata=False
data.obj, 370, return_inferencedata=False, keep_size=False
)
inference_data = to_inference_data(
trace=data.obj,
Expand Down