ensemble.evaluation.quantity_map

Contents

ensemble.evaluation.quantity_map#

quantity_map(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.

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.

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


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

# Results will be available to all snapshot compute functions
feature_extract_fns = {
    'feature': custom_feature_compute_fn
}

quantity_trajs = quantity_map(
    trajectory, quantities, reference_nbrs, dynamic_kwargs,
    energy_params, feature_extract_fns=feature_extract_fns
)
custom_quantity = quantity_trajs['custom_quantity']
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