diff --git a/docs/source/api/distributions.rst b/docs/source/api/distributions.rst index 26a1f01ff5..60bb36aa36 100644 --- a/docs/source/api/distributions.rst +++ b/docs/source/api/distributions.rst @@ -8,6 +8,7 @@ Distributions distributions/discrete distributions/multivariate distributions/mixture + distributions/simulator distributions/timeseries distributions/transforms distributions/utilities diff --git a/docs/source/api/distributions/simulator.rst b/docs/source/api/distributions/simulator.rst new file mode 100644 index 0000000000..8b9bcd0894 --- /dev/null +++ b/docs/source/api/distributions/simulator.rst @@ -0,0 +1,11 @@ +********** +Simulator +********** + +.. currentmodule:: pymc3.distributions.simulator +.. autosummary:: + + Simulator + +.. automodule:: pymc3.distributions.simulator + :members: diff --git a/pymc3/aesaraf.py b/pymc3/aesaraf.py index 97bdd07577..0cc47b113e 100644 --- a/pymc3/aesaraf.py +++ b/pymc3/aesaraf.py @@ -350,15 +350,24 @@ def rvs_to_value_vars( """ + # Avoid circular dependency + from pymc3.distributions import NoDistribution + def transform_replacements(var, replacements): rv_var, rv_value_var = extract_rv_and_value_vars(var) if rv_value_var is None: - warnings.warn( - f"No value variable found for {rv_var}; " - "the random variable will not be replaced." - ) - return [] + # If RandomVariable does not have a value_var and corresponds to + # a NoDistribution, we allow further replacements in upstream graph + if isinstance(rv_var.owner.op, NoDistribution): + return rv_var.owner.inputs + + else: + warnings.warn( + f"No value variable found for {rv_var}; " + "the random variable will not be replaced." + ) + return [] transform = getattr(rv_value_var.tag, "transform", None) @@ -886,6 +895,22 @@ def compile_rv_inplace(inputs, outputs, mode=None, **kwargs): Using this function ensures that compiled functions containing random variables will produce new samples on each call. """ + + # Avoid circular dependency + from pymc3.distributions import NoDistribution + + # Set the default update of a NoDistribution RNG so that it is automatically + # updated after every function call + output_to_list = outputs if isinstance(outputs, list) else [outputs] + for rv in ( + node + for node in walk_model(output_to_list, walk_past_rvs=True) + if node.owner and isinstance(node.owner.op, NoDistribution) + ): + rng = rv.owner.inputs[0] + if not hasattr(rng, "default_update"): + rng.default_update = rv.owner.outputs[0] + mode = get_mode(mode) opt_qry = mode.provided_optimizer.including("random_make_inplace") mode = Mode(linker=mode.linker, optimizer=opt_qry) diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index e72a90bed1..fda75386a7 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -23,7 +23,6 @@ from typing import Optional import aesara -import aesara.tensor as at from aesara.tensor.random.op import RandomVariable from aesara.tensor.random.var import RandomStateSharedVariable @@ -353,47 +352,6 @@ def get_moment(rv: TensorVariable) -> TensorVariable: return _get_moment(rv.owner.op, rv, size, *rv.owner.inputs[3:]) -class NoDistribution(Distribution): - def __init__( - self, - shape, - dtype, - initval=None, - defaults=(), - parent_dist=None, - *args, - **kwargs, - ): - super().__init__( - shape=shape, dtype=dtype, initval=initval, defaults=defaults, *args, **kwargs - ) - self.parent_dist = parent_dist - - def __getattr__(self, name): - # Do not use __getstate__ and __setstate__ from parent_dist - # to avoid infinite recursion during unpickling - if name.startswith("__"): - raise AttributeError("'NoDistribution' has no attribute '%s'" % name) - return getattr(self.parent_dist, name) - - def logp(self, x): - """Calculate log probability. - - Parameters - ---------- - x: numeric - Value for which log-probability is calculated. - - Returns - ------- - TensorVariable - """ - return at.zeros_like(x) - - def _distr_parameters_for_repr(self): - return [] - - class Discrete(Distribution): """Base class for discrete distributions""" @@ -409,6 +367,13 @@ class Continuous(Distribution): """Base class for continuous distributions""" +class NoDistribution(Distribution): + """Base class for artifical distributions + + RandomVariables that share this type are allowed in logprob graphs + """ + + class DensityDist(Distribution): """Distribution based on a given log density function. diff --git a/pymc3/distributions/simulator.py b/pymc3/distributions/simulator.py index 0b0fba1d30..0832e92809 100644 --- a/pymc3/distributions/simulator.py +++ b/pymc3/distributions/simulator.py @@ -14,112 +14,244 @@ import logging +import aesara +import aesara.tensor as at import numpy as np +from aesara.graph.op import Apply, Op +from aesara.tensor.random.op import RandomVariable +from aesara.tensor.var import TensorVariable from scipy.spatial import cKDTree +from pymc3.aesaraf import floatX from pymc3.distributions.distribution import NoDistribution +from pymc3.distributions.logprob import _logp __all__ = ["Simulator"] _log = logging.getLogger("pymc3") +class SimulatorRV(RandomVariable): + """ + Base class for SimulatorRVs + + This should be subclassed when defining custom Simulator objects. + """ + + name = "SimulatorRV" + ndim_supp = None + ndims_params = None + dtype = "floatX" + _print_name = ("Simulator", "\\operatorname{Simulator}") + + fn = None + _distance = None + _sum_stat = None + epsilon = None + + @classmethod + def rng_fn(cls, *args, **kwargs): + return cls.fn(*args, **kwargs) + + @classmethod + def distance(cls, *args, **kwargs): + return cls._distance(*args, **kwargs) + + @classmethod + def sum_stat(cls, *args, **kwargs): + return cls._sum_stat(*args, **kwargs) + + class Simulator(NoDistribution): r""" - Define a simulator, from a Python function, to be used in ABC methods. + Simulator distribution, used for Approximate Bayesian Inference (ABC) + with Sequential Monte Carlo (SMC) sampling via ``pm.sample_smc``. + + Simulator distributions have a stochastic pseudo-loglikelihood defined by + a distance metric between the observed and simulated data, and tweaked + by a hyper-parameter ``epsilon``. Parameters ---------- - function: callable - Python function defined by the user. + fn: callable + Python random simulator function. Should expect the following signature + ``(rng, arg1, arg2, ... argn, size)``, where rng is a ``numpy.random.RandomStream()`` + and ``size`` defines the size of the desired sample. params: list - Parameters passed to function. - distance: str or callable - Distance functions. Available options are "gaussian" (default), "laplacian", - "kullback_leibler" or a user defined function that takes epsilon, the summary statistics of - observed_data and the summary statistics of simulated_data as input. - ``gaussian`` :math: `-0.5 \left(\left(\frac{xo - xs}{\epsilon}\right)^2\right)` - ``laplace`` :math: `{\left(\frac{|xo - xs|}{\epsilon}\right)}` - ``kullback_leibler`` `:math: d \sum(-\log(\frac{\nu_d} {\rho_d}) / \epsilon) + log_r` - gaussian + ``sum_stat="sort"`` is equivalent to the 1D 2-wasserstein distance - laplace + ``sum_stat="sort"`` is equivalent to the the 1D 1-wasserstein distance - sum_stat: str or callable - Summary statistics. Available options are ``indentity``, ``sort``, ``mean``, ``median``. - If a callable is based it should return a number or a 1d numpy array. + Parameters used by the Simulator random function. Parameters can also + be passed by order, in which case the keyword argument ``params`` is + ignored. Alternatively, each parameter can be passed by order after fn, + ``param1, param2, ..., paramN`` + distance : Aesara Op, callable or str + Distance function. Available options are ``"gaussian"`` (default), ``"laplace"``, + ``"kullback_leibler"`` or a user defined function (or Aesara Op) that takes + ``epsilon``, the summary statistics of observed_data and the summary statistics + of simulated_data as input. + + ``gaussian``: :math:`-0.5 \left(\left(\frac{xo - xs}{\epsilon}\right)^2\right)` + + ``laplace``: :math:`-{\left(\frac{|xo - xs|}{\epsilon}\right)}` + + ``kullback_leibler``: :math:`\frac{d}{n} \frac{1}{\epsilon} \sum_i^n -\log \left( \frac{{\nu_d}_i}{{\rho_d}_i} \right) + \log_r` [1]_ + + ``distance="gaussian"`` + ``sum_stat="sort"`` is equivalent to the 1D 2-wasserstein distance + + ``distance="laplace"`` + ``sum_stat="sort"`` is equivalent to the the 1D 1-wasserstein distance + sum_stat: Aesara Op, callable or str + Summary statistic function. Available options are ``"indentity"`` (default), + ``"sort"``, ``"mean"``, ``"median"``. If a callable (or Aesara Op) is defined, + it should return a 1d numpy array (or Aesara vector). epsilon: float or array - Scaling parameter for the distance functions. It should be a float or an array of the - same size of the output of ``sum_stat``. - *args and **kwargs: - Arguments and keywords arguments that the function takes. + Scaling parameter for the distance functions. It should be a float or + an array of the same size of the output of ``sum_stat``. Defaults to ``1.0`` + ndim_supp : int + Number of dimensions of the SimulatorRV (0 for scalar, 1 for vector, etc.) + Defaults to ``0``. + ndims_params: list[int] + Number of minimum dimensions of each parameter of the RV. For example, + if the Simulator accepts two scalar inputs, it should be ``[0, 0]``. + Defaults to ``0`` for each parameter. + + Examples + -------- + .. code-block:: python + + def my_simulator_fn(rng, loc, scale, size): + return rng.normal(loc, scale, size=size) + + with pm.Model() as m: + loc = pm.Normal("loc", 0, 1) + scale = pm.HalfNormal("scale", 1) + simulator = pm.Simulator("sim", my_simulator, loc, scale, observed=data) + idata = pm.sample_smc() + + References + ---------- + .. [1] Pérez-Cruz, F. (2008, July). Kullback-Leibler divergence + estimation of continuous distributions. In 2008 IEEE international + symposium on information theory (pp. 1666-1670). IEEE. + `link `__ + """ - def __init__( - self, - function, - *args, + def __new__( + cls, + name, + fn, + *unnamed_params, params=None, distance="gaussian", sum_stat="identity", epsilon=1, + ndim_supp=0, + ndims_params=None, + dtype="floatX", **kwargs, ): - self.function = function - self.params = params - observed = self.data - self.epsilon = epsilon - - if distance == "gaussian": - self.distance = gaussian - elif distance == "laplace": - self.distance = laplace - elif distance == "kullback_leibler": - self.distance = KullbackLiebler(observed) - if sum_stat != "identity": - _log.info(f"Automatically setting sum_stat to identity as expected by {distance}") - sum_stat = "identity" - elif hasattr(distance, "__call__"): - self.distance = distance - else: - raise ValueError(f"The distance metric {distance} is not implemented") - - if sum_stat == "identity": - self.sum_stat = identity - elif sum_stat == "sort": - self.sum_stat = np.sort - elif sum_stat == "mean": - self.sum_stat = np.mean - elif sum_stat == "median": - self.sum_stat = np.median - elif hasattr(sum_stat, "__call__"): - self.sum_stat = sum_stat + + if not isinstance(distance, Op): + if distance == "gaussian": + distance = gaussian + elif distance == "laplace": + distance = laplace + elif distance == "kullback_leibler": + raise NotImplementedError("KL not refactored yet") + # TODO: Wrap KL in aesara OP + # distance = KullbackLiebler(observed) + # if sum_stat != "identity": + # _log.info(f"Automatically setting sum_stat to identity as expected by {distance}") + # sum_stat = "identity" + elif callable(distance): + distance = create_distance_op_from_fn(distance) + else: + raise ValueError(f"The distance metric {distance} is not implemented") + + if not isinstance(sum_stat, Op): + if sum_stat == "identity": + sum_stat = identity + elif sum_stat == "sort": + sum_stat = at.sort + elif sum_stat == "mean": + sum_stat = at.mean + elif sum_stat == "median": + # Missing in Aesara, see aesara/issues/525 + sum_stat = create_sum_stat_op_from_fn(np.median) + elif callable(sum_stat): + sum_stat = create_sum_stat_op_from_fn(sum_stat) + else: + raise ValueError(f"The summary statistic {sum_stat} is not implemented") + + epsilon = at.as_tensor_variable(floatX(epsilon)) + + if params is None: + params = unnamed_params else: - raise ValueError(f"The summary statistics {sum_stat} is not implemented") - - super().__init__(shape=np.prod(observed.shape), dtype=observed.dtype, *args, **kwargs) - - def random(self, point=None, size=None): - """ - Draw random values from Simulator. - - Parameters - ---------- - point: dict, optional - Dict of variable values on which random values are to be conditioned (uses default - point if not specified). - size: int, optional - Desired size of random sample (returns one sample if not specified). - - Returns - ------- - array - """ - # size = to_tuple(size) - # params = draw_values([*self.params], point=point, size=size) - # if len(size) == 0: - # return self.function(*params) - # else: - # return np.array([self.function(*params) for _ in range(size[0])]) + if unnamed_params: + raise ValueError("Cannot pass both unnamed parameters and `params`") + + # Assume scalar ndims_params + if ndims_params is None: + ndims_params = [0] * len(params) + + sim_op = type( + f"Simulator_{name}", + (SimulatorRV,), + dict( + name="Simulator", + ndim_supp=ndim_supp, + ndims_params=ndims_params, + dtype=dtype, + inplace=False, + fn=fn, + _distance=distance, + _sum_stat=sum_stat, + epsilon=epsilon, + ), + )() + + # The logp function is registered to the more general SimulatorRV, + # in order to avoid issues with multiprocessing / pickling + + # rv_type = type(sim_op) + # NoDistribution.register(rv_type) + NoDistribution.register(SimulatorRV) + + # @_logp.register(rv_type) + @_logp.register(SimulatorRV) + def logp(op, sim_rv, rvs_to_values, *sim_params, **kwargs): + value_var = rvs_to_values.get(sim_rv, sim_rv) + return cls.logp( + value_var, + sim_rv, + ) + + cls.rv_op = sim_op + return super().__new__(cls, name, *params, **kwargs) + + @classmethod + def dist(cls, *params, **kwargs): + return super().dist(params, **kwargs) + + @classmethod + def logp(cls, value, sim_rv): + # Use a new rng to avoid non-randomness in parallel sampling + # TODO: Model rngs should be updated prior to multiprocessing split, + # in which case this would not be needed. However, that would have to be + # done for every sampler that may accomodate Simulators + rng = aesara.shared(np.random.default_rng()) + rng.tag.is_rng = True + + # Create a new simulatorRV with identical inputs as the original one + sim_op = sim_rv.owner.op + sim_value = sim_op.make_node(rng, *sim_rv.owner.inputs[1:]).default_output() + sim_value.name = "sim_value" + + return sim_op.distance( + sim_op.epsilon, + sim_op.sum_stat(value), + sim_op.sum_stat(sim_value), + ) def identity(x): @@ -134,7 +266,7 @@ def gaussian(epsilon, obs_data, sim_data): def laplace(epsilon, obs_data, sim_data): """Laplace kernel.""" - return -np.abs((obs_data - sim_data) / epsilon) + return -at.abs_((obs_data - sim_data) / epsilon) class KullbackLiebler: @@ -155,3 +287,53 @@ def __call__(self, epsilon, obs_data, sim_data): sim_data = sim_data[:, None] nu_d, _ = cKDTree(sim_data).query(self.obs_data, 1) return self.d_n * np.sum(-np.log(nu_d / self.rho_d) / epsilon) + self.log_r + + +scalarX = at.dscalar if aesara.config.floatX == "float64" else at.fscalar +vectorX = at.dvector if aesara.config.floatX == "float64" else at.fvector + + +def create_sum_stat_op_from_fn(fn): + # Check if callable returns TensorVariable with dummy inputs + try: + res = fn(vectorX()) + if isinstance(res, TensorVariable): + return fn + except Exception: + pass + + # Otherwise, automatically wrap in Aesara Op + class SumStat(Op): + def make_node(self, x): + x = at.as_tensor_variable(x) + return Apply(self, [x], [vectorX()]) + + def perform(self, node, inputs, outputs): + (x,) = inputs + outputs[0][0] = np.atleast_1d(fn(x)).astype(aesara.config.floatX) + + return SumStat() + + +def create_distance_op_from_fn(fn): + # Check if callable returns TensorVariable with dummy inputs + try: + res = fn(scalarX(), vectorX(), vectorX()) + if isinstance(res, TensorVariable): + return fn + except Exception: + pass + + # Otherwise, automatically wrap in Aesara Op + class Distance(Op): + def make_node(self, epsilon, obs_data, sim_data): + epsilon = at.as_tensor_variable(epsilon) + obs_data = at.as_tensor_variable(obs_data) + sim_data = at.as_tensor_variable(sim_data) + return Apply(self, [epsilon, obs_data, sim_data], [vectorX()]) + + def perform(self, node, inputs, outputs): + eps, obs_data, sim_data = inputs + outputs[0][0] = np.atleast_1d(fn(eps, obs_data, sim_data)).astype(aesara.config.floatX) + + return Distance() diff --git a/pymc3/smc/sample_smc.py b/pymc3/smc/sample_smc.py index 1927cebe75..1f48857285 100644 --- a/pymc3/smc/sample_smc.py +++ b/pymc3/smc/sample_smc.py @@ -182,15 +182,6 @@ def sample_smc( if parallel is False: cores = 1 - _log = logging.getLogger("pymc3") - _log.info("Initializing SMC sampler...") - - model = modelcontext(model) - if model.name: - raise NotImplementedError( - "The SMC implementation currently does not support named models. " - "See https://github.com/pymc-devs/pymc3/pull/4365." - ) if cores is None: cores = _cpu_count() @@ -199,11 +190,6 @@ def sample_smc( else: cores = min(chains, cores) - _log.info( - f"Sampling {chains} chain{'s' if chains > 1 else ''} " - f"in {cores} job{'s' if cores > 1 else ''}" - ) - if random_seed == -1: random_seed = None if chains == 1 and isinstance(random_seed, int): @@ -215,6 +201,15 @@ def sample_smc( if not isinstance(random_seed, Iterable): raise TypeError("Invalid value for `random_seed`. Must be tuple, list or int") + model = modelcontext(model) + + _log = logging.getLogger("pymc3") + _log.info("Initializing SMC sampler...") + _log.info( + f"Sampling {chains} chain{'s' if chains > 1 else ''} " + f"in {cores} job{'s' if cores > 1 else ''}" + ) + params = ( draws, kernel, diff --git a/pymc3/smc/smc.py b/pymc3/smc/smc.py index 420884a9ff..5840d951d3 100644 --- a/pymc3/smc/smc.py +++ b/pymc3/smc/smc.py @@ -19,12 +19,12 @@ import aesara.tensor as at import numpy as np -from aesara import function as aesara_function from aesara.graph.basic import clone_replace from scipy.special import logsumexp from scipy.stats import multivariate_normal from pymc3.aesaraf import ( + compile_rv_inplace, floatX, inputvars, join_nonshared_inputs, @@ -575,6 +575,6 @@ def _logp_forw(point, out_vars, vars, shared): vars = new_vars out_list, inarray0 = join_nonshared_inputs(point, out_vars, vars, shared) - f = aesara_function([inarray0], out_list[0]) + f = compile_rv_inplace([inarray0], out_list[0]) f.trust_input = True return f diff --git a/pymc3/tests/test_smc.py b/pymc3/tests/test_smc.py index 512eb008ef..736b518184 100644 --- a/pymc3/tests/test_smc.py +++ b/pymc3/tests/test_smc.py @@ -12,11 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import aesara import aesara.tensor as at import numpy as np import pytest import scipy.stats as st +from aesara.graph.basic import ancestors +from aesara.tensor.random.op import RandomVariable +from aesara.tensor.random.var import ( + RandomGeneratorSharedVariable, + RandomStateSharedVariable, +) +from aesara.tensor.sort import SortOp from arviz.data.inference_data import InferenceData import pymc3 as pm @@ -236,82 +244,256 @@ def test_deprecated_abc_args(self): pm.sample_smc(draws=10, chains=1, save_log_pseudolikelihood=True) -@pytest.mark.xfail(reason="SMC-ABC not refactored yet") -class TestSMCABC(SeededTest): +class TestSimulator(SeededTest): + """ + Tests for pm.Simulator. They are included in this file because Simulator was + designed primarily to be used with SMC sampling. + """ + + @staticmethod + def count_rvs(end_node): + return len( + [ + node + for node in ancestors([end_node]) + if node.owner is not None and isinstance(node.owner.op, RandomVariable) + ] + ) + + @staticmethod + def normal_sim(rng, a, b, size): + return rng.normal(a, b, size=size) + + @staticmethod + def abs_diff(eps, obs_data, sim_data): + return np.mean(np.abs((obs_data - sim_data) / eps)) + + @staticmethod + def quantiles(x): + return np.quantile(x, [0.25, 0.5, 0.75]) + def setup_class(self): super().setup_class() self.data = np.random.normal(loc=0, scale=1, size=1000) - def normal_sim(a, b): - return np.random.normal(a, b, 1000) - with pm.Model() as self.SMABC_test: a = pm.Normal("a", mu=0, sigma=1) b = pm.HalfNormal("b", sigma=1) - s = pm.Simulator( - "s", normal_sim, params=(a, b), sum_stat="sort", epsilon=1, observed=self.data - ) + s = pm.Simulator("s", self.normal_sim, a, b, sum_stat="sort", observed=self.data) self.s = s - def quantiles(x): - return np.quantile(x, [0.25, 0.5, 0.75]) + with pm.Model() as self.SMABC_potential: + a = pm.Normal("a", mu=0, sigma=1, initval=0.5) + b = pm.HalfNormal("b", sigma=1) + c = pm.Potential("c", pm.math.switch(a > 0, 0, -np.inf)) + s = pm.Simulator("s", self.normal_sim, a, b, observed=self.data) + + def test_one_gaussian(self): + assert self.count_rvs(self.SMABC_test.logpt) == 1 + + with self.SMABC_test: + trace = pm.sample_smc(draws=1000, return_inferencedata=False) + pr_p = pm.sample_prior_predictive(1000) + po_p = pm.sample_posterior_predictive(trace, 1000) + + assert abs(self.data.mean() - trace["a"].mean()) < 0.05 + assert abs(self.data.std() - trace["b"].mean()) < 0.05 - def abs_diff(eps, obs_data, sim_data): - return np.mean(np.abs((obs_data - sim_data) / eps)) + assert pr_p["s"].shape == (1000, 1000) + assert abs(0 - pr_p["s"].mean()) < 0.10 + assert abs(1.4 - pr_p["s"].std()) < 0.10 - with pm.Model() as self.SMABC_test2: + assert po_p["s"].shape == (1000, 1000) + assert abs(self.data.mean() - po_p["s"].mean()) < 0.10 + assert abs(self.data.std() - po_p["s"].std()) < 0.10 + + def test_custom_dist_sum_stat(self): + with pm.Model() as m: a = pm.Normal("a", mu=0, sigma=1) b = pm.HalfNormal("b", sigma=1) s = pm.Simulator( "s", - normal_sim, - params=(a, b), - distance=abs_diff, - sum_stat=quantiles, - epsilon=1, + self.normal_sim, + a, + b, + distance=self.abs_diff, + sum_stat=self.quantiles, observed=self.data, ) - with pm.Model() as self.SMABC_potential: - a = pm.Normal("a", mu=0, sigma=1) - b = pm.HalfNormal("b", sigma=1) - c = pm.Potential("c", pm.math.switch(a > 0, 0, -np.inf)) - s = pm.Simulator( - "s", normal_sim, params=(a, b), sum_stat="sort", epsilon=1, observed=self.data - ) + assert self.count_rvs(m.logpt) == 1 - def test_one_gaussian(self): - with self.SMABC_test: - trace = pm.sample_smc(draws=1000, kernel="ABC") + with m: + pm.sample_smc(draws=100) - np.testing.assert_almost_equal(self.data.mean(), trace["a"].mean(), decimal=2) - np.testing.assert_almost_equal(self.data.std(), trace["b"].mean(), decimal=1) + def test_custom_dist_sum_stat_scalar(self): + """ + Test that automatically wrapped functions cope well with scalar inputs + """ + scalar_data = 5 - def test_sim_data_ppc(self): - with self.SMABC_test: - trace, sim_data = pm.sample_smc(draws=1000, kernel="ABC", chains=2, save_sim_data=True) - pr_p = pm.sample_prior_predictive(1000) - po_p = pm.sample_posterior_predictive(trace, 1000) + with pm.Model() as m: + s = pm.Simulator( + "s", + self.normal_sim, + 0, + 1, + distance=self.abs_diff, + sum_stat=self.quantiles, + observed=scalar_data, + ) + assert self.count_rvs(m.logpt) == 1 - assert sim_data["s"].shape == (2, 1000, 1000) - np.testing.assert_almost_equal(self.data.mean(), sim_data["s"].mean(), decimal=2) - np.testing.assert_almost_equal(self.data.std(), sim_data["s"].std(), decimal=1) - assert pr_p["s"].shape == (1000, 1000) - np.testing.assert_almost_equal(0, pr_p["s"].mean(), decimal=1) - np.testing.assert_almost_equal(1.4, pr_p["s"].std(), decimal=1) - assert po_p["s"].shape == (1000, 1000) - np.testing.assert_almost_equal(0, po_p["s"].mean(), decimal=2) - np.testing.assert_almost_equal(1, po_p["s"].std(), decimal=1) + with pm.Model() as m: + s = pm.Simulator( + "s", + self.normal_sim, + 0, + 1, + distance=self.abs_diff, + sum_stat="mean", + observed=scalar_data, + ) + assert self.count_rvs(m.logpt) == 1 - def test_custom_dist_sum(self): - with self.SMABC_test2: - trace = pm.sample_smc(draws=1000, kernel="ABC") + def test_model_with_potential(self): + assert self.count_rvs(self.SMABC_potential.logpt) == 1 - def test_potential(self): with self.SMABC_potential: - trace = pm.sample_smc(draws=1000, kernel="ABC") + trace = pm.sample_smc(draws=100, chains=1, return_inferencedata=False) assert np.all(trace["a"] >= 0) + def test_simulator_metropolis_mcmc(self): + with self.SMABC_test as m: + step = pm.Metropolis([m.rvs_to_values[m["a"]], m.rvs_to_values[m["b"]]]) + trace = pm.sample(step=step, return_inferencedata=False) + + assert abs(self.data.mean() - trace["a"].mean()) < 0.05 + assert abs(self.data.std() - trace["b"].mean()) < 0.05 + + def test_multiple_simulators(self): + true_a = 2 + true_b = -2 + + data1 = np.random.normal(true_a, 0.1, size=1000) + data2 = np.random.normal(true_b, 0.1, size=1000) + + with pm.Model() as m: + a = pm.Normal("a", mu=0, sigma=3) + b = pm.Normal("b", mu=0, sigma=3) + sim1 = pm.Simulator( + "sim1", + self.normal_sim, + a, + 0.1, + distance="gaussian", + sum_stat="sort", + observed=data1, + ) + sim2 = pm.Simulator( + "sim2", + self.normal_sim, + b, + 0.1, + distance="laplace", + sum_stat="mean", + epsilon=0.1, + observed=data2, + ) + + assert self.count_rvs(m.logpt) == 2 + + # Check that the logps use the correct methods + a_val = m.rvs_to_values[a] + sim1_val = m.rvs_to_values[sim1] + logp_sim1 = pm.logp(sim1, sim1_val) + logp_sim1_fn = aesara.function([sim1_val, a_val], logp_sim1) + + b_val = m.rvs_to_values[b] + sim2_val = m.rvs_to_values[sim2] + logp_sim2 = pm.logp(sim2, sim2_val) + logp_sim2_fn = aesara.function([sim2_val, b_val], logp_sim2) + + assert any( + node for node in logp_sim1_fn.maker.fgraph.toposort() if isinstance(node.op, SortOp) + ) + + assert not any( + node for node in logp_sim2_fn.maker.fgraph.toposort() if isinstance(node.op, SortOp) + ) + + with m: + trace = pm.sample_smc(return_inferencedata=False) + + assert abs(true_a - trace["a"].mean()) < 0.05 + assert abs(true_b - trace["b"].mean()) < 0.05 + + def test_nested_simulators(self): + true_a = 2 + rng = self.get_random_state() + data = rng.normal(true_a, 0.1, size=1000) + + with pm.Model() as m: + sim1 = pm.Simulator( + "sim1", + self.normal_sim, + params=(0, 4), + distance="gaussian", + sum_stat="identity", + ) + sim2 = pm.Simulator( + "sim2", + self.normal_sim, + params=(sim1, 0.1), + distance="gaussian", + sum_stat="mean", + epsilon=0.1, + observed=data, + ) + + assert self.count_rvs(m.logpt) == 2 + + with m: + trace = pm.sample_smc(return_inferencedata=False) + + assert np.abs(true_a - trace["sim1"].mean()) < 0.1 + + def test_upstream_rngs_not_in_compiled_logp(self): + smc = IMH(model=self.SMABC_test) + smc.initialize_population() + smc._initialize_kernel() + likelihood_func = smc.likelihood_logp_func + + # Test graph is stochastic + inarray = floatX(np.array([0, 0])) + assert likelihood_func(inarray) != likelihood_func(inarray) + + # Test only one shared RNG is present + compiled_graph = likelihood_func.maker.fgraph.outputs + shared_rng_vars = [ + node + for node in ancestors(compiled_graph) + if isinstance(node, (RandomStateSharedVariable, RandomGeneratorSharedVariable)) + ] + assert len(shared_rng_vars) == 1 + + def test_simulator_error_msg(self): + msg = "The distance metric not_real is not implemented" + with pytest.raises(ValueError, match=msg): + with pm.Model() as m: + sim = pm.Simulator("sim", self.normal_sim, 0, 1, distance="not_real") + + msg = "The summary statistic not_real is not implemented" + with pytest.raises(ValueError, match=msg): + with pm.Model() as m: + sim = pm.Simulator("sim", self.normal_sim, 0, 1, sum_stat="not_real") + + msg = "Cannot pass both unnamed parameters and `params`" + with pytest.raises(ValueError, match=msg): + with pm.Model() as m: + sim = pm.Simulator("sim", self.normal_sim, 0, params=(1)) + + @pytest.mark.xfail(reason="KL not refactored") def test_automatic_use_of_sort(self): with pm.Model() as model: s_k = pm.Simulator( @@ -324,31 +506,26 @@ def test_automatic_use_of_sort(self): ) assert s_k.distribution.sum_stat is pm.distributions.simulator.identity - def test_repr_latex(self): - expected = "$\\text{s} \\sim \\text{Simulator}(\\text{normal_sim}(a, b), \\text{gaussian}, \\text{sort})$" - assert expected == self.s._repr_latex_() - assert self.s._repr_latex_() == self.s.__latex__() - assert self.SMABC_test.model._repr_latex_() == self.SMABC_test.model.__latex__() - def test_name_is_string_type(self): with self.SMABC_potential: assert not self.SMABC_potential.name - trace = pm.sample_smc(draws=10, kernel="ABC") + trace = pm.sample_smc(draws=10, chains=1, return_inferencedata=False) assert isinstance(trace._straces[0].name, str) - def test_named_models_are_unsupported(self): - def normal_sim(a, b): - return np.random.normal(a, b, 1000) - - with pm.Model(name="NamedModel"): + def test_named_model(self): + # Named models used to fail with Simulator because the arguments to the + # random fn used to be passed by name. This is no longer true. + # https://github.com/pymc-devs/pymc3/pull/4365#issuecomment-761221146 + name = "NamedModel" + with pm.Model(name=name): a = pm.Normal("a", mu=0, sigma=1) b = pm.HalfNormal("b", sigma=1) - c = pm.Potential("c", pm.math.switch(a > 0, 0, -np.inf)) - s = pm.Simulator( - "s", normal_sim, params=(a, b), sum_stat="sort", epsilon=1, observed=self.data - ) - with pytest.raises(NotImplementedError, match="named models"): - pm.sample_smc(draws=10, kernel="ABC") + s = pm.Simulator("s", self.normal_sim, a, b, observed=self.data) + + trace = pm.sample_smc(draws=10, chains=2, return_inferencedata=False) + assert f"{name}_a" in trace.varnames + assert f"{name}_b" in trace.varnames + assert f"{name}_b_log__" in trace.varnames class TestMHKernel(SeededTest):