def startup(self, scenario: Scenario, n: int, initial_state: cdict,
initial_extra: cdict, **kwargs) -> Tuple[cdict, cdict]:
initial_state, initial_extra = super().startup(scenario, n,
initial_state,
initial_extra, **kwargs)
if self.parameters.ensemble_batchsize is None:
self.parameters.ensemble_batchsize = n
initial_extra.parameters.ensemble_batchsize = n
if self.parameters.ensemble_batchsize == n:
self.get_batch_inds = lambda _: jnp.repeat(
jnp.arange(n)[None], n, axis=0)
else:
self.get_batch_inds = lambda rk: random.choice(
rk, n, shape=(
n,
self.parameters.ensemble_batchsize,
))
del initial_extra.parameters.stepsize
random_keys = random.split(initial_extra.random_key, n + 1)
initial_extra.random_key = random_keys[-1]
initial_state.potential, initial_state.grad_potential = vmap(
scenario.potential_and_grad)(initial_state.value, random_keys[:n])
initial_state, initial_extra = self.adapt(initial_state, initial_extra)
self.opt_init, self.opt_update, self.get_params = self.optimiser(
step_size=self.parameters.stepsize,
**initial_extra.parameters.optim_params)
initial_extra.opt_state = self.opt_init(initial_state.value)
return initial_state, initial_extra
def update(self, scenario: Scenario, ensemble_state: cdict,
extra: cdict) -> Tuple[cdict, cdict]:
n = ensemble_state.value.shape[0]
extra.iter = extra.iter + 1
random_keys = random.split(extra.random_key, n + 2)
batch_inds = self.get_batch_inds(random_keys[-1])
extra.random_key = random_keys[-2]
phi_hat = self.kernelised_grad_matrix(ensemble_state.value,
ensemble_state.grad_potential,
extra.parameters.kernel_params,
batch_inds)
extra.opt_state = self.opt_update(extra.iter, -phi_hat,
extra.opt_state)
ensemble_state.value = self.get_params(extra.opt_state)
ensemble_state.potential, ensemble_state.grad_potential \
= vmap(scenario.potential_and_grad)(ensemble_state.value, random_keys[:n])
ensemble_state, extra = self.adapt(ensemble_state, extra)
return ensemble_state, extra
