jax_sgmc.scheduler

Schedules parameters of the integrator and solver.

The scheduler organizes the independent variables of the update equation, such as the temperature and the step size, which are organized by multiple specific schedulers.

Base Scheduler

jax_sgmc.scheduler.init_scheduler(step_size=None, temperature=None, burn_in=None, thinning=None, progress_bar=True, progress_bar_steps=20)[source]

Initializes the scheduler.

The scheduler combines the specific schedules for each variable. It updates them and gets them at a central place and makes it possible to combine them or provide default values.

Parameters:

step_size (Optional[specific_scheduler]) – Triplet from step-size scheduler initialization
temperature (Optional[specific_scheduler]) – Triplet from temperature scheduler initialization
burn_in (Optional[specific_scheduler]) – Triplet from burn-in scheduler initialization
thinning (Optional[specific_scheduler]) – Triplet from thinning scheduler initialization
progress_bar (bool) – Show the percentage of completed steps

Return type:

Tuple[Callable, Callable, Callable]

Returns:

Returns a triplet of (init_fn, update_fn, get_fn).

jax_sgmc.scheduler.scheduler_state(state, step_size_state, temperature_state, burn_in_state, thinning_state, progress_bar_state)

Collects the states of the specific schedulers.

jax_sgmc.scheduler.state: State of the base scheduler, e.g. to keep track of current iteration

jax_sgmc.scheduler.step_size_state: State of the step size scheduler

jax_sgmc.scheduler.temperature_state: State of the temperature scheduler

jax_sgmc.scheduler.burn_in_state: State of the burn in scheduler

jax_sgmc.scheduler.thinning_state: State of thinning

jax_sgmc.scheduler.progress_bar_state: State of the progress bar

jax_sgmc.scheduler.schedule(step_size, temperature, burn_in, accept)

Auxiliary variables for integrator.

jax_sgmc.scheduler.step_size: Learning rate

jax_sgmc.scheduler.temperature: Scaling the magnitude of the additional noise

jax_sgmc.scheduler.burn_in: Bool, whether current step can be accepted

jax_sgmc.scheduler.accept: Bool, whether current sample should be saved

jax_sgmc.scheduler.static_information(samples_collected)

Information which is constant during the run.

jax_sgmc.scheduler.samples_collected: Number of samples saved.

Specific Schedulers

class jax_sgmc.scheduler.specific_scheduler(init, update, get)

Bundles the specific scheduler as described above.

init_fn: Function to initialize the specific scheduler

update_fn: Calculate the next state of the scheduler

get_fn: Get the scheduled parameter

Step-size

`polynomial_step_size`([a, b, gamma])	Polynomial decreasing step size schedule.
`polynomial_step_size_first_last`([first, ...])	Initializes polynomial step size schedule via first and last step.
`adaptive_step_size`([burn_in, ...])	Dual averaging scheme to tune step size for schemes with MH-step.

Temperature

`constant_temperature`([tau])	Scales the added noise with an unchanged constant.
`cyclic_temperature`([beta, k])	Cyclic switch of the temperature between 0.0 and 1.0.

Burn In

`cyclic_burn_in`([beta, k])	Discards samples at the beginning of each cycle.
`initial_burn_in`([n])	Discards the first n steps.

Thinning

random_thinning(step_size_schedule, ...[, key])

Random thinning weighted by the step size.