def __init__(self, vars, proposal="uniform", order="random", model=None):
model = pm.modelcontext(model)
vars = [model.rvs_to_values.get(var, var) for var in vars]
vars = pm.inputvars(vars)
initial_point = model.initial_point()
dimcats = []
# The above variable is a list of pairs (aggregate dimension, number
# of categories). For example, if vars = [x, y] with x being a 2-D
# variable with M categories and y being a 3-D variable with N
# categories, we will have dimcats = [(0, M), (1, M), (2, N), (3, N), (4, N)].
for v in vars:
v_init_val = initial_point[v.name]
rv_var = model.values_to_rvs[v]
distr = getattr(rv_var.owner, "op", None)
if isinstance(distr, CategoricalRV):
k_graph = rv_var.owner.inputs[3].shape[-1]
(k_graph, ), _ = rvs_to_value_vars((k_graph, ),
apply_transforms=True)
k = model.compile_fn(k_graph,
inputs=model.value_vars,
on_unused_input="ignore")(initial_point)
elif isinstance(distr, BernoulliRV):
k = 2
else:
raise ValueError(
"All variables must be categorical or binary" +
"for CategoricalGibbsMetropolis")
start = len(dimcats)
dimcats += [(dim, k)
for dim in range(start, start + v_init_val.size)]
if order == "random":
self.shuffle_dims = True
self.dimcats = dimcats
else:
if sorted(order) != list(range(len(dimcats))):
raise ValueError("Argument 'order' has to be a permutation")
self.shuffle_dims = False
self.dimcats = [dimcats[j] for j in order]
if proposal == "uniform":
self.astep = self.astep_unif
elif proposal == "proportional":
# Use the optimized "Metropolized Gibbs Sampler" described in Liu96.
self.astep = self.astep_prop
else:
raise ValueError(
"Argument 'proposal' should either be 'uniform' or 'proportional'"
)
super().__init__(vars, [model.compile_logp()])
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)
