From 1097eee28007ba703438c67e57752b82d8fe9a85 Mon Sep 17 00:00:00 2001 From: Jing Xie Date: Thu, 21 Jan 2021 17:22:54 -0500 Subject: [PATCH] Updates for new PyMC and Theano-PyMC versions --- pymc3_hmm/distributions.py | 51 ++++++++++++++++++------------------- pymc3_hmm/step_methods.py | 29 +++++++++++---------- setup.py | 4 +-- tests/test_distributions.py | 4 +-- tests/test_step_methods.py | 2 +- tests/utils.py | 5 ++-- 6 files changed, 47 insertions(+), 48 deletions(-) diff --git a/pymc3_hmm/distributions.py b/pymc3_hmm/distributions.py index 980af71..2871b7b 100644 --- a/pymc3_hmm/distributions.py +++ b/pymc3_hmm/distributions.py @@ -1,21 +1,17 @@ -import numpy as np +from copy import copy +import numpy as np +import pymc3 as pm import theano import theano.tensor as tt - -import pymc3 as pm - -from copy import copy - -from theano.gof.op import get_test_value - -from pymc3.distributions.mixture import all_discrete, _conversion_map -from pymc3.distributions.distribution import draw_values, _DrawValuesContext +from pymc3.distributions.distribution import _DrawValuesContext, draw_values +from pymc3.distributions.mixture import _conversion_map, all_discrete +from theano.graph.op import get_test_value from pymc3_hmm.utils import ( broadcast_to, - tt_expand_dims, tt_broadcast_arrays, + tt_expand_dims, vsearchsorted, ) @@ -218,21 +214,24 @@ def random(self, point=None, size=None): ------- array """ - with _DrawValuesContext() as draw_context: - - # TODO FIXME: Very, very lame... - term_smpl = draw_context.drawn_vars.get((self.states, 1), None) - if term_smpl is not None: - point[self.states.name] = term_smpl - - # `draw_values` is inconsistent and will not use the `size` - # parameter if the variables aren't random variables. - if hasattr(self.states, "distribution"): - (states,) = draw_values([self.states], point=point, size=size) - else: - states = pm.Constant.dist(self.states).random(point=point, size=size) - - # states = states.T + with _DrawValuesContext(): + + (states,) = draw_values([self.states], point=point, size=size) + + if size: + # `draw_values` will not honor the `size` parameter if its arguments + # don't contain random variables, so, when our `self.states` are + # constants, we have to broadcast `states` so that it matches `size + + # self.shape`. + # Unfortunately, this means that our sampler relies on + # `self.shape`, which is bad for other, arguable more important + # reasons (e.g. when this `Distribution`'s symbolic inputs change + # shape, we now have to manually update `Distribution.shape` so + # that the sampler is consistent) + + states = np.broadcast_to( + states, tuple(np.atleast_1d(size)) + tuple(self.shape) + ) samples = np.empty(states.shape) diff --git a/pymc3_hmm/step_methods.py b/pymc3_hmm/step_methods.py index d705fc0..5af4819 100644 --- a/pymc3_hmm/step_methods.py +++ b/pymc3_hmm/step_methods.py @@ -7,12 +7,12 @@ from pymc3.step_methods.arraystep import ArrayStep, Competence from pymc3.util import get_untransformed_name from theano.compile import optdb -from theano.gof.fg import FunctionGraph -from theano.gof.graph import Variable -from theano.gof.graph import inputs as tt_inputs -from theano.gof.op import get_test_value as test_value -from theano.gof.opt import OpRemove -from theano.gof.optdb import Query +from theano.graph.basic import Variable, graph_inputs +from theano.graph.fg import FunctionGraph +from theano.graph.op import get_test_value as test_value +from theano.graph.opt import OpRemove, pre_greedy_local_optimizer +from theano.graph.optdb import Query +from theano.tensor.elemwise import DimShuffle, Elemwise from theano.tensor.subtensor import AdvancedIncSubtensor1 from theano.tensor.var import TensorConstant @@ -135,7 +135,7 @@ def __init__(self, var, values=None, model=None): self.dependent_rvs = [ v for v in model.basic_RVs - if v is not var and var in tt.gof.graph.inputs([v.logpt]) + if v is not var and var in graph_inputs([v.logpt]) ] # We compile a function--from a Theano graph--that computes the @@ -234,7 +234,7 @@ def __init__(self, model_vars, values=None, model=None, rng=None): v for v in model.vars + model.observed_RVs if isinstance(v.distribution, DiscreteMarkovChain) - and all(d in tt_inputs([v.distribution.Gammas]) for d in dir_priors) + and all(d in graph_inputs([v.distribution.Gammas]) for d in dir_priors) ] if not self.dir_priors_untrans or not len(state_seqs) == 1: @@ -300,18 +300,19 @@ def _set_row_mappings(self, Gamma, dir_priors, model): """ # Remove unimportant `Op`s from the transition matrix graph - Gamma = tt.gof.opt.pre_greedy_local_optimizer( + Gamma = pre_greedy_local_optimizer( + FunctionGraph([], []), [ - OpRemove(tt.elemwise.Elemwise(ts.Cast(ts.float32))), - OpRemove(tt.elemwise.Elemwise(ts.Cast(ts.float64))), - OpRemove(tt.elemwise.Elemwise(ts.identity)), + OpRemove(Elemwise(ts.Cast(ts.float32))), + OpRemove(Elemwise(ts.Cast(ts.float64))), + OpRemove(Elemwise(ts.identity)), ], Gamma, ) # Canonicalize the transition matrix graph fg = FunctionGraph( - tt_inputs([Gamma] + self.dir_priors_untrans), + list(graph_inputs([Gamma] + self.dir_priors_untrans)), [Gamma] + self.dir_priors_untrans, clone=True, ) @@ -323,7 +324,7 @@ def _set_row_mappings(self, Gamma, dir_priors, model): Gamma_DimShuffle = Gamma.owner - if not (isinstance(Gamma_DimShuffle.op, tt.elemwise.DimShuffle)): + if not (isinstance(Gamma_DimShuffle.op, DimShuffle)): raise TypeError("The transition matrix should be non-time-varying") Gamma_Join = Gamma_DimShuffle.inputs[0].owner diff --git a/setup.py b/setup.py index 3b69268..3f6b0aa 100644 --- a/setup.py +++ b/setup.py @@ -16,8 +16,8 @@ install_requires=[ "numpy>=1.18.1", "scipy>=1.4.0", - "pymc3>=3.10.0", - "theano-pymc>=1.0.11", + "pymc3==3.11.0", + "theano-pymc==1.1.0", ], tests_require=["pytest"], long_description=open("README.md").read() if exists("README.md") else "", diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 6a4af3b..e17809e 100644 --- a/tests/test_distributions.py +++ b/tests/test_distributions.py @@ -398,8 +398,8 @@ def test_PoissonZeroProcess_random(): # six Poisson means. assert np.array_equal(test_dist.shape, (6, 6)) test_sample = test_dist.random() - assert np.array_equal(test_sample.shape, test_states.squeeze().shape) - assert np.all(test_sample[..., test_states.squeeze() > 0] > 0) + assert np.array_equal(test_sample.shape, test_states.shape) + assert np.all(test_sample[..., test_states > 0] > 0) test_states = np.c_[0, 0, 1, 1, 0, 1] test_dist = PoissonZeroProcess.dist(test_mus, test_states) diff --git a/tests/test_step_methods.py b/tests/test_step_methods.py index 3670115..a3287b9 100644 --- a/tests/test_step_methods.py +++ b/tests/test_step_methods.py @@ -9,7 +9,7 @@ import pymc3 as pm -from theano.gof.op import get_test_value +from theano.graph.op import get_test_value from tests.utils import simulate_poiszero_hmm diff --git a/tests/utils.py b/tests/utils.py index ba1c571..a757438 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -26,8 +26,7 @@ def simulate_poiszero_hmm( sample_point = pm.sample_prior_predictive(samples=1) - # TODO FIXME: Why is `pm.sample_prior_predictive` adding an extra - # dimension to the `Y_rv` result? - sample_point[Y_rv.name] = sample_point[Y_rv.name].squeeze() + # Remote the extra "sampling" dimension from the sample results + sample_point = {k: v.squeeze(0) for k, v in sample_point.items()} return sample_point, test_model