def astep(self,
q0: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
point_map_info = q0.point_map_info
q0 = q0.data
if not self.steps_until_tune and self.tune:
# Tune scaling parameter
self.scaling = tune(self.scaling,
self.accepted_sum / float(self.tune_interval))
# Reset counter
self.steps_until_tune = self.tune_interval
self.accepted_sum[:] = 0
delta = self.proposal_dist() * self.scaling
if self.any_discrete:
if self.all_discrete:
delta = np.round(delta, 0).astype("int64")
q0 = q0.astype("int64")
q = (q0 + delta).astype("int64")
else:
delta[self.discrete] = np.round(delta[self.discrete], 0)
q = q0 + delta
else:
q = floatX(q0 + delta)
if self.elemwise_update:
q_temp = q0.copy()
# Shuffle order of updates (probably we don't need to do this in every step)
np.random.shuffle(self.enum_dims)
for i in self.enum_dims:
q_temp[i] = q[i]
accept_rate_i = self.delta_logp(q_temp, q0)
q_temp_, accepted_i = metrop_select(accept_rate_i, q_temp, q0)
q_temp[i] = q_temp_[i]
self.accept_rate_iter[i] = accept_rate_i
self.accepted_iter[i] = accepted_i
self.accepted_sum[i] += accepted_i
q = q_temp
else:
accept_rate = self.delta_logp(q, q0)
q, accepted = metrop_select(accept_rate, q, q0)
self.accept_rate_iter = accept_rate
self.accepted_iter = accepted
self.accepted_sum += accepted
self.steps_until_tune -= 1
stats = {
"tune": self.tune,
"scaling": np.mean(self.scaling),
"accept": np.mean(np.exp(self.accept_rate_iter)),
"accepted": np.mean(self.accepted_iter),
}
return RaveledVars(q, point_map_info), [stats]
def astep(self, q0: RaveledVars,
logp) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
logp_q0 = logp(q0)
point_map_info = q0.point_map_info
q0 = q0.data
# Convert adaptive_scale_factor to a jump probability
p_jump = 1.0 - 0.5**self.scaling
rand_array = nr.random(q0.shape)
q = np.copy(q0)
# Locations where switches occur, according to p_jump
switch_locs = rand_array < p_jump
q[switch_locs] = True - q[switch_locs]
logp_q = logp(RaveledVars(q, point_map_info))
accept = logp_q - logp_q0
q_new, accepted = metrop_select(accept, q, q0)
self.accepted += accepted
stats = {
"tune": self.tune,
"accept": np.exp(accept),
"p_jump": p_jump,
}
q_new = RaveledVars(q_new, point_map_info)
return q_new, [stats]
def astep(self, q0, logp):
# sample from kernel based Gaussian proposal
q_dist = self.construct_proposal(q0)
q = np.ravel(q_dist.sample())
# evaluate target log probability
logp_q = logp(q)
# MH accept/reject step
if self.q_dist is None:
q_new = q
else:
q_new = metrop_select(logp_q + q_dist.log_pdf(q0) \
- self.log_pdf_target - self.q_dist.log_pdf(q), q, q0)
# adapt
if self.tune and not self.steps_until_tune:
# tune scaling parameter using metropolis method
self.nu2 = tune(self.nu2, self.accepted / float(self.tune_interval))
# Reset counter
self.steps_until_tune = self.tune_interval
self.accepted = 0
# update log-pdf and proposal distribution object on accept
if any(q_new != q0):
self.accepted += 1
self.q_dist = q_dist
self.log_pdf_target = logp_q
self.steps_until_tune -= 1
return q_new
def astep(self,
q0: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
point_map_info = q0.point_map_info
q0 = q0.data
# same tuning scheme as DEMetropolis
if not self.steps_until_tune and self.tune:
if self.tune_target == "scaling":
self.scaling = tune(self.scaling,
self.accepted / float(self.tune_interval))
elif self.tune_target == "lambda":
self.lamb = tune(self.lamb,
self.accepted / float(self.tune_interval))
# Reset counter
self.steps_until_tune = self.tune_interval
self.accepted = 0
epsilon = self.proposal_dist() * self.scaling
it = len(self._history)
# use the DE-MCMC-Z proposal scheme as soon as the history has 2 entries
if it > 1:
# differential evolution proposal
# select two other chains
iz1 = np.random.randint(it)
iz2 = np.random.randint(it)
while iz2 == iz1:
iz2 = np.random.randint(it)
z1 = self._history[iz1]
z2 = self._history[iz2]
# propose a jump
q = floatX(q0 + self.lamb * (z1 - z2) + epsilon)
else:
# propose just with noise in the first 2 iterations
q = floatX(q0 + epsilon)
accept = self.delta_logp(q, q0)
q_new, accepted = metrop_select(accept, q, q0)
self.accepted += accepted
self._history.append(q_new)
self.steps_until_tune -= 1
stats = {
"tune": self.tune,
"scaling": self.scaling,
"lambda": self.lamb,
"accept": np.exp(accept),
"accepted": accepted,
}
q_new = RaveledVars(q_new, point_map_info)
return q_new, [stats]
def astep(self,
q0: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
point_map_info = q0.point_map_info
q0 = q0.data
if not self.steps_until_tune and self.tune:
if self.tune == "scaling":
self.scaling = tune(self.scaling,
self.accepted / float(self.tune_interval))
elif self.tune == "lambda":
self.lamb = tune(self.lamb,
self.accepted / float(self.tune_interval))
# Reset counter
self.steps_until_tune = self.tune_interval
self.accepted = 0
epsilon = self.proposal_dist() * self.scaling
# differential evolution proposal
# select two other chains
ir1, ir2 = np.random.choice(self.other_chains, 2, replace=False)
r1 = DictToArrayBijection.map(self.population[ir1])
r2 = DictToArrayBijection.map(self.population[ir2])
# propose a jump
q = floatX(q0 + self.lamb * (r1.data - r2.data) + epsilon)
accept = self.delta_logp(q, q0)
q_new, accepted = metrop_select(accept, q, q0)
self.accepted += accepted
self.steps_until_tune -= 1
stats = {
"tune": self.tune,
"scaling": self.scaling,
"lambda": self.lamb,
"accept": np.exp(accept),
"accepted": accepted,
}
q_new = RaveledVars(q_new, point_map_info)
return q_new, [stats]
def astep(self,
q0: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
point_map_info = q0.point_map_info
q0 = q0.data
if not self.steps_until_tune and self.tune:
# Tune scaling parameter
self.scaling = tune(self.scaling,
self.accepted / float(self.tune_interval))
# Reset counter
self.steps_until_tune = self.tune_interval
self.accepted = 0
delta = self.proposal_dist() * self.scaling
if self.any_discrete:
if self.all_discrete:
delta = np.round(delta, 0).astype("int64")
q0 = q0.astype("int64")
q = (q0 + delta).astype("int64")
else:
delta[self.discrete] = np.round(delta[self.discrete], 0)
q = q0 + delta
else:
q = floatX(q0 + delta)
accept = self.delta_logp(q, q0)
q_new, accepted = metrop_select(accept, q, q0)
self.accepted += accepted
self.steps_until_tune -= 1
stats = {
"tune": self.tune,
"scaling": self.scaling,
"accept": np.exp(accept),
"accepted": accepted,
}
q_new = RaveledVars(q_new, point_map_info)
return q_new, [stats]
def astep_unif(self, q0: RaveledVars, logp) -> RaveledVars:
point_map_info = q0.point_map_info
q0 = q0.data
dimcats = self.dimcats
if self.shuffle_dims:
nr.shuffle(dimcats)
q = RaveledVars(np.copy(q0), point_map_info)
logp_curr = logp(q)
for dim, k in dimcats:
curr_val, q.data[dim] = q.data[dim], sample_except(k, q.data[dim])
logp_prop = logp(q)
q.data[dim], accepted = metrop_select(logp_prop - logp_curr, q.data[dim], curr_val)
if accepted:
logp_curr = logp_prop
return q
def astep(self, q0: RaveledVars, logp: Callable[[RaveledVars], np.ndarray]) -> RaveledVars:
order = self.order
if self.shuffle_dims:
nr.shuffle(order)
q = RaveledVars(np.copy(q0.data), q0.point_map_info)
logp_curr = logp(q)
for idx in order:
# No need to do metropolis update if the same value is proposed,
# as you will get the same value regardless of accepted or reject
if nr.rand() < self.transit_p:
curr_val, q.data[idx] = q.data[idx], True - q.data[idx]
logp_prop = logp(q)
q.data[idx], accepted = metrop_select(logp_prop - logp_curr, q.data[idx], curr_val)
if accepted:
logp_curr = logp_prop
return q
def astep(self,
q0: RaveledVars) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
"""One MLDA step, given current sample q0"""
# Check if the tuning flag has been changed and if yes,
# change the proposal's tuning flag and reset self.accepted
# This is triggered by _iter_sample while the highest-level MLDA step
# method is running. It then propagates to all levels.
if self.proposal_dist.tune != self.tune:
self.proposal_dist.tune = self.tune
# set tune in sub-methods of compound stepper explicitly because
# it is not set within sample.py (only the CompoundStep's tune flag is)
if isinstance(self.step_method_below, CompoundStep):
for method in self.step_method_below.methods:
method.tune = self.tune
self.accepted = 0
# Set subchain_selection (which sample from the coarse chain
# is passed as a proposal to the fine chain). If variance
# reduction is used, a random sample is selected as proposal.
# If variance reduction is not used, the last sample is
# selected as proposal.
if self.variance_reduction:
self.subchain_selection = np.random.randint(
0, self.subsampling_rate)
else:
self.subchain_selection = self.subsampling_rate - 1
self.proposal_dist.subchain_selection = self.subchain_selection
# Call the recursive DA proposal to get proposed sample
# and convert dict -> numpy array
q = self.proposal_dist(q0)
# Evaluate MLDA acceptance log-ratio
# If proposed sample from lower levels is the same as current one,
# do not calculate likelihood, just set accept to 0.0
if (q.data == q0.data).all():
accept = np.float64(0.0)
skipped_logp = True
else:
# NB! The order and sign of the first term are swapped compared
# to the convention to make sure the proposal is evaluated last.
accept = -self.delta_logp(q0.data, q.data) + self.delta_logp_below(
q0.data, q.data)
skipped_logp = False
# Accept/reject sample - next sample is stored in q_new
q_new, accepted = metrop_select(accept, q, q0)
if skipped_logp:
accepted = False
# if sample is accepted, update self.Q_last with the sample's Q value
# runs only for VR or when store_Q_fine is True
if self.variance_reduction or self.store_Q_fine:
if accepted and not skipped_logp:
self.Q_last = self.model.Q.get_value()
# Variance reduction
if self.variance_reduction:
self.update_vr_variables(accepted, skipped_logp)
# Adaptive error model - runs only during tuning.
if self.tune and self.adaptive_error_model:
self.update_error_estimate(accepted, skipped_logp)
# Update acceptance counter
self.accepted += accepted
stats = {
"tune": self.tune,
"accept": np.exp(accept),
"accepted": accepted
}
# Save the VR statistics to the stats dictionary (only happens in the
# top MLDA level)
if (self.variance_reduction
or self.store_Q_fine) and not self.is_child:
q_stats = {}
if self.variance_reduction:
m = self
for level in range(self.num_levels - 1, 0, -1):
# save the Q differences for this level and iteration
q_stats[f"Q_{level}_{level - 1}"] = np.array(m.Q_diff)
# this makes sure Q_diff is reset for
# the next iteration
m.Q_diff = []
if level == 1:
break
m = m.step_method_below
q_stats["Q_0"] = np.array(m.Q_base_full)
m.Q_base_full = []
if self.store_Q_fine:
q_stats["Q_" + str(self.num_levels - 1)] = np.array(
self.Q_last)
stats = {**stats, **q_stats}
# Capture the base tuning stats from the level below.
self.base_tuning_stats = []
if isinstance(self.step_method_below, MLDA):
self.base_tuning_stats = self.step_method_below.base_tuning_stats
elif isinstance(self.step_method_below, MetropolisMLDA):
self.base_tuning_stats.append(
{"base_scaling": self.step_method_below.scaling})
elif isinstance(self.step_method_below, DEMetropolisZMLDA):
self.base_tuning_stats.append({
"base_scaling":
self.step_method_below.scaling,
"base_lambda":
self.step_method_below.lamb,
})
elif isinstance(self.step_method_below, CompoundStep):
# Below method is CompoundStep
for method in self.step_method_below.methods:
if isinstance(method, MetropolisMLDA):
self.base_tuning_stats.append(
{"base_scaling": method.scaling})
elif isinstance(method, DEMetropolisZMLDA):
self.base_tuning_stats.append({
"base_scaling": method.scaling,
"base_lambda": method.lamb
})
return q_new, [stats] + self.base_tuning_stats
