def _pick_random_and_init(batch_size, move_cumulative, out, sections):
if out.size != batch_size:
out = _np.empty(batch_size, dtype=out.dtype)
sections = _np.empty(batch_size, dtype=out.dtype)
for i in range(batch_size):
p = _random.uniform()
out[i] = _np.searchsorted(move_cumulative, p)
return out, sections
def _choose_and_return(state_1, x_prime, mels, sections, log_prob_corr):
low = 0
for i in range(state_1.shape[0]):
p = _random.uniform()
exit_state = 0
cumulative_prob = mels[low].real
while p > cumulative_prob:
exit_state += 1
cumulative_prob += mels[low + exit_state].real
state_1[i] = x_prime[low + exit_state]
low = sections[i]
log_prob_corr.fill(0.0)
def acceptance_kernel(state, state1, log_values, log_values_1,
log_prob_corr, machine_pow):
accepted = 0
for i in range(state.shape[0]):
prob = _np.exp(
machine_pow *
(log_values_1[i] - log_values[i] + log_prob_corr[i]).real)
assert not math.isnan(prob)
if prob > _random.uniform(0, 1):
log_values[i] = log_values_1[i]
state[i] = state1[i]
accepted += 1
return accepted
def _exchange_step_kernel(log_values, machine_pow, beta, proposed_beta,
prob, beta_stats, accepted_samples):
# Choose a random swap order (odd/even swap)
swap_order = _random.randint(0, 2, size=()).item()
n_replicas = beta.shape[0]
for i in range(swap_order, n_replicas, 2):
inn = (i + 1) % n_replicas
proposed_beta[i] = beta[inn]
proposed_beta[inn] = beta[i]
for i in range(n_replicas):
prob[i] = math.exp(machine_pow * (proposed_beta[i] - beta[i]) *
log_values[i].real)
for i in range(swap_order, n_replicas, 2):
inn = (i + 1) % n_replicas
prob[i] *= prob[inn]
if prob[i] > _random.uniform(0, 1):
# swapping status
beta[i], beta[inn] = beta[inn], beta[i]
accepted_samples[i], accepted_samples[inn] = (
accepted_samples[inn],
accepted_samples[i],
)
if beta_stats[0] == i:
beta_stats[0] = inn
elif beta_stats[0] == inn:
beta_stats[0] = i
# Update statistics to compute diffusion coefficient of replicas
# Total exchange steps performed
beta_stats[-1] += 1
delta = beta_stats[0] - beta_stats[1]
beta_stats[1] += delta / float(beta_stats[-1])
delta2 = beta_stats[0] - beta_stats[1]
beta_stats[2] += delta * delta2
def _fixed_beta_acceptance_kernel(
state,
state1,
log_values,
log_values_1,
log_prob_corr,
machine_pow,
beta,
accepted,
):
for i in range(state.shape[0]):
prob = math.exp(
machine_pow * beta[i] *
(log_values_1[i] - log_values[i] + log_prob_corr[i]).real)
assert not math.isnan(prob)
accept = prob > _random.uniform(0, 1)
if accept:
log_values[i] = log_values_1[i]
state[i] = state1[i]
accepted[i] += 1
def _pick_random_and_init(batch_size, move_cumulative, out):
for i in range(batch_size):
p = _random.uniform()
out[i] = np.searchsorted(move_cumulative, p)