def test_run(self): model, coarse_models = self.build_models() with model: step = pm.MLDA(subsampling_rates=2, coarse_models=coarse_models) pm.sample(draws=50, chains=2, tune=50, step=step) step = pm.MLDA( subsampling_rates=2, coarse_models=coarse_models, base_sampler="Metropolis" ) pm.sample(draws=50, chains=2, tune=50, step=step)
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 self.is_child: warnings.warn( "The MLDA implementation in PyMC3 is still immature. You should be particularly critical of its results." ) 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) 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, tt.sharedvar.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, tt.sharedvar.TensorSharedVariable) and isinstance(self.model_below.Sigma_B, tt.sharedvar.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.vars vars = pm.inputvars(vars) self.vars = vars self.var_names = [var.name for var in self.vars] self.accepted = 0 # Construct theano function for current-level model likelihood # (for use in acceptance) shared = pm.make_shared_replacements(vars, model) self.delta_logp = delta_logp_inverse(model.logpt, vars, shared) # Construct theano 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.vars if var.name in self.var_names ] vars_below = pm.inputvars(vars_below) shared_below = pm.make_shared_replacements(vars_below, model_below) self.delta_logp_below = delta_logp(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.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.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": np.bool, "tune": np.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, coarse_models: List[Model], vars: Optional[list] = None, base_S: Optional = None, base_proposal_dist: Optional[Type[Proposal]] = None, base_scaling: Union[float, int] = 1.0, tune: bool = True, base_tune_interval: int = 100, model: Optional[Model] = None, mode: Optional = None, subsampling_rates: List[int] = 5, base_blocked: bool = False, **kwargs) -> None: warnings.warn("The MLDA implementation in PyMC3 is very young. " "You should be extra critical about its results.") model = pm.modelcontext(model) # assign internal state self.coarse_models = coarse_models if not isinstance(coarse_models, list): raise ValueError( "MLDA step method cannot use coarse_models if it is not a list" ) if len(self.coarse_models) == 0: raise ValueError("MLDA step method was given an empty " "list of coarse models. Give at least " "one coarse model.") 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.num_levels = len(self.coarse_models) + 1 self.base_S = base_S self.base_proposal_dist = base_proposal_dist self.base_scaling = base_scaling self.tune = tune self.base_tune_interval = base_tune_interval self.model = model self.next_model = self.coarse_models[-1] self.mode = mode self.base_blocked = base_blocked self.base_scaling_stats = None # Process model variables if vars is None: vars = model.vars vars = pm.inputvars(vars) self.vars = vars self.var_names = [var.name for var in self.vars] self.accepted = 0 # Construct theano function for current-level model likelihood # (for use in acceptance) shared = pm.make_shared_replacements(vars, model) self.delta_logp = delta_logp(model.logpt, vars, shared) # Construct theano function for next-level model likelihood # (for use in acceptance) next_model = pm.modelcontext(self.next_model) vars_next = [ var for var in next_model.vars if var.name in self.var_names ] vars_next = pm.inputvars(vars_next) shared_next = pm.make_shared_replacements(vars_next, next_model) self.delta_logp_next = delta_logp(next_model.logpt, vars_next, shared_next) super().__init__(vars, shared) # initialise complete step method hierarchy if self.num_levels == 2: with self.next_model: # make sure the correct variables are selected from next_model vars_next = [ var for var in self.next_model.vars if var.name in self.var_names ] # MetropolisMLDA sampler in base level (level=0), targeting self.next_model self.next_step_method = pm.MetropolisMLDA( vars=vars_next, 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, blocked=self.base_blocked, ) else: # drop the last coarse model next_coarse_models = self.coarse_models[:-1] next_subsampling_rates = self.subsampling_rates[:-1] with self.next_model: # make sure the correct variables are selected from next_model vars_next = [ var for var in self.next_model.vars if var.name in self.var_names ] # MLDA sampler in some intermediate level, targeting self.next_model self.next_step_method = pm.MLDA( vars=vars_next, base_S=self.base_S, base_proposal_dist=self.base_proposal_dist, base_scaling=self.base_scaling, tune=self.tune, base_tune_interval=self.base_tune_interval, model=None, mode=self.mode, subsampling_rates=next_subsampling_rates, coarse_models=next_coarse_models, base_blocked=self.base_blocked, **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 and not MLDA) self.proposal_dist = RecursiveDAProposal( self.next_step_method, self.next_model, self.tune, self.subsampling_rates[-1], )