Unverified Commit b5891be9 authored by Luciano Paz's avatar Luciano Paz Committed by GitHub

Merge pull request #3597 from rpgoldman/vectorize-posterior-predictive

Second try to vectorize sample_posterior_predictive.
parents b49340b8 e2c334ee
......@@ -36,3 +36,7 @@ benchmarks/results/
# VSCode
.vscode/
.mypy_cache
pytestdebug.log
.dir-locals.el
.pycheckers
......@@ -7,6 +7,7 @@
- `DEMetropolis` can now tune both `lambda` and `scaling` parameters, but by default neither of them are tuned. See [#3743](https://github.com/pymc-devs/pymc3/pull/3743) for more info.
- `DEMetropolisZ`, an improved variant of `DEMetropolis` brings better parallelization and higher efficiency with fewer chains with a slower initial convergence. This implementation is experimental. See [#3784](https://github.com/pymc-devs/pymc3/pull/3784) for more info.
- Notebooks that give insight into `DEMetropolis`, `DEMetropolisZ` and the `DifferentialEquation` interface are now located in the [Tutorials/Deep Dive](https://docs.pymc.io/nb_tutorials/index.html) section.
- Add `fast_sample_posterior_predictive`, a vectorized alternative to `sample_posterior_predictive`. This alternative is substantially faster for large models.
### Maintenance
- Remove `sample_ppc` and `sample_ppc_w` that were deprecated in 3.6.
......
......@@ -44,7 +44,7 @@ from .exceptions import *
from . import sampling
from .backends.tracetab import *
from .backends import save_trace, load_trace
from .backends import save_trace, load_trace, point_list_to_multitrace
from .plots import *
from .tests import test
......
......@@ -131,7 +131,7 @@ defined that returns a MultiTrace object.
For specific examples, see pymc3.backends.{ndarray,text,sqlite}.py.
"""
from ..backends.ndarray import NDArray, save_trace, load_trace
from ..backends.ndarray import NDArray, save_trace, load_trace, point_list_to_multitrace
from ..backends.text import Text
from ..backends.sqlite import SQLite
from ..backends.hdf5 import HDF5
......
......@@ -19,13 +19,14 @@ creating custom backends).
"""
import itertools as itl
import logging
from typing import List
from typing import Dict, List, Optional
from abc import ABC
import numpy as np
import warnings
import theano.tensor as tt
from ..model import modelcontext
from ..model import modelcontext, Model
from .report import SamplerReport, merge_reports
logger = logging.getLogger('pymc3')
......@@ -35,7 +36,7 @@ class BackendError(Exception):
pass
class BaseTrace:
class BaseTrace(ABC):
"""Base trace object
Parameters
......
......@@ -20,12 +20,12 @@ import glob
import json
import os
import shutil
from typing import Optional, Dict, Any
from typing import Optional, Dict, Any, List
import numpy as np
from pymc3.backends import base
from pymc3.backends.base import MultiTrace
from pymc3.model import Model
from pymc3.model import Model, modelcontext
from pymc3.exceptions import TraceDirectoryError
......@@ -366,3 +366,19 @@ def _slice_as_ndarray(strace, idx):
sliced.draw_idx = (stop - start) // step
return sliced
def point_list_to_multitrace(point_list: List[Dict[str, np.ndarray]], model: Optional[Model]=None) -> MultiTrace:
'''transform point list into MultiTrace'''
_model = modelcontext(model)
varnames = list(point_list[0].keys())
with _model:
chain = NDArray(model=_model, vars=[_model[vn] for vn in varnames])
chain.setup(draws=len(point_list), chain=0)
# since we are simply loading a trace by hand, we need only a vacuous function for
# chain.record() to use. This crushes the default.
def point_fun(point):
return [point[vn] for vn in varnames]
chain.fn = point_fun
for point in point_list:
chain.record(point)
return MultiTrace([chain])
......@@ -16,6 +16,8 @@ from . import timeseries
from . import transforms
from . import shape_utils
from .posterior_predictive import fast_sample_posterior_predictive
from .continuous import Uniform
from .continuous import Flat
from .continuous import HalfFlat
......@@ -168,5 +170,6 @@ __all__ = ['Uniform',
'Interpolated',
'Bound',
'Rice',
'Simulator'
'Simulator',
'fast_sample_posterior_predictive'
]
......@@ -13,6 +13,10 @@
# limitations under the License.
import numbers
import contextvars
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Optional, Callable
import numpy as np
import theano.tensor as tt
......@@ -33,6 +37,7 @@ from .shape_utils import (
__all__ = ['DensityDist', 'Distribution', 'Continuous', 'Discrete',
'NoDistribution', 'TensorType', 'draw_values', 'generate_samples']
vectorized_ppc = contextvars.ContextVar('vectorized_ppc', default=None) # type: contextvars.ContextVar[Optional[Callable]]
class _Unpickling:
pass
......@@ -530,11 +535,18 @@ def draw_values(params, point=None, size=None):
a) are named parameters in the point
b) are RVs with a random method
"""
# The following check intercepts and redirects calls to
# draw_values in the context of sample_posterior_predictive
ppc_sampler = vectorized_ppc.get(None)
if ppc_sampler is not None:
# this is being done inside new, vectorized sample_posterior_predictive
return ppc_sampler(params, trace=point, samples=size)
if point is None:
point = {}
# Get fast drawable values (i.e. things in point or numbers, arrays,
# constants or shares, or things that were already drawn in related
# contexts)
if point is None:
point = {}
with _DrawValuesContext() as context:
params = dict(enumerate(params))
drawn = context.drawn_vars
......
This diff is collapsed.
......@@ -1420,6 +1420,11 @@ def _get_scaling(total_size, shape, ndim):
class FreeRV(Factor, PyMC3Variable):
"""Unobserved random variable that a model is specified in terms of."""
dshape = None # type: Tuple[int, ...]
size = None # type: int
distribution = None # type: Optional[Distribution]
model = None # type: Optional[Model]
def __init__(self, type=None, owner=None, index=None, name=None,
distribution=None, total_size=None, model=None):
"""
......
......@@ -33,6 +33,7 @@ from theano.tensor import Tensor
from .backends.base import BaseTrace, MultiTrace
from .backends.ndarray import NDArray
from .distributions.distribution import draw_values
from .distributions.posterior_predictive import fast_sample_posterior_predictive
from .model import modelcontext, Point, all_continuous, Model
from .step_methods import (
NUTS,
......@@ -71,6 +72,7 @@ __all__ = [
"sample_posterior_predictive_w",
"init_nuts",
"sample_prior_predictive",
"fast_sample_posterior_predictive",
]
STEP_METHODS = (
......@@ -1552,8 +1554,9 @@ def sample_posterior_predictive(
raise IncorrectArgumentsError("Should not specify both vars and var_names arguments.")
else:
vars = [model[x] for x in var_names]
elif vars is not None: # var_names is None, and vars is not.
warnings.warn("vars argument is deprecated in favor of var_names.", DeprecationWarning)
elif vars is not None: # var_names is None, and vars is not.
warnings.warn("vars argument is deprecated in favor of var_names.",
DeprecationWarning)
if vars is None:
vars = model.observed_RVs
......
......@@ -40,19 +40,30 @@ class TestData(SeededTest):
trace = pm.sample(1000, init=None, tune=1000, chains=1)
pp_trace0 = pm.sample_posterior_predictive(trace, 1000)
pp_trace01 = pm.fast_sample_posterior_predictive(trace, 1000)
x_shared.set_value(x_pred)
pp_trace1 = pm.sample_posterior_predictive(trace, samples=1000)
pp_trace11 = pm.fast_sample_posterior_predictive(trace, samples=1000)
prior_trace1 = pm.sample_prior_predictive(1000)
assert prior_trace0['b'].shape == (1000,)
assert prior_trace0['obs'].shape == (1000, 100)
assert prior_trace1['obs'].shape == (1000, 200)
assert pp_trace0['obs'].shape == (1000, 100)
assert pp_trace01['obs'].shape == (1000, 100)
np.testing.assert_allclose(x, pp_trace0['obs'].mean(axis=0), atol=1e-1)
np.testing.assert_allclose(x, pp_trace01['obs'].mean(axis=0), atol=1e-1)
assert pp_trace1['obs'].shape == (1000, 200)
assert pp_trace11['obs'].shape == (1000, 200)
assert prior_trace1['b'].shape == (1000,)
assert prior_trace1['obs'].shape == (1000, 200)
np.testing.assert_allclose(x_pred, pp_trace1['obs'].mean(axis=0),
atol=1e-1)
np.testing.assert_allclose(x_pred, pp_trace11['obs'].mean(axis=0),
atol=1e-1)
def test_sample_posterior_predictive_after_set_data(self):
with pm.Model() as model:
......@@ -66,10 +77,14 @@ class TestData(SeededTest):
x_test = [5, 6, 9]
pm.set_data(new_data={'x': x_test})
y_test = pm.sample_posterior_predictive(trace)
y_test1 = pm.fast_sample_posterior_predictive(trace)
assert y_test['obs'].shape == (1000, 3)
assert y_test1['obs'].shape == (1000, 3)
np.testing.assert_allclose(x_test, y_test['obs'].mean(axis=0),
atol=1e-1)
np.testing.assert_allclose(x_test, y_test1['obs'].mean(axis=0),
atol=1e-1)
def test_sample_after_set_data(self):
with pm.Model() as model:
......@@ -85,10 +100,14 @@ class TestData(SeededTest):
pm.set_data(new_data={'x': new_x, 'y': new_y})
new_trace = pm.sample(1000, init=None, tune=1000, chains=1)
pp_trace = pm.sample_posterior_predictive(new_trace, 1000)
pp_tracef = pm.fast_sample_posterior_predictive(new_trace, 1000)
assert pp_trace['obs'].shape == (1000, 3)
assert pp_tracef['obs'].shape == (1000, 3)
np.testing.assert_allclose(new_y, pp_trace['obs'].mean(axis=0),
atol=1e-1)
np.testing.assert_allclose(new_y, pp_tracef['obs'].mean(axis=0),
atol=1e-1)
def test_creation_of_data_outside_model_context(self):
with pytest.raises((IndexError, TypeError)) as error:
......
......@@ -948,6 +948,64 @@ def test_mixture_random_shape():
assert ppc['like2'].shape == (200, 20)
assert ppc['like3'].shape == (200, 20)
@pytest.mark.xfail
def test_mixture_random_shape_fast():
# test the shape broadcasting in mixture random
y = np.concatenate([nr.poisson(5, size=10),
nr.poisson(9, size=10)])
with pm.Model() as m:
comp0 = pm.Poisson.dist(mu=np.ones(2))
w0 = pm.Dirichlet('w0', a=np.ones(2))
like0 = pm.Mixture('like0',
w=w0,
comp_dists=comp0,
observed=y)
comp1 = pm.Poisson.dist(mu=np.ones((20, 2)),
shape=(20, 2))
w1 = pm.Dirichlet('w1', a=np.ones(2))
like1 = pm.Mixture('like1',
w=w1,
comp_dists=comp1,
observed=y)
comp2 = pm.Poisson.dist(mu=np.ones(2))
w2 = pm.Dirichlet('w2',
a=np.ones(2),
shape=(20, 2))
like2 = pm.Mixture('like2',
w=w2,
comp_dists=comp2,
observed=y)
comp3 = pm.Poisson.dist(mu=np.ones(2),
shape=(20, 2))
w3 = pm.Dirichlet('w3',
a=np.ones(2),
shape=(20, 2))
like3 = pm.Mixture('like3',
w=w3,
comp_dists=comp3,
observed=y)
rand0, rand1, rand2, rand3 = draw_values([like0, like1, like2, like3],
point=m.test_point,
size=100)
assert rand0.shape == (100, 20)
assert rand1.shape == (100, 20)
assert rand2.shape == (100, 20)
assert rand3.shape == (100, 20)
# I *think* that the mixture means that this is not going to work,
# but I could be wrong. [2019/08/22:rpg]
with m:
ppc = pm.fast_sample_posterior_predictive([m.test_point], samples=200)
assert ppc['like0'].shape == (200, 20)
assert ppc['like1'].shape == (200, 20)
assert ppc['like2'].shape == (200, 20)
assert ppc['like3'].shape == (200, 20)
class TestDensityDist():
@pytest.mark.parametrize("shape", [(), (3,), (3, 2)], ids=str)
......@@ -968,6 +1026,10 @@ class TestDensityDist():
ppc = pm.sample_posterior_predictive(trace, samples=samples, model=model, size=size)
assert ppc['density_dist'].shape == (samples, size) + obs.distribution.shape
# ppc = pm.fast_sample_posterior_predictive(trace, samples=samples, model=model, size=size)
# assert ppc['density_dist'].shape == (samples, size) + obs.distribution.shape
@pytest.mark.parametrize("shape", [(), (3,), (3, 2)], ids=str)
def test_density_dist_with_random_sampleable_failure(self, shape):
with pm.Model() as model:
......@@ -987,6 +1049,10 @@ class TestDensityDist():
with pytest.raises(RuntimeError):
pm.sample_posterior_predictive(trace, samples=samples, model=model, size=100)
with pytest.raises((TypeError, RuntimeError)):
pm.fast_sample_posterior_predictive(trace, samples=samples, model=model, size=100)
@pytest.mark.parametrize("shape", [(), (3,), (3, 2)], ids=str)
def test_density_dist_with_random_sampleable_hidden_error(self, shape):
with pm.Model() as model:
......@@ -1008,6 +1074,11 @@ class TestDensityDist():
assert len(ppc['density_dist']) == samples
assert ((samples,) + obs.distribution.shape) != ppc['density_dist'].shape
ppc = pm.fast_sample_posterior_predictive(trace, samples=samples, model=model)
assert len(ppc['density_dist']) == samples
assert ((samples,) + obs.distribution.shape) != ppc['density_dist'].shape
def test_density_dist_with_random_sampleable_handcrafted_success(self):
with pm.Model() as model:
mu = pm.Normal('mu', 0, 1)
......@@ -1027,6 +1098,28 @@ class TestDensityDist():
ppc = pm.sample_posterior_predictive(trace, samples=samples, model=model, size=size)
assert ppc['density_dist'].shape == (samples, size) + obs.distribution.shape
@pytest.mark.xfail
def test_density_dist_with_random_sampleable_handcrafted_success_fast(self):
with pm.Model() as model:
mu = pm.Normal('mu', 0, 1)
normal_dist = pm.Normal.dist(mu, 1)
rvs = pm.Normal.dist(mu, 1, shape=100).random
obs = pm.DensityDist(
'density_dist',
normal_dist.logp,
observed=np.random.randn(100),
random=rvs,
wrap_random_with_dist_shape=False
)
trace = pm.sample(100)
samples = 500
size = 100
ppc = pm.fast_sample_posterior_predictive(trace, samples=samples, model=model, size=size)
assert ppc['density_dist'].shape == (samples, size) + obs.distribution.shape
def test_density_dist_without_random_not_sampleable(self):
with pm.Model() as model:
mu = pm.Normal('mu', 0, 1)
......@@ -1038,6 +1131,9 @@ class TestDensityDist():
with pytest.raises(ValueError):
pm.sample_posterior_predictive(trace, samples=samples, model=model, size=100)
with pytest.raises((TypeError, ValueError)):
pm.fast_sample_posterior_predictive(trace, samples=samples, model=model, size=100)
class TestNestedRandom(SeededTest):
def build_model(self, distribution, shape, nested_rvs_info):
......
......@@ -15,7 +15,7 @@
from ..model import Model
from ..distributions.continuous import Flat, Normal
from ..distributions.timeseries import EulerMaruyama, AR1, AR, GARCH11
from ..sampling import sample, sample_posterior_predictive
from ..sampling import sample, sample_posterior_predictive, fast_sample_posterior_predictive
from ..theanof import floatX
import numpy as np
......@@ -141,9 +141,12 @@ def test_linear():
trace = sample(init='advi+adapt_diag', chains=1)
ppc = sample_posterior_predictive(trace, model=model)
ppcf = fast_sample_posterior_predictive(trace, model=model)
# test
p95 = [2.5, 97.5]
lo, hi = np.percentile(trace[lamh], p95, axis=0)
assert (lo < lam) and (lam < hi)
lo, hi = np.percentile(ppc['zh'], p95, axis=0)
assert ((lo < z) * (z < hi)).mean() > 0.95
lo, hi = np.percentile(ppcf['zh'], p95, axis=0)
assert ((lo < z) * (z < hi)).mean() > 0.95
......@@ -264,12 +264,17 @@ class TestSaveLoad:
np.random.seed(seed)
with TestSaveLoad.model():
ppc = pm.sample_posterior_predictive(self.trace)
ppcf = pm.fast_sample_posterior_predictive(self.trace)
seed = 10
np.random.seed(seed)
with TestSaveLoad.model():
trace2 = pm.load_trace(directory)
ppc2 = pm.sample_posterior_predictive(trace2)
ppc2f = pm.sample_posterior_predictive(trace2)
for key, value in ppc.items():
assert (value == ppc2[key]).all()
for key, value in ppcf.items():
assert (value == ppc2f[key]).all()
import pymc3 as pm
from pymc3.distributions.posterior_predictive import _TraceDict
import numpy as np
from pymc3.backends.ndarray import point_list_to_multitrace
def test_translate_point_list():
with pm.Model() as model:
mu = pm.Normal("mu", 0.0, 1.0)
a = pm.Normal("a", mu=mu, sigma=1, observed=0.0)
mt = point_list_to_multitrace([model.test_point], model)
assert isinstance(mt, pm.backends.base.MultiTrace)
assert set(["mu"]) == set(mt.varnames)
assert len(mt) == 1
def test_build_TraceDict():
with pm.Model() as model:
mu = pm.Normal("mu", 0.0, 1.0)
a = pm.Normal("a", mu=mu, sigma=1, observed=np.array([0.5, 0.2]))
trace = pm.sample(chains=2, draws=500)
dict = _TraceDict(multi_trace=trace)
assert isinstance(dict, _TraceDict)
assert len(dict) == 1000
np.testing.assert_array_equal(trace['mu'], dict['mu'])
assert set(trace.varnames) == set(dict.varnames) == set(["mu"])
def test_build_TraceDict_point_list():
with pm.Model() as model:
mu = pm.Normal("mu", 0.0, 1.0)
a = pm.Normal("a", mu=mu, sigma=1, observed=np.array([0.5, 0.2]))
dict = _TraceDict(point_list=[model.test_point])
assert set(dict.varnames) == set(["mu"])
assert len(dict) == 1
assert len(dict["mu"]) == 1
assert dict["mu"][0] == 0.0
This diff is collapsed.
......@@ -41,15 +41,19 @@ class TestShared(SeededTest):
trace = pm.sample(1000, init=None, tune=1000, chains=1)
pp_trace0 = pm.sample_posterior_predictive(trace, 1000)
pp_trace01 = pm.fast_sample_posterior_predictive(trace, 1000)
x_shared.set_value(x_pred)
prior_trace1 = pm.sample_prior_predictive(1000)
pp_trace1 = pm.sample_posterior_predictive(trace, 1000)
pp_trace11 = pm.fast_sample_posterior_predictive(trace, 1000)
assert prior_trace0['b'].shape == (1000,)
assert prior_trace0['obs'].shape == (1000, 100)
np.testing.assert_allclose(x, pp_trace0['obs'].mean(axis=0), atol=1e-1)
np.testing.assert_allclose(x, pp_trace01['obs'].mean(axis=0), atol=1e-1)
assert prior_trace1['b'].shape == (1000,)
assert prior_trace1['obs'].shape == (1000, 200)
np.testing.assert_allclose(x_pred, pp_trace1['obs'].mean(axis=0), atol=1e-1)
np.testing.assert_allclose(x_pred, pp_trace11['obs'].mean(axis=0), atol=1e-1)
......@@ -6,3 +6,5 @@ pandas>=0.18.0
patsy>=0.5.1
fastprogress>=0.2.0
h5py>=2.7.0
typing-extensions>=3.7.4
contextvars; python_version < '3.7'
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment