ensemble.evaluation.quantity_multimap

ensemble.evaluation.quantity_multimap#

quantity_multimap(*states, quantities, nbrs=None, state_kwargs=None, constant_state_kwargs=None, energy_params=None, batch_size=1, feature_extract_fns=None)[source]#

Computes quantities of interest for all states in a trajectory.

This function extends quantity_map() to quantities with respect to multiple reference states. Therefore, the quantity function signature changes to

def quantity_fn(*states, neighbor=None, energy_params=None, **kwargs):
    ...

The keywords arguments, i.e. the neighbor list, are with respect to the first state of *states.

Parameters:

  • states (State) – System states, concatenated along the first dimensions of the arrays.

  • quantities (Dict[str, QuantityComputeFunction]) – The quantity dict containing for each target quantity the snapshot compute function

  • nbrs (NeighborList) – Reference neighbor list to compute new neighbor list

  • state_kwargs (Dict[str, Array]) – Kwargs to supply reference 'kT' and/or 'pressure' to the energy function or the quantity functions.

  • constant_state_kwargs (Dict[str, Array]) – Kwargs to supply information to the energy function that is constant over all states.

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

  • batch_size (int) – Number of batches for vmap

  • feature_extract_fns (Dict[str, Callable]) – Callables to compute features accessible to all snapshot compute functions.

Returns:

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

Contents