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[State]

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 
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def run_mcmc[State: IsMCMCState](
    key: Array,
    propagator: Propagator[State],
    state: State,
    config: RunConfig,
    logged_data: MCMCLoggedData[State],
) -> 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.
    """

    def postfix(sample: MCMCStepData) -> dict[str, str]:
        return {"Loading": str(sample.particle_count.data.sum())}

    chain = key_chain(key)
    cycle_fn = make_sampled_cycle_function(propagator, logged_data.per_step.view)
    logging.info("Warming up (%d cycles)...", config.num_warmup_cycles)
    state = run_sampled_cycles(
        next(chain),
        cycle_fn,
        state,
        config.num_warmup_cycles,
        TqdmLogger(config.num_warmup_cycles),
        config.cycles_per_call,
    )
    logging.info("Production run (%d cycles)...", config.num_cycles)
    # The fixed snapshot must outlive donation of the simulation buffers.
    fixed = jax.tree.map(
        lambda array: array.copy(), jax.device_get(logged_data.fixed.view(state))
    )
    sample_config = MCMCLoggedData[MCMCStepData](
        fixed=WriterGroupConfig(
            lambda _: fixed,
            logged_data.fixed.logging_frequency,
            compression=logged_data.fixed.compression,
        ),
        per_step=WriterGroupConfig(
            lambda sample: sample,
            logged_data.per_step.logging_frequency,
            compression=logged_data.per_step.compression,
        ),
    )
    logger = CompositeLogger[MCMCStepData](
        HDF5StorageWriter(
            config.out_file,
            sample_config,
            logged_data.per_step.view(state),
            config.num_cycles,
            compile_views=False,
        ),
        TqdmLogger(config.num_cycles, postfix=postfix),
    )
    state = run_sampled_cycles(
        next(chain),
        cycle_fn,
        state,
        config.num_cycles,
        logger,
        config.cycles_per_call,
    )
    logging.info("Done.")
    return state