From 31a9f14079634bc74b4bbd1f65e08d68f2eebd3e Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Tue, 21 Sep 2021 22:57:20 +0200 Subject: [PATCH 1/4] Remove deprecated ElemwiseCategorical --- pymc3/step_methods/__init__.py | 1 - pymc3/step_methods/gibbs.py | 105 --------------------------------- 2 files changed, 106 deletions(-) delete mode 100644 pymc3/step_methods/gibbs.py diff --git a/pymc3/step_methods/__init__.py b/pymc3/step_methods/__init__.py index 1cd281752ed..e37a24fe689 100644 --- a/pymc3/step_methods/__init__.py +++ b/pymc3/step_methods/__init__.py @@ -14,7 +14,6 @@ from pymc3.step_methods.compound import CompoundStep from pymc3.step_methods.elliptical_slice import EllipticalSlice -from pymc3.step_methods.gibbs import ElemwiseCategorical from pymc3.step_methods.hmc import NUTS, HamiltonianMC from pymc3.step_methods.metropolis import ( BinaryGibbsMetropolis, diff --git a/pymc3/step_methods/gibbs.py b/pymc3/step_methods/gibbs.py deleted file mode 100644 index 14fb6eaa18c..00000000000 --- a/pymc3/step_methods/gibbs.py +++ /dev/null @@ -1,105 +0,0 @@ -# Copyright 2020 The PyMC Developers -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -Created on May 12, 2012 - -@author: john -""" -from warnings import warn - -import aesara.tensor as at - -from aesara.graph.basic import graph_inputs -from numpy import arange, array, cumsum, empty, exp, max, nested_iters, searchsorted -from numpy.random import uniform - -from pymc3.distributions import logpt -from pymc3.distributions.discrete import Categorical -from pymc3.model import modelcontext -from pymc3.step_methods.arraystep import ArrayStep, Competence - -__all__ = ["ElemwiseCategorical"] - - -class ElemwiseCategorical(ArrayStep): - """ - Gibbs sampling for categorical variables that only have ElemwiseCategoricalise effects - the variable can't be indexed into or transposed or anything otherwise that will mess things up - - """ - - # TODO: It would be great to come up with a way to make - # ElemwiseCategorical more general (handling more complex elementwise - # variables) - - def __init__(self, vars, values=None, model=None): - warn( - "ElemwiseCategorical is deprecated, switch to CategoricalGibbsMetropolis.", - DeprecationWarning, - stacklevel=2, - ) - model = modelcontext(model) - self.var = vars[0] - # XXX: This needs to be refactored - self.sh = None # ones(self.var.dshape, self.var.dtype) - if values is None: - self.values = arange(self.var.distribution.k) - else: - self.values = values - - super().__init__(vars, [elemwise_logp(model, self.var)]) - - def astep(self, q, logp): - p = array([logp(v * self.sh) for v in self.values]) - # XXX: This needs to be refactored - shape = None # self.var.dshape - return categorical(p, shape) - - @staticmethod - def competence(var, has_grad): - dist = getattr(var.owner, "op", None) - if isinstance(dist, Categorical): - return Competence.COMPATIBLE - return Competence.INCOMPATIBLE - - -def elemwise_logp(model, var): - terms = [] - for v in model.basic_RVs: - v_logp = logpt(v) - if var in graph_inputs([v_logp]): - terms.append(v_logp) - return model.fn(at.add(*terms)) - - -def categorical(prob, shape): - out = empty([1] + list(shape)) - - n = len(shape) - it0, it1 = nested_iters( - [prob, out], - [list(range(1, n + 1)), [0]], - op_flags=[["readonly"], ["readwrite"]], - flags=["reduce_ok"], - ) - - for _ in it0: - p, o = it1.itviews - p = cumsum(exp(p - max(p, axis=0))) - r = uniform() * p[-1] - - o[0] = searchsorted(p, r) - - return out[0, ...] From d926746ba972baad14435c7480c1d81765d30db9 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Wed, 22 Sep 2021 02:36:22 +0200 Subject: [PATCH 2/4] Add vars property to CompoundStep --- pymc3/sampling.py | 6 +----- pymc3/step_methods/compound.py | 4 ++++ 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pymc3/sampling.py b/pymc3/sampling.py index e8131ce642f..44097741964 100644 --- a/pymc3/sampling.py +++ b/pymc3/sampling.py @@ -187,11 +187,7 @@ def assign_step_methods(model, step=None, methods=STEP_METHODS, step_kwargs=None except TypeError: steps.append(step) for step in steps: - try: - assigned_vars = assigned_vars.union(set(step.vars)) - except AttributeError: - for method in step.methods: - assigned_vars = assigned_vars.union(set(method.vars)) + assigned_vars = assigned_vars.union(set(step.vars)) # Use competence classmethods to select step methods for remaining # variables diff --git a/pymc3/step_methods/compound.py b/pymc3/step_methods/compound.py index a92569bd30e..9a032609606 100644 --- a/pymc3/step_methods/compound.py +++ b/pymc3/step_methods/compound.py @@ -71,3 +71,7 @@ def reset_tuning(self): for method in self.methods: if hasattr(method, "reset_tuning"): method.reset_tuning() + + @property + def vars(self): + return [var for method in self.methods for var in method.vars] From 5fcc3e5490b51b929ad9123630805bd6069056a7 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Wed, 22 Sep 2021 01:56:23 +0200 Subject: [PATCH 3/4] Convert RVs to value vars in step methods --- pymc3/step_methods/arraystep.py | 6 +-- pymc3/step_methods/elliptical_slice.py | 5 ++- pymc3/step_methods/hmc/base_hmc.py | 2 + pymc3/step_methods/hmc/hmc.py | 2 +- pymc3/step_methods/hmc/nuts.py | 2 +- pymc3/step_methods/metropolis.py | 16 +++++-- pymc3/step_methods/mlda.py | 8 +++- pymc3/step_methods/pgbart.py | 2 +- pymc3/step_methods/sgmcmc.py | 4 +- pymc3/step_methods/slicer.py | 4 +- pymc3/tests/test_step.py | 62 +++++++++++++++++++++++++- 11 files changed, 98 insertions(+), 15 deletions(-) diff --git a/pymc3/step_methods/arraystep.py b/pymc3/step_methods/arraystep.py index 4a367424023..4673b19234e 100644 --- a/pymc3/step_methods/arraystep.py +++ b/pymc3/step_methods/arraystep.py @@ -134,7 +134,7 @@ class ArrayStep(BlockedStep): Parameters ---------- vars: list - List of variables for sampler. + List of value variables for sampler. fs: list of logp Aesara functions allvars: Boolean (default False) blocked: Boolean (default True) @@ -190,7 +190,7 @@ def __init__(self, vars, shared, blocked=True): """ Parameters ---------- - vars: list of sampling variables + vars: list of sampling value variables shared: dict of Aesara variable -> shared variable blocked: Boolean (default True) """ @@ -235,7 +235,7 @@ def __init__(self, vars, shared, blocked=True): """ Parameters ---------- - vars: list of sampling variables + vars: list of sampling value variables shared: dict of Aesara variable -> shared variable blocked: Boolean (default True) """ diff --git a/pymc3/step_methods/elliptical_slice.py b/pymc3/step_methods/elliptical_slice.py index ea88d716598..dfc24631cfc 100644 --- a/pymc3/step_methods/elliptical_slice.py +++ b/pymc3/step_methods/elliptical_slice.py @@ -16,6 +16,7 @@ import numpy as np import numpy.random as nr +from pymc3.aesaraf import inputvars from pymc3.model import modelcontext from pymc3.step_methods.arraystep import ArrayStep, Competence @@ -61,7 +62,7 @@ class EllipticalSlice(ArrayStep): Parameters ---------- vars: list - List of variables for sampler. + List of value variables for sampler. prior_cov: array, optional Covariance matrix of the multivariate Gaussian prior. prior_chol: array, optional @@ -88,6 +89,8 @@ def __init__(self, vars=None, prior_cov=None, prior_chol=None, model=None, **kwa if vars is None: vars = self.model.cont_vars + else: + vars = [self.model.rvs_to_values.get(var, var) for var in vars] vars = inputvars(vars) super().__init__(vars, [self.model.fastlogp], **kwargs) diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py index aaaaa9f4b2c..074d7e3721f 100644 --- a/pymc3/step_methods/hmc/base_hmc.py +++ b/pymc3/step_methods/hmc/base_hmc.py @@ -89,6 +89,8 @@ def __init__( if vars is None: vars = self._model.cont_vars + else: + vars = [self._model.rvs_to_values.get(var, var) for var in vars] super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **aesara_kwargs) diff --git a/pymc3/step_methods/hmc/hmc.py b/pymc3/step_methods/hmc/hmc.py index 1cc8ef335a2..4f3c42c7cab 100644 --- a/pymc3/step_methods/hmc/hmc.py +++ b/pymc3/step_methods/hmc/hmc.py @@ -60,7 +60,7 @@ def __init__(self, vars=None, path_length=2.0, max_steps=1024, **kwargs): Parameters ---------- vars: list, default=None - List of Aesara variables. If None, all continuous RVs from the + List of value variables. If None, all continuous RVs from the model are included. path_length: float, default=2 Total length to travel diff --git a/pymc3/step_methods/hmc/nuts.py b/pymc3/step_methods/hmc/nuts.py index 267c20659fa..210608e2b0f 100644 --- a/pymc3/step_methods/hmc/nuts.py +++ b/pymc3/step_methods/hmc/nuts.py @@ -115,7 +115,7 @@ def __init__(self, vars=None, max_treedepth=10, early_max_treedepth=8, **kwargs) Parameters ---------- vars: list, default=None - List of Aesara variables. If None, all continuous RVs from the + List of value variables. If None, all continuous RVs from the model are included. Emax: float, default 1000 Maximum energy change allowed during leapfrog steps. Larger diff --git a/pymc3/step_methods/metropolis.py b/pymc3/step_methods/metropolis.py index 24b88f7ee89..0d5b9c4b81b 100644 --- a/pymc3/step_methods/metropolis.py +++ b/pymc3/step_methods/metropolis.py @@ -130,7 +130,7 @@ def __init__( Parameters ---------- vars: list - List of variables for sampler + List of value variables for sampler S: standard deviation or covariance matrix Some measure of variance to parameterize proposal distribution proposal_dist: function @@ -153,6 +153,8 @@ def __init__( if vars is None: vars = model.value_vars + else: + vars = [model.rvs_to_values.get(var, var) for var in vars] vars = pm.inputvars(vars) if S is None: @@ -288,7 +290,7 @@ class BinaryMetropolis(ArrayStep): Parameters ---------- vars: list - List of variables for sampler + List of value variables for sampler scaling: scalar or array Initial scale factor for proposal. Defaults to 1. tune: bool @@ -321,6 +323,8 @@ def __init__(self, vars, scaling=1.0, tune=True, tune_interval=100, model=None): self.steps_until_tune = tune_interval self.accepted = 0 + vars = [model.rvs_to_values.get(var, var) for var in vars] + if not all([v.dtype in pm.discrete_types for v in vars]): raise ValueError("All variables must be Bernoulli for BinaryMetropolis") @@ -388,7 +392,7 @@ class BinaryGibbsMetropolis(ArrayStep): Parameters ---------- vars: list - List of variables for sampler + List of value variables for sampler order: list or 'random' List of integers indicating the Gibbs update order e.g., [0, 2, 1, ...]. Default is random @@ -410,6 +414,7 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None): self.transit_p = transit_p initial_point = model.initial_point + vars = [model.rvs_to_values.get(var, var) for var in vars] self.dim = sum(initial_point[v.name].size for v in vars) if order == "random": @@ -490,6 +495,7 @@ def __init__(self, vars, proposal="uniform", order="random", model=None): model = pm.modelcontext(model) + vars = [model.rvs_to_values.get(var, var) for var in vars] vars = pm.inputvars(vars) initial_point = model.initial_point @@ -697,6 +703,8 @@ def __init__( if vars is None: vars = model.cont_vars + else: + vars = [model.rvs_to_values.get(var, var) for var in vars] vars = pm.inputvars(vars) if S is None: @@ -846,6 +854,8 @@ def __init__( if vars is None: vars = model.cont_vars + else: + vars = [model.rvs_to_values.get(var, var) for var in vars] vars = pm.inputvars(vars) if S is None: diff --git a/pymc3/step_methods/mlda.py b/pymc3/step_methods/mlda.py index 0b9dcdc8f28..c2cf6493370 100644 --- a/pymc3/step_methods/mlda.py +++ b/pymc3/step_methods/mlda.py @@ -74,6 +74,8 @@ def __init__(self, *args, **kwargs): value_vars = kwargs.get("vars", None) if value_vars is None: value_vars = model.value_vars + else: + value_vars = [model.rvs_to_values.get(var, var) for var in value_vars] value_vars = pm.inputvars(value_vars) shared = pm.make_shared_replacements(initial_values, value_vars, model) @@ -142,6 +144,8 @@ def __init__(self, *args, **kwargs): value_vars = kwargs.get("vars", None) if value_vars is None: value_vars = model.value_vars + else: + value_vars = [model.rvs_to_values.get(var, var) for var in value_vars] value_vars = pm.inputvars(value_vars) shared = pm.make_shared_replacements(initial_values, value_vars, model) @@ -218,7 +222,7 @@ class MLDA(ArrayStepShared): Note this list excludes the model passed to the model argument above, which is the finest available. vars : list - List of variables for sampler + List of value variables for sampler base_sampler : string Sampler used in the base (coarsest) chain. Can be 'Metropolis' or 'DEMetropolisZ'. Defaults to 'DEMetropolisZ'. @@ -549,6 +553,8 @@ def __init__( # Process model variables if value_vars is None: value_vars = model.value_vars + else: + value_vars = [model.rvs_to_values.get(var, var) for var in value_vars] value_vars = pm.inputvars(value_vars) self.vars = value_vars self.var_names = [var.name for var in self.vars] diff --git a/pymc3/step_methods/pgbart.py b/pymc3/step_methods/pgbart.py index 351f1ae8a26..6c556be95b0 100644 --- a/pymc3/step_methods/pgbart.py +++ b/pymc3/step_methods/pgbart.py @@ -39,7 +39,7 @@ class PGBART(ArrayStepShared): Parameters ---------- vars: list - List of variables for sampler + List of value variables for sampler num_particles : int Number of particles for the conditional SMC sampler. Defaults to 10 max_stages : int diff --git a/pymc3/step_methods/sgmcmc.py b/pymc3/step_methods/sgmcmc.py index 800c2da540c..19308eb6693 100644 --- a/pymc3/step_methods/sgmcmc.py +++ b/pymc3/step_methods/sgmcmc.py @@ -87,7 +87,7 @@ class BaseStochasticGradient(ArrayStepShared): Parameters ---------- vars: list - List of variables for sampler + List of value variables for sampler batch_size`: int Batch Size for each step total_size: int @@ -132,6 +132,8 @@ def __init__( if vars is None: vars = model.value_vars + else: + vars = [model.rvs_to_values.get(var, var) for var in vars] vars = inputvars(vars) diff --git a/pymc3/step_methods/slicer.py b/pymc3/step_methods/slicer.py index 5651d6e78ac..6074c9cc421 100644 --- a/pymc3/step_methods/slicer.py +++ b/pymc3/step_methods/slicer.py @@ -35,7 +35,7 @@ class Slice(ArrayStep): Parameters ---------- vars: list - List of variables for sampler. + List of value variables for sampler. w: float Initial width of slice (Defaults to 1). tune: bool @@ -57,6 +57,8 @@ def __init__(self, vars=None, w=1.0, tune=True, model=None, iter_limit=np.inf, * if vars is None: vars = self.model.cont_vars + else: + vars = [self.model.rvs_to_values.get(var, var) for var in vars] vars = inputvars(vars) super().__init__(vars, [self.model.fastlogp], **kwargs) diff --git a/pymc3/tests/test_step.py b/pymc3/tests/test_step.py index c3893b939fd..1d79491536c 100644 --- a/pymc3/tests/test_step.py +++ b/pymc3/tests/test_step.py @@ -618,8 +618,8 @@ def test_step_categorical(self): check = (("x", np.mean, mu, unc / 10.0), ("x", np.std, unc, unc / 10.0)) with model: steps = ( - CategoricalGibbsMetropolis(model.x, proposal="uniform"), - CategoricalGibbsMetropolis(model.x, proposal="proportional"), + CategoricalGibbsMetropolis([model.x], proposal="uniform"), + CategoricalGibbsMetropolis([model.x], proposal="proportional"), ) for step in steps: idata = sample(8000, tune=0, step=step, start=start, model=model, random_seed=1) @@ -1767,3 +1767,61 @@ def perform(self, node, inputs, outputs): ) assert Q_1_0.mean(axis=1) == 0.0 assert Q_2_1.mean(axis=1) == 0.0 + + +class TestRVsAssignmentSteps: + """ + Test that step methods convert input RVs to respective value vars + Step methods are tested with one and two variables to cover compound + the special branches in `BlockedStep.__new__` + """ + + @pytest.mark.parametrize( + "step, step_kwargs", + [ + (NUTS, {}), + (HamiltonianMC, {}), + (Metropolis, {}), + (Slice, {}), + (EllipticalSlice, {"prior_cov": np.eye(1)}), + (DEMetropolis, {}), + (DEMetropolisZ, {}), + # (MLDA, {}), # TODO + ], + ) + def test_continuous_steps(self, step, step_kwargs): + with Model() as m: + c1 = HalfNormal("c1") + c2 = HalfNormal("c2") + + assert [m.rvs_to_values[c1]] == step([c1], **step_kwargs).vars + assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set( + step([c1, c2], **step_kwargs).vars + ) + + @pytest.mark.parametrize( + "step, step_kwargs", + [ + (BinaryGibbsMetropolis, {}), + (CategoricalGibbsMetropolis, {}), + ], + ) + def test_discrete_steps(self, step, step_kwargs): + with Model() as m: + d1 = Bernoulli("d1", p=0.5) + d2 = Bernoulli("d2", p=0.5) + + assert [m.rvs_to_values[d1]] == step([d1], **step_kwargs).vars + assert {m.rvs_to_values[d1], m.rvs_to_values[d2]} == set( + step([d1, d2], **step_kwargs).vars + ) + + def test_compound_step(self): + with Model() as m: + c1 = HalfNormal("c1") + c2 = HalfNormal("c2") + + step1 = NUTS([c1]) + step2 = NUTS([c2]) + step = CompoundStep([step1, step2]) + assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(step.vars) From 2f902a5ca8b5e79cfb28e3bd12d477a060128362 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Wed, 22 Sep 2021 02:59:42 +0200 Subject: [PATCH 4/4] Do not allow logp_dlogp_function to receive RVs --- pymc3/model.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pymc3/model.py b/pymc3/model.py index aae6e35f9c4..b3d7848fb5c 100644 --- a/pymc3/model.py +++ b/pymc3/model.py @@ -714,14 +714,11 @@ def logp_dlogp_function(self, grad_vars=None, tempered=False, **kwargs): `alpha` can be changed using `ValueGradFunction.set_weights([alpha])`. """ if grad_vars is None: - grad_vars = [v.tag.value_var for v in typefilter(self.free_RVs, continuous_types)] + grad_vars = [self.rvs_to_values[v] for v in typefilter(self.free_RVs, continuous_types)] else: for i, var in enumerate(grad_vars): if var.dtype not in continuous_types: raise ValueError(f"Can only compute the gradient of continuous types: {var}") - # We allow one to pass the random variable terms as arguments - if hasattr(var.tag, "value_var"): - grad_vars[i] = var.tag.value_var if tempered: with self: