def test_make_shared_replacements(self): """Check if pm.make_shared_replacements preserves broadcasting.""" with pm.Model() as test_model: test1 = pm.Normal("test1", mu=0.0, sigma=1.0, size=(1, 10)) test2 = pm.Normal("test2", mu=0.0, sigma=1.0, size=(10, 1)) # Replace test1 with a shared variable, keep test 2 the same replacement = pm.make_shared_replacements( test_model.compute_initial_point(), [test_model.test2], test_model) assert (test_model.test1.broadcastable == replacement[ test_model.test1.tag.value_var].broadcastable)
def __init__(self, vars=None, S=None, proposal_dist=None, lamb=None, scaling=0.001, tune=None, tune_interval=100, model=None, mode=None, **kwargs): model = pm.modelcontext(model) initial_values = model.initial_point() initial_values_size = sum(initial_values[n.name].size for n in model.value_vars) 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: S = np.ones(initial_values_size) if proposal_dist is not None: self.proposal_dist = proposal_dist(S) else: self.proposal_dist = UniformProposal(S) self.scaling = np.atleast_1d(scaling).astype("d") if lamb is None: # default to the optimal lambda for normally distributed targets lamb = 2.38 / np.sqrt(2 * initial_values_size) self.lamb = float(lamb) if tune not in {None, "scaling", "lambda"}: raise ValueError( 'The parameter "tune" must be one of {None, scaling, lambda}') self.tune = tune self.tune_interval = tune_interval self.steps_until_tune = tune_interval self.accepted = 0 self.mode = mode shared = pm.make_shared_replacements(initial_values, vars, model) self.delta_logp = delta_logp(initial_values, model.logpt(), vars, shared) super().__init__(vars, shared)
def __init__(self, vars=None, S=None, proposal_dist=None, scaling=1.0, tune=True, tune_interval=100, model=None, mode=None, **kwargs): """Create an instance of a Metropolis stepper Parameters ---------- vars: list List of value variables for sampler S: standard deviation or covariance matrix Some measure of variance to parameterize proposal distribution proposal_dist: function Function that returns zero-mean deviates when parameterized with S (and n). Defaults to normal. scaling: scalar or array Initial scale factor for proposal. Defaults to 1. tune: bool Flag for tuning. Defaults to True. tune_interval: int The frequency of tuning. Defaults to 100 iterations. model: PyMC Model Optional model for sampling step. Defaults to None (taken from context). mode: string or `Mode` instance. compilation mode passed to Aesara functions """ model = pm.modelcontext(model) initial_values = model.initial_point() 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: S = np.ones(sum(initial_values[v.name].size for v in vars)) if proposal_dist is not None: self.proposal_dist = proposal_dist(S) elif S.ndim == 1: self.proposal_dist = NormalProposal(S) elif S.ndim == 2: self.proposal_dist = MultivariateNormalProposal(S) else: raise ValueError("Invalid rank for variance: %s" % S.ndim) self.scaling = np.atleast_1d(scaling).astype("d") self.tune = tune self.tune_interval = tune_interval self.steps_until_tune = tune_interval self.accepted = 0 # Determine type of variables self.discrete = np.concatenate([[v.dtype in pm.discrete_types] * (initial_values[v.name].size or 1) for v in vars]) self.any_discrete = self.discrete.any() self.all_discrete = self.discrete.all() # remember initial settings before tuning so they can be reset self._untuned_settings = dict(scaling=self.scaling, steps_until_tune=tune_interval, accepted=self.accepted) self.mode = mode shared = pm.make_shared_replacements(initial_values, vars, model) self.delta_logp = delta_logp(initial_values, model.logpt(), vars, shared) super().__init__(vars, shared)
def __init__( self, coarse_models: List[Model], vars: Optional[list] = None, base_sampler="DEMetropolisZ", base_S: Optional = None, base_proposal_dist: Optional[Type[Proposal]] = None, base_scaling: Optional = None, tune: bool = True, base_tune_target: str = "lambda", base_tune_interval: int = 100, base_lamb: Optional = None, base_tune_drop_fraction: float = 0.9, model: Optional[Model] = None, mode: Optional = None, subsampling_rates: List[int] = 5, base_blocked: bool = False, variance_reduction: bool = False, store_Q_fine: bool = False, adaptive_error_model: bool = False, **kwargs, ) -> None: # this variable is used to identify MLDA objects which are # not in the finest level (i.e. child MLDA objects) self.is_child = kwargs.get("is_child", False) if not isinstance(coarse_models, list): raise ValueError( "MLDA step method cannot use coarse_models if it is not a list" ) if len(coarse_models) == 0: raise ValueError("MLDA step method was given an empty " "list of coarse models. Give at least " "one coarse model.") # assign internal state model = pm.modelcontext(model) initial_values = model.compute_initial_point() self.model = model self.coarse_models = coarse_models self.model_below = self.coarse_models[-1] self.num_levels = len(self.coarse_models) + 1 # set up variance reduction. self.variance_reduction = variance_reduction self.store_Q_fine = store_Q_fine # check that certain requirements hold # for the variance reduction feature to work if self.variance_reduction or self.store_Q_fine: if not hasattr(self.model, "Q"): raise AttributeError("Model given to MLDA does not contain" "variable 'Q'. You need to include" "the variable in the model definition" "for variance reduction to work or" "for storing the fine Q." "Use pm.Data() to define it.") if not isinstance(self.model.Q, TensorSharedVariable): raise TypeError( "The variable 'Q' in the model definition is not of type " "'TensorSharedVariable'. Use pm.Data() to define the" "variable.") if self.is_child and self.variance_reduction: # this is the subsampling rate applied to the current level # it is stored in the level above and transferred here self.subsampling_rate_above = kwargs.pop("subsampling_rate_above", None) # set up adaptive error model self.adaptive_error_model = adaptive_error_model # check that certain requirements hold # for the adaptive error model feature to work if self.adaptive_error_model: if not hasattr(self.model_below, "mu_B"): raise AttributeError( "Model below in hierarchy does not contain" "variable 'mu_B'. You need to include" "the variable in the model definition" "for adaptive error model to work." "Use pm.Data() to define it.") if not hasattr(self.model_below, "Sigma_B"): raise AttributeError( "Model below in hierarchy does not contain" "variable 'Sigma_B'. You need to include" "the variable in the model definition" "for adaptive error model to work." "Use pm.Data() to define it.") if not (isinstance(self.model_below.mu_B, TensorSharedVariable) and isinstance(self.model_below.Sigma_B, TensorSharedVariable)): raise TypeError( "At least one of the variables 'mu_B' and 'Sigma_B' " "in the definition of the below model is not of type " "'TensorSharedVariable'. Use pm.Data() to define those " "variables.") # this object is used to recursively update the mean and # variance of the bias correction given new differences # between levels self.bias = RecursiveSampleMoments( self.model_below.mu_B.get_value(), self.model_below.Sigma_B.get_value()) # this list holds the bias objects from all levels # it is gradually constructed when MLDA objects are # created and then shared between all levels self.bias_all = kwargs.pop("bias_all", None) if self.bias_all is None: self.bias_all = [self.bias] else: self.bias_all.append(self.bias) # variables used for adaptive error model self.last_synced_output_diff = None self.adaptation_started = False # set up subsampling rates. if isinstance(subsampling_rates, int): self.subsampling_rates = [subsampling_rates] * len( self.coarse_models) else: if len(subsampling_rates) != len(self.coarse_models): raise ValueError( f"List of subsampling rates needs to have the same " f"length as list of coarse models but the lengths " f"were {len(subsampling_rates)}, {len(self.coarse_models)}" ) self.subsampling_rates = subsampling_rates self.subsampling_rate = self.subsampling_rates[-1] self.subchain_selection = None # set up base sampling self.base_sampler = base_sampler # VR is not compatible with compound base samplers so an automatic conversion # to a block sampler happens here if if self.variance_reduction and self.base_sampler == "Metropolis" and not base_blocked: warnings.warn( "Variance reduction is not compatible with non-blocked (compound) samplers." "Automatically switching to a blocked Metropolis sampler.") self.base_blocked = True else: self.base_blocked = base_blocked self.base_S = base_S self.base_proposal_dist = base_proposal_dist if base_scaling is None: if self.base_sampler == "Metropolis": self.base_scaling = 1.0 else: self.base_scaling = 0.001 else: self.base_scaling = float(base_scaling) self.tune = tune if not self.tune and self.base_sampler == "DEMetropolisZ": raise ValueError( f"The argument tune was set to False while using" f" a 'DEMetropolisZ' base sampler. 'DEMetropolisZ' " f" tune needs to be True.") self.base_tune_target = base_tune_target self.base_tune_interval = base_tune_interval self.base_lamb = base_lamb self.base_tune_drop_fraction = float(base_tune_drop_fraction) self.base_tuning_stats = None self.mode = mode # Process model variables 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) self.vars = vars self.var_names = [var.name for var in self.vars] self.accepted = 0 # Construct Aesara function for current-level model likelihood # (for use in acceptance) shared = pm.make_shared_replacements(initial_values, vars, model) self.delta_logp = delta_logp(initial_values, model.logpt(), vars, shared) # Construct Aesara function for below-level model likelihood # (for use in acceptance) model_below = pm.modelcontext(self.model_below) vars_below = [ var for var in model_below.value_vars if var.name in self.var_names ] vars_below = pm.inputvars(vars_below) shared_below = pm.make_shared_replacements(initial_values, vars_below, model_below) self.delta_logp_below = delta_logp(initial_values, model_below.logpt(), vars_below, shared_below) super().__init__(vars, shared) # initialise complete step method hierarchy if self.num_levels == 2: with self.model_below: # make sure the correct variables are selected from model_below vars_below = [ var for var in self.model_below.value_vars if var.name in self.var_names ] # create kwargs if self.variance_reduction: base_kwargs = { "mlda_subsampling_rate_above": self.subsampling_rate, "mlda_variance_reduction": True, } else: base_kwargs = {} if self.base_sampler == "Metropolis": # MetropolisMLDA sampler in base level (level=0), targeting self.model_below self.step_method_below = pm.MetropolisMLDA( vars=vars_below, proposal_dist=self.base_proposal_dist, S=self.base_S, scaling=self.base_scaling, tune=self.tune, tune_interval=self.base_tune_interval, model=None, mode=self.mode, blocked=self.base_blocked, **base_kwargs, ) else: # DEMetropolisZMLDA sampler in base level (level=0), targeting self.model_below self.step_method_below = pm.DEMetropolisZMLDA( vars=vars_below, S=self.base_S, proposal_dist=self.base_proposal_dist, lamb=self.base_lamb, scaling=self.base_scaling, tune=self.base_tune_target, tune_interval=self.base_tune_interval, tune_drop_fraction=self.base_tune_drop_fraction, model=None, mode=self.mode, **base_kwargs, ) else: # drop the last coarse model coarse_models_below = self.coarse_models[:-1] subsampling_rates_below = self.subsampling_rates[:-1] with self.model_below: # make sure the correct variables are selected from model_below vars_below = [ var for var in self.model_below.value_vars if var.name in self.var_names ] # create kwargs if self.variance_reduction: mlda_kwargs = { "is_child": True, "subsampling_rate_above": self.subsampling_rate, } else: mlda_kwargs = {"is_child": True} if self.adaptive_error_model: mlda_kwargs = { **mlda_kwargs, **{ "bias_all": self.bias_all } } # MLDA sampler in some intermediate level, targeting self.model_below self.step_method_below = pm.MLDA( vars=vars_below, base_S=self.base_S, base_sampler=self.base_sampler, base_proposal_dist=self.base_proposal_dist, base_scaling=self.base_scaling, tune=self.tune, base_tune_target=self.base_tune_target, base_tune_interval=self.base_tune_interval, base_lamb=self.base_lamb, base_tune_drop_fraction=self.base_tune_drop_fraction, model=None, mode=self.mode, subsampling_rates=subsampling_rates_below, coarse_models=coarse_models_below, base_blocked=self.base_blocked, variance_reduction=self.variance_reduction, store_Q_fine=False, adaptive_error_model=self.adaptive_error_model, **mlda_kwargs, ) # instantiate the recursive DA proposal. # this is the main proposal used for # all levels (Recursive Delayed Acceptance) # (except for level 0 where the step method is MetropolisMLDA # or DEMetropolisZMLDA - not MLDA) self.proposal_dist = RecursiveDAProposal(self.step_method_below, self.model_below, self.tune, self.subsampling_rate) # set up data types of stats. if isinstance(self.step_method_below, MLDA): # get the stat types from the level below if that level is MLDA self.stats_dtypes = self.step_method_below.stats_dtypes else: # otherwise, set it up from scratch. self.stats_dtypes = [{ "accept": np.float64, "accepted": bool, "tune": bool }] if isinstance(self.step_method_below, MetropolisMLDA): self.stats_dtypes.append({"base_scaling": np.float64}) elif isinstance(self.step_method_below, DEMetropolisZMLDA): self.stats_dtypes.append({ "base_scaling": np.float64, "base_lambda": np.float64 }) elif isinstance(self.step_method_below, CompoundStep): for method in self.step_method_below.methods: if isinstance(method, MetropolisMLDA): self.stats_dtypes.append({"base_scaling": np.float64}) elif isinstance(method, DEMetropolisZMLDA): self.stats_dtypes.append({ "base_scaling": np.float64, "base_lambda": np.float64 }) # initialise necessary variables for doing variance reduction if self.variance_reduction: self.sub_counter = 0 self.Q_diff = [] if self.is_child: self.Q_reg = [np.nan] * self.subsampling_rate_above if self.num_levels == 2: self.Q_base_full = [] if not self.is_child: for level in range(self.num_levels - 1, 0, -1): self.stats_dtypes[0][f"Q_{level}_{level - 1}"] = object self.stats_dtypes[0]["Q_0"] = object # initialise necessary variables for doing variance reduction or storing fine Q if self.variance_reduction or self.store_Q_fine: self.Q_last = np.nan self.Q_diff_last = np.nan if self.store_Q_fine and not self.is_child: self.stats_dtypes[0][f"Q_{self.num_levels - 1}"] = object
def __init__(self, vars=None, S=None, proposal_dist=None, scaling=1.0, tune=True, tune_interval=100, model=None, mode=None, **kwargs): """Create an instance of a Metropolis stepper Parameters ---------- vars: list List of value variables for sampler S: standard deviation or covariance matrix Some measure of variance to parameterize proposal distribution proposal_dist: function Function that returns zero-mean deviates when parameterized with S (and n). Defaults to normal. scaling: scalar or array Initial scale factor for proposal. Defaults to 1. tune: bool Flag for tuning. Defaults to True. tune_interval: int The frequency of tuning. Defaults to 100 iterations. model: PyMC Model Optional model for sampling step. Defaults to None (taken from context). mode: string or `Mode` instance. compilation mode passed to Aesara functions """ model = pm.modelcontext(model) initial_values = model.initial_point() 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) initial_values_shape = [initial_values[v.name].shape for v in vars] if S is None: S = np.ones(int(sum(np.prod(ivs) for ivs in initial_values_shape))) if proposal_dist is not None: self.proposal_dist = proposal_dist(S) elif S.ndim == 1: self.proposal_dist = NormalProposal(S) elif S.ndim == 2: self.proposal_dist = MultivariateNormalProposal(S) else: raise ValueError("Invalid rank for variance: %s" % S.ndim) self.scaling = np.atleast_1d(scaling).astype("d") self.tune = tune self.tune_interval = tune_interval self.steps_until_tune = tune_interval # Determine type of variables self.discrete = np.concatenate([[v.dtype in pm.discrete_types] * (initial_values[v.name].size or 1) for v in vars]) self.any_discrete = self.discrete.any() self.all_discrete = self.discrete.all() # Metropolis will try to handle one batched dimension at a time This, however, # is not safe for discrete multivariate distributions (looking at you Multinomial), # due to high dependency among the support dimensions. For continuous multivariate # distributions we assume they are being transformed in a way that makes each # dimension semi-independent. is_scalar = len( initial_values_shape) == 1 and initial_values_shape[0] == () self.elemwise_update = not (is_scalar or (self.any_discrete and max( getattr(model.values_to_rvs[var].owner.op, "ndim_supp", 1) for var in vars) > 0)) if self.elemwise_update: dims = int(sum(np.prod(ivs) for ivs in initial_values_shape)) else: dims = 1 self.enum_dims = np.arange(dims, dtype=int) self.accept_rate_iter = np.zeros(dims, dtype=float) self.accepted_iter = np.zeros(dims, dtype=bool) self.accepted_sum = np.zeros(dims, dtype=int) # remember initial settings before tuning so they can be reset self._untuned_settings = dict(scaling=self.scaling, steps_until_tune=tune_interval) # TODO: This is not being used when compiling the logp function! self.mode = mode shared = pm.make_shared_replacements(initial_values, vars, model) self.delta_logp = delta_logp(initial_values, model.logp(), vars, shared) super().__init__(vars, shared)
def __init__(self, vars=None, S=None, proposal_dist=None, lamb=None, scaling=0.001, tune="lambda", tune_interval=100, tune_drop_fraction: float = 0.9, model=None, mode=None, **kwargs): model = pm.modelcontext(model) initial_values = model.recompute_initial_point() initial_values_size = sum(initial_values[n.name].size for n in model.value_vars) 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: S = np.ones(initial_values_size) if proposal_dist is not None: self.proposal_dist = proposal_dist(S) else: self.proposal_dist = UniformProposal(S) self.scaling = np.atleast_1d(scaling).astype("d") if lamb is None: # default to the optimal lambda for normally distributed targets lamb = 2.38 / np.sqrt(2 * initial_values_size) self.lamb = float(lamb) if tune not in {None, "scaling", "lambda"}: raise ValueError( 'The parameter "tune" must be one of {None, scaling, lambda}') self.tune = True self.tune_target = tune self.tune_interval = tune_interval self.tune_drop_fraction = tune_drop_fraction self.steps_until_tune = tune_interval self.accepted = 0 # cache local history for the Z-proposals self._history = [] # remember initial settings before tuning so they can be reset self._untuned_settings = dict( scaling=self.scaling, lamb=self.lamb, steps_until_tune=tune_interval, accepted=self.accepted, ) self.mode = mode shared = pm.make_shared_replacements(initial_values, vars, model) self.delta_logp = delta_logp(initial_values, model.logpt, vars, shared) super().__init__(vars, shared)