<commit_before>#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Thermodynamic ensemble utilities.

DESCRIPTION

This module provides a utilities for simulating physical thermodynamic ensembles.

Provided classes include:

* ThermodynamicState - Data specifying a thermodynamic state obeying Boltzmann statistics (System/temperature/pressure/pH)

DEPENDENCIES

TODO

COPYRIGHT

@author <NAME> <<EMAIL>>

"""

#=============================================================================================
# GLOBAL IMPORTS
#=============================================================================================

import copy
import numpy as np

import simtk.openmm as mm
import simtk.unit as units
from openmmtools import testsystems

import logging
logger = logging.getLogger(__name__)

#=============================================================================================
# REVISION CONTROL
#=============================================================================================

__version__ = "$Revision: $"

#=============================================================================================
# MODULE CONSTANTS
#=============================================================================================

kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA # Boltzmann constant

#=============================================================================================
# Thermodynamic state description
#=============================================================================================


class ThermodynamicState(object):
    """Object describing a thermodynamic state obeying Boltzmann statistics.

    Examples
    --------

    Specify an NVT state for a water box at 298 K.

    >>> system_container = testsystems.WaterBox()
    >>> (system, positions) = system_container.system, system_container.positions
    >>> state = ThermodynamicState(system=system, temperature=298.0*units.kelvin)

    Get the inverse temperature
    
    >>> beta = state.beta

    Specify an NPT state at 298 K and 1 atm pressure.

    >>> state = ThermodynamicState(system=system, temperature=298.0*units.kelvin, pressure=1.0*units.atmospheres)

    Note that the pressure is only relevant for periodic systems.

    A barostat will be added to the system if none is attached.

    Notes
    -----

    This state object cannot describe states obeying non-Boltzamnn statistics, such as Tsallis statistics.

    ToDo
    ----

    * Implement a more fundamental ProbabilityState as a base class?
    * Implement pH.

    """
    def __init__(self, system, temperature, pressure=None):
        """Construct a thermodynamic state with given system and temperature.

        Parameters
        ----------

        system : simtk.openmm.System
            System object describing the potential energy function
            for the system (default: None)
        temperature : simtk.unit.Quantity, compatible with 'kelvin'
            Temperature for a system with constant temperature
        pressure : simtk.unit.Quantity,  compatible with 'atmospheres', optional, default=None
            If not None, specifies the pressure for constant-pressure systems.


        """

        # Initialize.
        self.system = None          # the System object governing the potential energy computation
        self.temperature = None     # the temperature
        self.pressure = None        # the pressure, or None if not isobaric

        self._cache_context = True  # if True, try to cache Context object
        self._context = None        # cached Context
        self._integrator = None     # cached Integrator

        # Store provided values.
        if system is not None:
            if type(system) is not mm.System:
                raise(TypeError("system must be an OpenMM System; instead found %s" % type(system)))
            self.system = copy.deepcopy(system) # TODO: Do this when deep copy works.
            # self.system = system # we make a shallow copy for now, which can cause trouble later
        if temperature is not None:
            self.temperature = temperature
        if pressure is not None:
            self.pressure = pressure

        # If temperature and pressure are specified, make sure MonteCarloBarostat is attached.
        if temperature and pressure:
            # Try to find barostat.
            barostat = False
            for force_index in range(self.system.getNumForces()):
                force = self.system.getForce(force_index)
                # Dispatch forces
                if isinstance(force, mm.MonteCarloBarostat):
                    barostat = force
                    break
            if barostat:
                # Set temperature.
                # TODO: Set pressure too, once that option is available.
                barostat.setTemperature(temperature)
            else:
                # Create barostat.
                barostat = mm.MonteCarloBarostat(pressure, temperature)
                self.system.addForce(barostat)

        return

    def _create_context(self, platform=None):
        """Create Integrator and Context objects if they do not already exist.
        """

        # Check if we already have a Context defined.
        if self._context:
            #if platform and (platform != self._context.getPlatform()): # TODO: Figure out why requested and cached platforms differed in tests.
            if platform and (platform.getName() != self._context.getPlatform().getName()): # DEBUG: Only compare Platform names for now; change this later to incorporate GPU IDs.
                # Platform differs from the one requested; destroy it.
                logger.info((platform.getName(), self._context.getPlatform().getName()))
                logger.debug("Platform differs from the one requested; destroying and recreating...")
                del self._context, self._integrator
            else:
                # Cached context is what we expect; do nothing.
                return

        # Create an integrator.
        timestep = 1.0 * units.femtosecond
        self._integrator = mm.VerletIntegrator(timestep)

        # Create a new OpenMM context.
        if platform:
            self._context = mm.Context(self.system, self._integrator, platform)<commit_msg>Correctly set pressure and temperature in ThermodynamicState<commit_after>#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Thermodynamic ensemble utilities.

DESCRIPTION

This module provides a utilities for simulating physical thermodynamic ensembles.

Provided classes include:

* ThermodynamicState - Data specifying a thermodynamic state obeying Boltzmann statistics (System/temperature/pressure/pH)

DEPENDENCIES

TODO

COPYRIGHT

@author <NAME> <<EMAIL>>

"""

#=============================================================================================
# GLOBAL IMPORTS
#=============================================================================================

import copy
import numpy as np

import simtk.openmm as mm
import simtk.unit as units
from openmmtools import testsystems

import logging
logger = logging.getLogger(__name__)

#=============================================================================================
# REVISION CONTROL
#=============================================================================================

__version__ = "$Revision: $"

#=============================================================================================
# MODULE CONSTANTS
#=============================================================================================

kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA # Boltzmann constant

#=============================================================================================
# Thermodynamic state description
#=============================================================================================


class ThermodynamicState(object):
    """Object describing a thermodynamic state obeying Boltzmann statistics.

    Examples
    --------

    Specify an NVT state for a water box at 298 K.

    >>> system_container = testsystems.WaterBox()
    >>> (system, positions) = system_container.system, system_container.positions
    >>> state = ThermodynamicState(system=system, temperature=298.0*units.kelvin)

    Get the inverse temperature

    >>> beta = state.beta

    Specify an NPT state at 298 K and 1 atm pressure.

    >>> state = ThermodynamicState(system=system, temperature=298.0*units.kelvin, pressure=1.0*units.atmospheres)

    Note that the pressure is only relevant for periodic systems.

    A barostat will be added to the system if none is attached.

    Notes
    -----

    This state object cannot describe states obeying non-Boltzamnn statistics, such as Tsallis statistics.

    ToDo
    ----

    * Implement a more fundamental ProbabilityState as a base class?
    * Implement pH.

    """
    def __init__(self, system, temperature, pressure=None):
        """Construct a thermodynamic state with given system and temperature.

        Parameters
        ----------

        system : simtk.openmm.System
            System object describing the potential energy function
            for the system (default: None)
        temperature : simtk.unit.Quantity, compatible with 'kelvin'
            Temperature for a system with constant temperature
        pressure : simtk.unit.Quantity,  compatible with 'atmospheres', optional, default=None
            If not None, specifies the pressure for constant-pressure systems.


        """

        # Initialize.
        self.system = None          # the System object governing the potential energy computation
        self.temperature = None     # the temperature
        self.pressure = None        # the pressure, or None if not isobaric

        self._cache_context = True  # if True, try to cache Context object
        self._context = None        # cached Context
        self._integrator = None     # cached Integrator

        # Store provided values.
        if system is not None:
            if type(system) is not mm.System:
                raise(TypeError("system must be an OpenMM System; instead found %s" % type(system)))
            self.system = copy.deepcopy(system) # TODO: Do this when deep copy works.
            # self.system = system # we make a shallow copy for now, which can cause trouble later
        if temperature is not None:
            self.temperature = temperature
        if pressure is not None:
            self.pressure = pressure

        # If temperature and pressure are specified, make sure MonteCarloBarostat is attached.
        if temperature and pressure:
            # Try to find barostat.
            barostat = False
            for force_index in range(self.system.getNumForces()):
                force = self.system.getForce(force_index)
                # Dispatch forces
                if isinstance(force, mm.MonteCarloBarostat):
                    barostat = force
                    break
            if barostat:
                # Set temperature.
                if hasattr(barostat, 'setDefaultTemperature'):
                    barostat.setDefaultTemperature(temperature)
                elif hasattr(barostat, 'setTemperature'):
                    barostat.setTemperature(temperature)
                else:
                    raise Exception("barostat does not have 'setTemperature' or 'setDefaultTemperature' interfaces!")
                # Set pressure
                barostat.setDefaultPressure(pressure)
            else:
                # Create barostat.
                barostat = mm.MonteCarloBarostat(pressure, temperature)
                self.system.addForce(barostat)

        return

    def _create_context(self, platform=None):
        """Create Integrator and Context objects if they do not already exist.
        """

        # Check if we already have a Context defined.
        if self._context:
            #if platform and (platform != self._context.getPlatform()): # TODO: Figure out why requested and cached platforms differed in tests.
            if platform and (platform.getName() != self._context.getPlatform().getName()): # DEBUG: Only compare Platform names for now; change this later to incorporate GPU IDs.
                # Platform differs from the one requested; destroy it.
                logger.info((platform.getName(), self._context.getPlatform().getName()))
                logger.debug("Platform differs from the one requested; destroying and recreating...")
                del self._context, self._integrator
            else:
                # Cached context is what we expect; do nothing.
                return

        # Create an integrator.
        timestep = 1.0 * units.femtosecond
        self._integrator = mm.VerletIntegrator(timestep)

        # Create a new OpenMM context.
        if platform:
            self._context = mm.Context(self.system, self._integrator, platform)
{'fraction_non_alphanumeric': 0.20296062518785693, 'fraction_numerical': 0.0028554253080853623, 'mean_word_length': 2.8972181551976575, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 39, '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': '4419838', 'n_tokens_mistral': 3174, 'n_tokens_neox': 2827, 'n_words': 1016}