Skip to content

kups.application.mcmc.simulation

Generic simulation loop for rigid-body MCMC simulations.

run_mcmc(key, propagator, state, config, logged_data)

Run a µVT MCMC simulation with warmup and production phases.

Parameters:

Name Type Description Default
key Array

JAX PRNG key.

 required
propagator Propagator[State]

Propagator, e.g. from :func:~kups.application.simulations.mcmc_rigid.make_propagator.

 required
state State

Initial simulation state.

 required
config RunConfig

Run configuration.

 required
logged_data MCMCLoggedData

Logging configuration with host/adsorbate split.

 required

Returns:

Type Description
State

Final simulation state after production run.
Source code in src/kups/application/mcmc/simulation.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def run_mcmc[State](
    key: Array,
    propagator: Propagator[State],
    state: State,
    config: RunConfig,
    logged_data: MCMCLoggedData,
) -> State:
    """Run a µVT MCMC simulation with warmup and production phases.

    Args:
        key: JAX PRNG key.
        propagator: Propagator, e.g. from :func:`~kups.application.simulations.mcmc_rigid.make_propagator`.
        state: Initial simulation state.
        config: Run configuration.
        logged_data: Logging configuration with host/adsorbate split.

    Returns:
        Final simulation state after production run.
    """
    chain = key_chain(key)
    logging.info("Warming up (%d cycles)...", config.num_warmup_cycles)
    state = run_warmup_cycles(next(chain), propagator, state, config.num_warmup_cycles)
    logging.info("Production run (%d cycles)...", config.num_cycles)
    logger = CompositeLogger(
        HDF5StorageWriter(config.out_file, logged_data, state, config.num_cycles),
        TqdmLogger(config.num_cycles),
    )
    state = run_simulation_cycles(
        next(chain), propagator, state, config.num_cycles, logger
    )
    logging.info("Done.")
    return state