ensemble.sampling.quantity_traj

Contents

ensemble.sampling.quantity_traj#

quantity_traj(traj_state, quantities, energy_params=None, batch_size=1)[source]#

Computes quantities of interest for all states in a trajectory.

Arbitrary quantity functions can be provided via the quantities-dict. The quantities dict provides the function to compute the quantity on a single snapshot. The resulting quantity trajectory will be saved in a dict under the same key as the input quantity function.

Note

This version exists for backward-compatibility. Consider using the more flexible version quantity_map() instead.

Example usage:
def custom_compute_fn(state, neighbor=None, **kwargs):
    ...
    return quantity_snapshot


quantities = {
    'energy': custom_quantity.energy_wrapper(energy_template_fn),
    'custom_quantity': custom_compute_fn
}

quantity_trajs = quantity_traj(traj_state, quantities, energy_params)
custom_quantity = quantity_trajs['custom_quantity']
Parameters:

  • traj_state – Trajectory state from the trajectory generator

  • quantities – The quantity dict containing for each target quantity the snapshot compute function

  • energy_params – Energy params for energy_fn_template to initialize the current energy_fn

  • batch_size – Number of batches for vmap

Returns:

A dict of quantity trajectories saved under the same key as the input quantity function.

Contents