<commit_before>"""OpenMM molecular dynamics runner with accessory classes.

OpenMM is a library with support for running molecular dynamics
simulations with specific support for fast GPU calculations. The
component based architecture of OpenMM makes it a perfect fit with
wepy.

In addition to the principle OpenMMRunner class there are a few
classes here that make using OpenMM runner more efficient.

First is a WalkerState class (OpenMMState) that wraps the openmm state
object directly, itself is a wrapper around the C++
datastructures. This gives better performance by not performing copies
to a WalkerState dictionary.

Second, is the OpenMMWalker which is identical to the Walker class
except that it enforces the state is an actual instantiation of
OpenMMState. Use of this is optional.

Finally, is the OpenMMGPUWorker class. This is to be used as the
worker type for the WorkerMapper work mapper. This is necessary to
allow passing of the device index to OpenMM for which GPU device to
use.

"""
from copy import copy
import random as rand
from warnings import warn
import logging
import time

import numpy as np

import simtk.openmm.app as omma
import simtk.openmm as omm
import simtk.unit as unit

from wepy.walker import Walker, WalkerState
from wepy.runners.runner import Runner
from wepy.work_mapper.worker import Worker
from wepy.work_mapper.task_mapper import WalkerTaskProcess
from wepy.reporter.reporter import Reporter
from wepy.util.util import box_vectors_to_lengths_angles

## Constants

KEYS = ('positions', 'velocities', 'forces', 'kinetic_energy',
        'potential_energy', 'time', 'box_vectors', 'box_volume',
        'parameters', 'parameter_derivatives')
"""Names of the fields of the OpenMMState."""

# when we use the get_state function from the simulation context we
# can pass options for what kind of data to get, this is the default
# to get all the data. TODO not really sure what the 'groups' keyword
# is for though
GET_STATE_KWARG_DEFAULTS = (('getPositions', True),
                            ('getVelocities', True),
                            ('getForces', True),
                            ('getEnergy', True),
                            ('getParameters', True),
                            ('getParameterDerivatives', True),
                            ('enforcePeriodicBox', True),)
"""Mapping of key word arguments to the simulation.context.getState
method for retrieving data for a simulation state. By default we set
each as True to retrieve all information. The presence or absence of
them is handled by the OpenMMState.

"""

# the Units objects that OpenMM uses internally and are returned from
# simulation data
UNITS = (('positions_unit', unit.nanometer),
         ('time_unit', unit.picosecond),
         ('box_vectors_unit', unit.nanometer),
         ('velocities_unit', unit.nanometer/unit.picosecond),
         ('forces_unit', unit.kilojoule / (unit.nanometer * unit.mole)),
         ('box_volume_unit', unit.nanometer),
         ('kinetic_energy_unit', unit.kilojoule / unit.mole),
         ('potential_energy_unit', unit.kilojoule / unit.mole),
        )
"""Mapping of units identifiers to the corresponding simtk.units Unit objects."""

# the names of the units from the units objects above. This is used
# for saving them to files
UNIT_NAMES = (('positions_unit', unit.nanometer.get_name()),
         ('time_unit', unit.picosecond.get_name()),
         ('box_vectors_unit', unit.nanometer.get_name()),
         ('velocities_unit', (unit.nanometer/unit.picosecond).get_name()),
         ('forces_unit', (unit.kilojoule / (unit.nanometer * unit.mole)).get_name()),
         ('box_volume_unit', unit.nanometer.get_name()),
         ('kinetic_energy_unit', (unit.kilojoule / unit.mole).get_name()),
         ('potential_energy_unit', (unit.kilojoule / unit.mole).get_name()),
        )
"""Mapping of unit identifier strings to the serialized string spec of the unit."""

# a random seed will be chosen from 1 to RAND_SEED_RANGE_MAX when the
# Langevin integrator is created. 0 is the default and special value
# which will then choose a random value when the integrator is created

#TODO: test this isn't needed
#RAND_SEED_RANGE_MAX = 1000000

# the runner for the simulation which runs the actual dynamics
class OpenMMRunner(Runner):
    """Runner for OpenMM simulations."""

    def __init__(self, system, topology, integrator,
                 platform=None,
                 enforce_box=False):<commit_msg>deprecate the UNITS constant<commit_after>"""OpenMM molecular dynamics runner with accessory classes.

OpenMM is a library with support for running molecular dynamics
simulations with specific support for fast GPU calculations. The
component based architecture of OpenMM makes it a perfect fit with
wepy.

In addition to the principle OpenMMRunner class there are a few
classes here that make using OpenMM runner more efficient.

First is a WalkerState class (OpenMMState) that wraps the openmm state
object directly, itself is a wrapper around the C++
datastructures. This gives better performance by not performing copies
to a WalkerState dictionary.

Second, is the OpenMMWalker which is identical to the Walker class
except that it enforces the state is an actual instantiation of
OpenMMState. Use of this is optional.

Finally, is the OpenMMGPUWorker class. This is to be used as the
worker type for the WorkerMapper work mapper. This is necessary to
allow passing of the device index to OpenMM for which GPU device to
use.

"""
from copy import copy
import random as rand
from warnings import warn
import logging
import time

import numpy as np

try:
    import simtk.openmm.app as omma
    import simtk.openmm as omm
    import simtk.unit as unit
except ModuleNotFoundError:
    raise ModuleNotFoundError("OpenMM has not been installed, which this runner requires.")

from wepy.walker import Walker, WalkerState
from wepy.runners.runner import Runner
from wepy.work_mapper.worker import Worker
from wepy.work_mapper.task_mapper import WalkerTaskProcess
from wepy.reporter.reporter import Reporter
from wepy.util.util import box_vectors_to_lengths_angles

## Constants

KEYS = ('positions', 'velocities', 'forces', 'kinetic_energy',
        'potential_energy', 'time', 'box_vectors', 'box_volume',
        'parameters', 'parameter_derivatives')
"""Names of the fields of the OpenMMState."""

# when we use the get_state function from the simulation context we
# can pass options for what kind of data to get, this is the default
# to get all the data. TODO not really sure what the 'groups' keyword
# is for though
GET_STATE_KWARG_DEFAULTS = (('getPositions', True),
                            ('getVelocities', True),
                            ('getForces', True),
                            ('getEnergy', True),
                            ('getParameters', True),
                            ('getParameterDerivatives', True),
                            ('enforcePeriodicBox', True),)
"""Mapping of key word arguments to the simulation.context.getState
method for retrieving data for a simulation state. By default we set
each as True to retrieve all information. The presence or absence of
them is handled by the OpenMMState.

"""

# the Units objects that OpenMM uses internally and are returned from
# simulation data

# TODO: this is never used and we only need the unit names. Its okay
# to use simtk.units here but other runners should use a units sytem
# like pint which is easier to install. So we should remove this since
# its not used.

# UNITS = (('positions_unit', unit.nanometer),
#          ('time_unit', unit.picosecond),
#          ('box_vectors_unit', unit.nanometer),
#          ('velocities_unit', unit.nanometer/unit.picosecond),
#          ('forces_unit', unit.kilojoule / (unit.nanometer * unit.mole)),
#          ('box_volume_unit', unit.nanometer),
#          ('kinetic_energy_unit', unit.kilojoule / unit.mole),
#          ('potential_energy_unit', unit.kilojoule / unit.mole),
#         )
# """Mapping of units identifiers to the corresponding simtk.units Unit objects."""

# the names of the units from the units objects above. This is used
# for saving them to files
UNIT_NAMES = (('positions_unit', unit.nanometer.get_name()),
         ('time_unit', unit.picosecond.get_name()),
         ('box_vectors_unit', unit.nanometer.get_name()),
         ('velocities_unit', (unit.nanometer/unit.picosecond).get_name()),
         ('forces_unit', (unit.kilojoule / (unit.nanometer * unit.mole)).get_name()),
         ('box_volume_unit', unit.nanometer.get_name()),
         ('kinetic_energy_unit', (unit.kilojoule / unit.mole).get_name()),
         ('potential_energy_unit', (unit.kilojoule / unit.mole).get_name()),
        )
"""Mapping of unit identifier strings to the serialized string spec of the unit."""

# a random seed will be chosen from 1 to RAND_SEED_RANGE_MAX when the
# Langevin integrator is created. 0 is the default and special value
# which will then choose a random value when the integrator is created

# TODO: test this isn't needed
# RAND_SEED_RANGE_MAX = 1000000

# the runner for the simulation which runs the actual dynamics
class OpenMMRunner(Runner):
    """Runner for OpenMM simulations."""

    def __init__(self, system, topology, integrator,
                 platform=None,
                 enforce_box=False):
{'fraction_non_alphanumeric': 0.08246871323136164, 'fraction_numerical': 0.0019253396085142796, 'mean_word_length': 3.525653436592449, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4197431', 'n_tokens_mistral': 2573, 'n_tokens_neox': 2496, 'n_words': 1031}