ensemble.sampling.trajectory_generator_init

ensemble.sampling.trajectory_generator_init#

trajectory_generator_init(simulator_template, energy_fn_template, ref_timings=None, quantities=None, vmap_batch=10, vmap_sim_batch=10)[source]#

Initializes a function to computes a trajectory.

Parameters:

  • simulator_template – Function returning new simulator given current energy function

  • energy_fn_template – Energy function template

  • ref_timings – Instance of TimingClass containing information about the times states need to be retained

  • quantities – Quantities dict to compute and store auxilary quantities alongside trajectory. This is particularly helpful for storing energy and pressure in a reweighting context.

  • vmap_batch – Batch size for computation of auxillary quantities.

The trajectory generation consists of the following steps:

  1. Evaluation of the dynamic kwargs at the specified simulation times. If the kwargs are constant, they are converted to constant functions.

  2. Running multiple short simulations, saving only the sim-states at the specified printout times.

  3. Computing auxilliary quantities for each of the saved simulation states.

Return type:

GenerateFn

Returns:

A function taking energy params and the current traj_state (including neighbor list) that runs the simulation forward generating the next TrajectoryState.

Contents