'''For extracting properties from OpenMM Contexts (e.g. positions, energies, etc)'''
__author__ = 'Timotej Bernat'
__email__ = 'timotej.bernat@colorado.edu'
from typing import Optional
from openmm import Context, State
from openmm.unit import Unit, Quantity
from .forces import impose_unique_force_groups, forces_by_force_group
# POSITIONS
[docs]
def get_context_positions(context : Context) -> Quantity:
'''Extract coordinates from the current state of a simulation'''
# NOTE : forcing numpy output for now, as OpenMM Vec3's don't seem particularly useful; may change this in the future
return context.getState(getPositions=True).getPositions(asNumpy=True)
# ENERGIES
[docs]
def get_openmm_energies(
context : Context,
preferred_unit : Optional[Unit]=None,
force_group_names : Optional[dict[int, str]]=None,
) -> dict[str, Quantity]:
'''
Evaluate energies of an OpenMM Context
Returns dict, keyed by contribution name, containing the total kinetic and potential energies,
as well as the individual potential energies contributed by each distinct force group
'''
if force_group_names is None:
force_group_names = {}
energies : dict[str, Quantity] = {}
## global
global_state : State = context.getState(getEnergy=True) # initialize shared global state
energies['Total potential energy'] = global_state.getPotentialEnergy()
energies['Total kinetic energy' ] = global_state.getKineticEnergy()
## by contribution
for group_id in forces_by_force_group(context.getSystem()):
local_state : State = context.getState(getEnergy=True, groups={group_id})
group_label = force_group_names.get(group_id, f'Group {group_id}')
energies[f'{group_label} potential energy'] = local_state.getPotentialEnergy()
if preferred_unit is not None:
for key, energy in energies.items():
energies[key] = energy.in_units_of(preferred_unit)
return energies
eval_openmm_energies = get_openmm_energies
[docs]
def get_openmm_energies_separated(context : Context, preferred_unit : Optional[Unit]=None) -> dict[str, Quantity]:
'''
Evaluate energies of an OpenMM Context
Enforces separation of each Force's contribution to the energy (i.e. by imposing unique force groups)
Returns dict, keyed by contribution name, containing the total kinetic and potential energies,
as well as the individual potential energies contributed by each distinct force group
'''
ommsys = context.getSystem()
impose_unique_force_groups(ommsys)
return get_openmm_energies(
context,
preferred_unit=preferred_unit,
force_group_names={
group_id : force.getName()
for group_id, force in forces_by_force_group(ommsys, denest=True).items()
},
)
eval_openmm_energies_separated = get_openmm_energies_separated