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 fdc6e45e71..df9ba0b9a6 100644 --- a/pymc3/aesaraf.py +++ b/pymc3/aesaraf.py @@ -354,11 +354,21 @@ 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 [] + + # TODO: Importing at top is creating a circular dependency + from pymc3.distributions.simulator import SimulatorRV + + # If orphan RandomVariable is a SimulatorRV, we allow for further + # replacements in upstream graph + if isinstance(rv_var.owner.op, SimulatorRV): + return var.owner.inputs[3:] + + 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) diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index 3f3f3d1e2d..3943f7483a 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -22,7 +22,6 @@ from typing import Optional import aesara -import aesara.tensor as at import dill from aesara.tensor.random.op import RandomVariable @@ -386,47 +385,6 @@ def _repr_latex_(self, *, formatting="latex_with_params", **kwargs): __latex__ = _repr_latex_ -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""" @@ -442,6 +400,10 @@ class Continuous(Distribution): """Base class for continuous distributions""" +class NoDistribution(Distribution): + """Base class for artifical distributions""" + + class DensityDist(Distribution): """Distribution based on a given log density function. diff --git a/pymc3/distributions/simulator.py b/pymc3/distributions/simulator.py index 8b5951b1ad..48b1beeb88 100644 --- a/pymc3/distributions/simulator.py +++ b/pymc3/distributions/simulator.py @@ -14,17 +14,52 @@ import logging +import aesara +import aesara.tensor as at import numpy as np +from aesara.graph.op import Op +from aesara.tensor.random.op import RandomVariable from scipy.spatial import cKDTree +from pymc3.aesaraf import floatX from pymc3.distributions.distribution import NoDistribution +from pymc3.distributions.logp import _logp __all__ = ["Simulator"] _log = logging.getLogger("pymc3") +class SimulatorRV(RandomVariable): + """A placeholder for Simulator RVs""" + + name = "SimulatorRV" + _print_name = ("Simulator", "\\operatorname{Simulator}") + fn = None + epsilon = None + distance = None + sum_stat = None + + @classmethod + def rng_fn(cls, *args, **kwargs): + if cls.fn is None: + raise ValueError(f"fn was not defined for {cls}") + return cls.fn(*args, **kwargs) + + @classmethod + def _distance(cls, epsilon, value, sim_value): + if cls.distance is None: + raise ValueError(f"distance function was not defined for {cls}") + return cls.distance(epsilon, value, sim_value) + + @classmethod + def _sum_stat(cls, value): + if cls.sum_stat is None: + raise ValueError(f"sum_stat function was not defined for {cls}") + return cls.sum_stat(value) + + class Simulator(NoDistribution): r""" Define a simulator, from a Python function, to be used in ABC methods. @@ -54,86 +89,104 @@ class Simulator(NoDistribution): Arguments and keywords arguments that the function takes. """ - def __init__( - self, - function, - *args, + def __new__( + cls, + name, + fn, + *, params=None, distance="gaussian", sum_stat="identity", epsilon=1, + observed=None, + 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 - 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])]) - - def _str_repr(self, name=None, dist=None, formatting="plain"): - if dist is None: - dist = self - name = name - function = dist.function.__name__ - params = ", ".join([var.name for var in dist.params]) - sum_stat = self.sum_stat.__name__ if hasattr(self.sum_stat, "__call__") else self.sum_stat - distance = getattr(self.distance, "__name__", self.distance.__class__.__name__) - - if "latex" in formatting: - return f"$\\text{{{name}}} \\sim \\text{{Simulator}}(\\text{{{function}}}({params}), \\text{{{distance}}}, \\text{{{sum_stat}}})$" - else: - return f"{name} ~ Simulator({function}({params}), {distance}, {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": + sum_stat = at.median + elif callable(sum_stat): + sum_stat = create_sum_stat_op_from_fn(sum_stat) + else: + raise ValueError(f"The summary statistics {sum_stat} is not implemented") + + epsilon = at.as_tensor_variable(floatX(epsilon)) + + if params is None: + 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, + # Specifc to Simulator + fn=fn, + distance=distance, + sum_stat=sum_stat, + epsilon=epsilon, + ), + )() + + # Register custom logp + rv_type = type(sim_op) + + @_logp.register(rv_type) + def logp(op, sim_rv, rvs_to_values, *sim_params, **kwargs): + value_var = rvs_to_values.get(sim_rv, sim_rv) + return Simulator.logp( + value_var, + sim_rv, + ) + + cls.rv_op = sim_op + return super().__new__(cls, name, params, observed=observed, **kwargs) + + @classmethod + def logp(cls, value, sim_rv): + # Create a new simulatorRV identically to the original one + sim_op = sim_rv.owner.op + sim_data = at.as_tensor_variable(sim_op.make_node(*sim_rv.owner.inputs)) + sim_data.name = "sim_data" + return sim_op._distance( + sim_op.epsilon, + sim_op._sum_stat(value), + sim_op._sum_stat(sim_data), + ) def identity(x): @@ -148,7 +201,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: @@ -169,3 +222,35 @@ 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 + + +def create_sum_stat_op_from_fn(fn): + class SumStat(Op): + if aesara.config.floatX == "float64": + itypes = [at.dvector] + otypes = [at.dvector] + else: + itypes = [at.fvector] + otypes = [at.fvector] + + 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): + class Distance(Op): + if aesara.config.floatX == "float64": + itypes = [at.dscalar, at.dvector, at.dvector] + otypes = [at.dvector] + else: + itypes = [at.fscalar, at.fvector, at.fvector] + otypes = [at.fvector] + + 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 e7d8884859..d386ebfac1 100644 --- a/pymc3/smc/sample_smc.py +++ b/pymc3/smc/sample_smc.py @@ -35,15 +35,15 @@ def sample_smc( draws=2000, - kernel="metropolis", + kernel=None, n_steps=25, *, start=None, tune_steps=True, p_acc_rate=0.85, threshold=0.5, - save_sim_data=False, - save_log_pseudolikelihood=True, + save_sim_data=None, + save_log_pseudolikelihood=None, model=None, random_seed=-1, parallel=None, @@ -62,9 +62,6 @@ def sample_smc( draws: int The number of samples to draw from the posterior (i.e. last stage). And also the number of independent chains. Defaults to 2000. - kernel: str - Kernel method for the SMC sampler. Available option are ``metropolis`` (default) and `ABC`. - Use `ABC` for likelihood free inference together with a ``pm.Simulator``. n_steps: int The number of steps of each Markov Chain. If ``tune_steps == True`` ``n_steps`` will be used for the first stage and for the others it will be determined automatically based on the @@ -82,13 +79,6 @@ def sample_smc( Determines the change of beta from stage to stage, i.e.indirectly the number of stages, the higher the value of `threshold` the higher the number of stages. Defaults to 0.5. It should be between 0 and 1. - save_sim_data : bool - Whether or not to save the simulated data. This parameter only works with the ABC kernel. - The stored data corresponds to a samples from the posterior predictive distribution. - save_log_pseudolikelihood : bool - Whether or not to save the log pseudolikelihood values. This parameter only works with the - ABC kernel. The stored data can be used to compute LOO or WAIC values. Computing LOO/WAIC - values from log pseudolikelihood values is experimental. model: Model (optional if in ``with`` context)). random_seed: int random seed @@ -156,6 +146,29 @@ def sample_smc( %282007%29133:7%28816%29>`__ """ + if kernel is not None: + warnings.warn( + "The argument kernel in sample_smc has been deprecated. It is no " + "longer needed to specify it.", + DeprecationWarning, + stacklevel=2, + ) + if save_sim_data is not None: + warnings.warn( + "save_sim_data has been deprecated. Use pm.sample_posterior_predictive " + "to obtain the same type of samples.", + DeprecationWarning, + stacklevel=2, + ) + + if save_log_pseudolikelihood is not None: + warnings.warn( + "save_log_pseudolikelihood has been deprecated. This information is " + "now saved as log_likelihood in models with Simulator distributions.", + DeprecationWarning, + stacklevel=2, + ) + if parallel is not None: warnings.warn( "The argument parallel is deprecated, use the argument cores instead.", @@ -198,22 +211,13 @@ def sample_smc( if not isinstance(random_seed, Iterable): raise TypeError("Invalid value for `random_seed`. Must be tuple, list or int") - if kernel.lower() == "abc": - if len(model.observed_RVs) != 1: - warnings.warn("SMC-ABC only works properly with models with one observed variable") - if model.potentials: - _log.info("Potentials will be added to the prior term") - params = ( draws, - kernel, n_steps, start, tune_steps, p_acc_rate, threshold, - save_sim_data, - save_log_pseudolikelihood, model, ) @@ -241,9 +245,7 @@ def sample_smc( ( traces, - sim_data, log_marginal_likelihoods, - log_pseudolikelihood, betas, accept_ratios, nsteps, @@ -259,7 +261,6 @@ def sample_smc( trace.report._n_draws = draws trace.report._n_tune = _n_tune trace.report.log_marginal_likelihood = log_marginal_likelihoods - trace.report.log_pseudolikelihood = log_pseudolikelihood trace.report.betas = betas trace.report.accept_ratios = accept_ratios trace.report.nsteps = nsteps @@ -309,23 +310,16 @@ def sample_smc( trace.report._run_convergence_checks(idata, model) trace.report._log_summary() - posterior = idata if return_inferencedata else trace - if save_sim_data: - return posterior, {modelcontext(model).observed_RVs[0].name: np.array(sim_data)} - else: - return posterior + return idata if return_inferencedata else trace def sample_smc_int( draws, - kernel, n_steps, start, tune_steps, p_acc_rate, threshold, - save_sim_data, - save_log_pseudolikelihood, model, random_seed, chain, @@ -334,14 +328,11 @@ def sample_smc_int( """Run one SMC instance.""" smc = SMC( draws=draws, - kernel=kernel, n_steps=n_steps, start=start, tune_steps=tune_steps, p_acc_rate=p_acc_rate, threshold=threshold, - save_sim_data=save_sim_data, - save_log_pseudolikelihood=save_log_pseudolikelihood, model=model, random_seed=random_seed, chain=chain, @@ -377,9 +368,7 @@ def sample_smc_int( return ( smc.posterior_to_trace(), - smc.sim_data, smc.log_marginal_likelihood, - smc.log_pseudolikelihood, betas, accept_ratios, nsteps, diff --git a/pymc3/smc/smc.py b/pymc3/smc/smc.py index 9ac914e335..f7b6d77253 100644 --- a/pymc3/smc/smc.py +++ b/pymc3/smc/smc.py @@ -16,15 +16,14 @@ from collections import OrderedDict -import aesara.tensor as at import numpy as np from aesara import config -from aesara import function as aesara_function from scipy.special import logsumexp from scipy.stats import multivariate_normal from pymc3.aesaraf import ( + compile_rv_inplace, floatX, inputvars, join_nonshared_inputs, @@ -42,28 +41,22 @@ class SMC: def __init__( self, draws=2000, - kernel="metropolis", n_steps=25, start=None, tune_steps=True, p_acc_rate=0.85, threshold=0.5, - save_sim_data=False, - save_log_pseudolikelihood=True, model=None, random_seed=-1, chain=0, ): self.draws = draws - self.kernel = kernel.lower() self.n_steps = n_steps self.start = start self.tune_steps = tune_steps self.p_acc_rate = p_acc_rate self.threshold = threshold - self.save_sim_data = save_sim_data - self.save_log_pseudolikelihood = save_log_pseudolikelihood self.model = model self.random_seed = random_seed self.chain = chain @@ -73,15 +66,22 @@ def __init__( if self.random_seed != -1: np.random.seed(self.random_seed) + self.var_info = None + self.posterior = None + self.prior_logp = None + self.likelihood_logp = None + self.posterior_logp = None + self.prior_logp_func = None + self.likelihood_logp_func = None + self.log_marginal_likelihood = 0 + self.beta = 0 self.max_steps = n_steps self.proposed = draws * n_steps self.acc_rate = 1 self.variables = inputvars(self.model.value_vars) self.weights = np.ones(self.draws) / self.draws - self.log_marginal_likelihood = 0 - self.sim_data = [] - self.log_pseudolikelihood = [] + self.cov = None def initialize_population(self): """Create an initial population from the prior distribution.""" @@ -102,7 +102,6 @@ def initialize_population(self): var_info[v.name] = (init[v.name].shape, init[v.name].size) for i in range(self.draws): - point = Point({v.name: init_rnd[v.name][i] for v in self.variables}, model=self.model) population.append(DictToArrayBijection.map(point).data) @@ -114,36 +113,12 @@ def setup_kernel(self): initial_values = self.model.initial_point shared = make_shared_replacements(initial_values, self.variables, self.model) - if self.kernel == "abc": - factors = [var.logpt for var in self.model.free_RVs] - factors += [at.sum(factor) for factor in self.model.potentials] - self.prior_logp_func = logp_forw( - initial_values, [at.sum(factors)], self.variables, shared - ) - simulator = self.model.observed_RVs[0] - distance = simulator.distribution.distance - sum_stat = simulator.distribution.sum_stat - self.likelihood_logp_func = PseudoLikelihood( - simulator.distribution.epsilon, - simulator.observations, - simulator.distribution.function, - [v.name for v in simulator.distribution.params], - self.model, - self.var_info, - self.variables, - distance, - sum_stat, - self.draws, - self.save_sim_data, - self.save_log_pseudolikelihood, - ) - elif self.kernel == "metropolis": - self.prior_logp_func = logp_forw( - initial_values, [self.model.varlogpt], self.variables, shared - ) - self.likelihood_logp_func = logp_forw( - initial_values, [self.model.datalogpt], self.variables, shared - ) + self.prior_logp_func = logp_forw( + initial_values, [self.model.varlogpt], self.variables, shared + ) + self.likelihood_logp_func = logp_forw( + initial_values, [self.model.datalogpt], self.variables, shared + ) def initialize_logp(self): """Initialize the prior and likelihood log probabilities.""" @@ -153,12 +128,6 @@ def initialize_logp(self): self.prior_logp = np.array(priors).squeeze() self.likelihood_logp = np.array(likelihoods).squeeze() - if self.kernel == "abc" and self.save_sim_data: - self.sim_data = self.likelihood_logp_func.get_data() - - if self.kernel == "abc" and self.save_log_pseudolikelihood: - self.log_pseudolikelihood = self.likelihood_logp_func.get_lpl() - def update_weights_beta(self): """Calculate the next inverse temperature (beta). @@ -201,8 +170,6 @@ def resample(self): self.prior_logp = self.prior_logp[resampling_indexes] self.likelihood_logp = self.likelihood_logp[resampling_indexes] self.posterior_logp = self.prior_logp + self.likelihood_logp * self.beta - if self.save_sim_data: - self.sim_data = self.sim_data[resampling_indexes] def update_proposal(self): """Update proposal based on the covariance matrix from tempered posterior.""" @@ -254,12 +221,6 @@ def mutate(self): self.prior_logp[accepted] = pl[accepted] self.likelihood_logp[accepted] = ll[accepted] - if self.kernel == "abc" and self.save_sim_data: - self.sim_data[accepted] = self.likelihood_logp_func.get_data()[accepted] - - if self.kernel == "abc" and self.save_log_pseudolikelihood: - self.log_pseudolikelihood[accepted] = self.likelihood_logp_func.get_lpl()[accepted] - self.acc_rate = np.mean(ac_) def posterior_to_trace(self): @@ -304,124 +265,8 @@ def logp_forw(point, out_vars, vars, shared): "together with aesara.config.floatX == `float32`", UserWarning, ) - f = aesara_function([inarray0], out_list[0], allow_input_downcast=True) + f = compile_rv_inplace([inarray0], out_list[0], allow_input_downcast=True) else: - f = aesara_function([inarray0], out_list[0]) - f.trust_input = False + f = compile_rv_inplace([inarray0], out_list[0]) + f.trust_input = True return f - - -class PseudoLikelihood: - """ - Pseudo Likelihood. - - epsilon: float - Standard deviation of the gaussian pseudo likelihood. - observations: array-like - observed data - function: python function - data simulator - params: list - names of the variables parameterizing the simulator. - model: PyMC3 model - var_info: dict - generated by ``SMC.initialize_population`` - variables: list - Model variables. - distance : str or callable - Distance function. - sum_stat: str or callable - Summary statistics. - size : int - Number of simulated datasets to save. When this number is exceeded the counter will be - restored to zero and it will start saving again. - save_sim_data : bool - whether to save or not the simulated data. - save_log_pseudolikelihood : bool - whether to save or not the log pseudolikelihood values. - """ - - def __init__( - self, - epsilon, - observations, - function, - params, - model, - var_info, - variables, - distance, - sum_stat, - size, - save_sim_data, - save_log_pseudolikelihood, - ): - self.epsilon = epsilon - self.function = function - self.params = params - self.model = model - self.var_info = var_info - self.variables = variables - self.varnames = [v.name for v in self.variables] - self.distance = distance - self.sum_stat = sum_stat - self.unobserved_RVs = [v.name for v in self.model.unobserved_RVs] - self.get_unobserved_fn = self.model.fastfn( - [v.tag.value_var for v in self.model.unobserved_RVs] - ) - self.size = size - self.save_sim_data = save_sim_data - self.save_log_pseudolikelihood = save_log_pseudolikelihood - self.sim_data_l = [] - self.lpl_l = [] - - self.observations = self.sum_stat(observations) - - def posterior_to_function(self, posterior): - """Turn posterior samples into function parameters to feed the simulator.""" - model = self.model - var_info = self.var_info - - varvalues = [] - samples = {} - size = 0 - for var in self.variables: - shape, new_size = var_info[var.name] - varvalues.append(posterior[size : size + new_size].reshape(shape)) - size += new_size - point = {k: v for k, v in zip(self.varnames, varvalues)} - for varname, value in zip(self.unobserved_RVs, self.get_unobserved_fn(point)): - if varname in self.params: - samples[varname] = value - return samples - - def save_data(self, sim_data): - """Save simulated data.""" - if len(self.sim_data_l) == self.size: - self.sim_data_l = [] - self.sim_data_l.append(sim_data) - - def get_data(self): - """Get simulated data.""" - return np.array(self.sim_data_l) - - def save_lpl(self, elemwise): - """Save log pseudolikelihood values.""" - if len(self.lpl_l) == self.size: - self.lpl_l = [] - self.lpl_l.append(elemwise) - - def get_lpl(self): - """Get log pseudolikelihood values.""" - return np.array(self.lpl_l) - - def __call__(self, posterior): - """Compute the pseudolikelihood.""" - func_parameters = self.posterior_to_function(posterior) - sim_data = self.function(**func_parameters) - if self.save_sim_data: - self.save_data(sim_data) - elemwise = self.distance(self.epsilon, self.observations, self.sum_stat(sim_data)) - if self.save_log_pseudolikelihood: - self.save_lpl(elemwise) - return elemwise.sum() diff --git a/pymc3/tests/test_smc.py b/pymc3/tests/test_smc.py index 3dbf0e37ca..fd0521f55b 100644 --- a/pymc3/tests/test_smc.py +++ b/pymc3/tests/test_smc.py @@ -19,11 +19,15 @@ import numpy as np import pytest +from aesara.graph.basic import ancestors +from aesara.tensor.random.op import RandomVariable +from aesara.tensor.sort import SortOp from arviz.data.inference_data import InferenceData import pymc3 as pm from pymc3.backends.base import MultiTrace +from pymc3.distributions.logp import _logp from pymc3.tests.helpers import SeededTest @@ -182,14 +186,13 @@ def test_depracated_parallel_arg(self): pm.sample_smc(draws=10, chains=1, parallel=False) -@pytest.mark.xfail(reason="SMC-ABC not refactored yet") class TestSMCABC(SeededTest): 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) + def normal_sim(rng, a, b, size): + return rng.normal(a, b, size=size) with pm.Model() as self.SMABC_test: a = pm.Normal("a", mu=0, sigma=1) @@ -219,7 +222,7 @@ def abs_diff(eps, obs_data, sim_data): ) with pm.Model() as self.SMABC_potential: - a = pm.Normal("a", mu=0, sigma=1) + 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( @@ -228,36 +231,32 @@ def abs_diff(eps, obs_data, sim_data): def test_one_gaussian(self): with self.SMABC_test: - trace = pm.sample_smc(draws=1000, kernel="ABC") - - 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_sim_data_ppc(self): - with self.SMABC_test: - trace, sim_data = pm.sample_smc(draws=1000, kernel="ABC", chains=2, save_sim_data=True) + trace = pm.sample_smc(draws=1000, chains=2, cores=1, return_inferencedata=False) pr_p = pm.sample_prior_predictive(1000) po_p = pm.sample_posterior_predictive(trace, 1000) - 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 abs(self.data.mean() - trace["a"].mean()) < 0.05 + assert abs(self.data.std() - trace["b"].mean()) < 0.05 + 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 abs(0 - pr_p["s"].mean()) < 0.10 + assert abs(1.4 - pr_p["s"].std()) < 0.10 + 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) + 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(self): with self.SMABC_test2: - trace = pm.sample_smc(draws=1000, kernel="ABC") + trace = pm.sample_smc(draws=100, chains=1) + @pytest.mark.xfail(reason="Potential is failing in SMC") def test_potential(self): with self.SMABC_potential: - trace = pm.sample_smc(draws=1000, kernel="ABC") + trace = pm.sample_smc(draws=1000, chains=1) assert np.all(trace["a"] >= 0) + @pytest.mark.xfail(reason="KL not refactored") def test_automatic_use_of_sort(self): with pm.Model() as model: s_k = pm.Simulator( @@ -270,16 +269,18 @@ def test_automatic_use_of_sort(self): ) assert s_k.distribution.sum_stat is pm.distributions.simulator.identity + @pytest.mark.xfail(reason="Latex not refactored") 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__() + @pytest.mark.xfail(reason="Potential is failing in SMC") 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, cores=1) assert isinstance(trace._straces[0].name, str) def test_named_models_are_unsupported(self): @@ -295,3 +296,94 @@ def normal_sim(a, b): ) with pytest.raises(NotImplementedError, match="named models"): pm.sample_smc(draws=10, kernel="ABC") + + def test_simulator_metropolis_mcmc(self): + with self.SMABC_potential 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): + def fn(rng, a, size): + return rng.normal(a, 0.1, size=size) + + true_a = 2 + true_b = -2 + + data1 = np.random.normal(true_a, 0.1, size=100) + data2 = np.random.normal(true_b, 0.1, size=100) + + with pm.Model() as m: + a = pm.Normal("a", 0, 1) + b = pm.Normal("b", 1) + + sim1 = pm.Simulator( + "sim1", + fn, + params=(a,), + sum_stat="sort", + distance="gaussian", + epsilon=1, + observed=data1, + ) + sim2 = pm.Simulator( + "sim2", + fn, + params=(b,), + sum_stat="identity", + distance="laplace", + epsilon=0.5, + observed=data2, + ) + + # Check that the logps use the correct methods + sim1_val = m.rvs_to_values[sim1] + logp_sim1 = _logp(sim1.owner.op, sim1, {sim1: sim1_val}) + logp_sim1_fn = aesara.function([sim1_val], logp_sim1) + + sim2_val = m.rvs_to_values[sim2] + logp_sim2 = _logp(sim2.owner.op, sim2, {sim2: sim2_val}) + logp_sim2_fn = aesara.function([sim2_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) + ) + + # Check that there are two RandomVariables in the graph + rvs = [ + node + for node in ancestors([m.logpt]) + if node.owner and isinstance(node.owner.op, RandomVariable) + ] + assert len(rvs) == 2 + + with m: + trace = pm.sample_smc(chains=1, return_inferencedata=False) + assert abs(true_a - trace["a"].mean()) < 0.05 + assert abs(true_b - trace["b"].mean()) < 0.05 + + def test_depracated_abc_args(self): + with self.SMABC_test: + with pytest.warns( + DeprecationWarning, + match="The argument kernel in sample_smc has been deprecated", + ): + pm.sample_smc(draws=10, chains=1, kernel="ABC") + + with pytest.warns( + DeprecationWarning, + match="save_sim_data has been deprecated", + ): + pm.sample_smc(draws=10, chains=1, save_sim_data=True) + + with pytest.warns( + DeprecationWarning, + match="save_log_pseudolikelihood has been deprecated", + ): + pm.sample_smc(draws=10, chains=1, save_log_pseudolikelihood=True)