ensemble.sampling.init_simulation_fn

ensemble.sampling.init_simulation_fn#

init_simulation_fn(run_to_printout_fn, timings, vmap_batch_size=1, devices=None)[source]#

Runs a simulation with frequently saved states.

Parameters:

  • run_to_printout_fn – Function to run simulation to next printout state

  • timings (TimingClass) – Instance of TimingClass containing information about which states to retain and simulation time.

  • vmap_batch_size (int) – If multiple simulation states provided, run multiple trajectories vectorized via vmap.

  • devices – If multiple devices provided, run simulations in parallel via shmap.

Returns:

Returns the final state after the simulation and the subsampled simulator states at the defined printout times.

Contents