#!/usr/local/bin/env python

#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Module to generate Systems and coordinates for simple reference molecular systems for testing.

DESCRIPTION

This module provides functions for building a number of simple test systems for testing purposes.

EXAMPLES

Import testsystems.

>>> import testsystems

Create a sodium chloride crystal.

>>> [system, coordinates] = testsystems.SodiumChlorideCrystal()

Create a small periodic box of Lennard-Jones particles.

>>> [system, coordinates] = testsystems.LennardJonesFluid()

COPYRIGHT

@author Randall J. Radmer <radmer@stanford.edu>
@author John D. Chodera <jchodera@gmail.com>
@author Kim M Branson <kim.branson@gmail.com>

All code in this repository is released under the GNU General Public License.

This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by the Free Software
Foundation, either version 3 of the License, or (at your option) any later
version.

This program is distributed in the hope that it will be useful, but WITHOUT ANY
WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
PARTICULAR PURPOSE.  See the GNU General Public License for more details.
 
You should have received a copy of the GNU General Public License along with
this program.  If not, see <http://www.gnu.org/licenses/>.

TODO

* Add units checking code to check arguments.
* Change default arguments to Quantity objects, rather than None.

"""

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

import os
import os.path
import numpy
import math

import simtk
import simtk.chem.openmm as openmm
import simtk.unit as units

#=============================================================================================
# 3D Harmonic Oscillator
#=============================================================================================

def HarmonicOscillator(K=None, mass=None, mm=None):
    """
    Create a 3D harmonic oscillator, with a single particle confined in an isotropic harmonic well.

    OPTIONAL ARGUMENTS
    
    K (simtk.unit.Quantity of energy/particle/distance^2) - harmonic restraining potential (default: 1.0 * units.kilojoules_per_mole/units.nanometer**2)    
    mass (simtk.unit.Quantity of mass) - particle mass (default: 39.948 * units.amu)
    mm (simtk.chem.openmm or compatible) - openmm implementation (default: simtk.chem.openmm)

    RETURNS

    system (simtk.chem.System) - the system object
    coordinates (simtk.unit.Quantity of distance, Nx3 numpy array) - coordinates of all particles in system

    NOTES

    The natural period of a harmonic oscillator is T = sqrt(m/K), so you will want to use an
    integration timestep smaller than ~ T/10.

    EXAMPLES

    Create a 3D harmonic oscillator with default parameters.

    >>> [system, coordinates] = HarmonicOscillator()

    Create a harmonic oscillator with specified mass and spring constant.

    >>> mass = 12.0 * units.amu
    >>> K = 1.0 * units.kilocalories_per_mole / units.angstroms**2
    >>> [system, coordinates] = HarmonicOscillator(K=K, mass=mass)

    Test the energy

    >>> # Create a Context.
    >>> temperature = 298.0 * units.kelvin
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy().in_units_of(units.kilocalories_per_mole)
    0.0 kcal/mol
    
    Integrate dynamics

    >>> nsteps = 100 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    # Use pyOpenMM by default.
    if mm is None:
        mm = simtk.chem.openmm

    # Default parameters
    if K       is None:  K           = 1.0 * units.kilojoules_per_mole / units.nanometer**2
    if mass    is None:  mass        = 39.948 * units.amu # arbitrary reference mass            

    # Create an empty system object.
    system = mm.System()

    # Add the particle to the system.
    system.addParticle(mass)

    # Set the coordinates.
    coordinates = units.Quantity(numpy.zeros([1,3], numpy.float32), units.angstroms)

    # Add a restrining potential centered at the origin.
    force = mm.CustomExternalForce('(K/2.0) * (x^2 + y^2 + z^2)')
    force.addGlobalParameter('K', K)
    force.addParticle(0, [])
    system.addForce(force)

    return (system, coordinates)

#=============================================================================================
# Free diatomic harmonic oscillator
#=============================================================================================

def Diatom(K=None, r0=None, m1=None, m2=None, mm=None):
    """
    Create a diatomic molecule with a single harmonic bond between the two atoms.

    OPTIONAL ARGUMENTS
    
    K (simtk.unit.Quantity of energy/particle/distance^2) - harmonic bond potential (default: 1.0 * units.kilojoules_per_mole/units.nanometer**2)
    r0 (simtk.unit.Quantity of distance) - bond length
    m1 (simtk.unit.Quantity of mass) - particle 1 mass (default: 39.948 * units.amu)
    m2 (simtk.unit.Quantity of mass) - particle 2 mass (default: 39.948 * units.amu)    
    mm (simtk.chem.openmm or compatible) - openmm implementation (default: simtk.chem.openmm)

    RETURNS

    system (simtk.chem.System) - the system object
    coordinates (simtk.unit.Quantity of distance, Nx3 numpy array) - coordinates of all particles in system

    NOTES

    The natural period of a harmonic oscillator is T = sqrt(m/K), so you will want to use an
    integration timestep smaller than ~ T/10.

    EXAMPLES

    Create a diatom with default parameters.

    >>> [system, coordinates] = Diatom()

    Test the energy

    >>> # Create a Context.
    >>> temperature = 298.0 * units.kelvin
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy().in_units_of(units.kilocalories_per_mole)
    0.0 kcal/mol
    
    Integrate dynamics

    >>> nsteps = 100 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    # Use pyOpenMM by default.
    if mm is None:
        mm = simtk.chem.openmm

    # Default parameters
    if K       is None:  K           = 1.0 * units.kilocalories_per_mole / units.nanometer**2
    if r0      is None:  r0          = 1.0 * units.nanometers
    if m1      is None:  m1          = 39.948 * units.amu 
    if m2      is None:  m2          = 39.948 * units.amu

    # Create an empty system object.
    system = mm.System()

    # Add two particles to the system.
    system.addParticle(m1)
    system.addParticle(m2)

    # Add a harmonic bond.
    force = mm.HarmonicBondForce()
    force.addBond(0, 1, r0, K)
    system.addForce(force)

    # Set the coordinates.
    coordinates = units.Quantity(numpy.zeros([2,3], numpy.float32), units.angstroms)
    coordinates[1,0] = r0

    # Add a central restraining potential.
    Kcentral = 1.0 * units.kilocalories_per_mole / units.nanometer**2
    force = mm.CustomExternalForce('(Kcentral/2.0) * (x^2 + y^2 + z^2)')
    force.addGlobalParameter('K', K)
    force.addParticle(0, [])
    force.addParticle(1, [])    
    #system.addForce(force)

    return (system, coordinates)

#=============================================================================================
# Constraint-coupled Pair of Harmonic Oscillators
#=============================================================================================

def ConstraintCoupledHarmonicOscillator(K=None, d=None, mass=None, mm=None):
    """
    Create a pair of particles in 3D harmonic oscillator wells, coupled by a constraint.

    @keyword K: harmonic restraining potential (default: 1.0 * units.kilojoules_per_mole/units.nanometer**2)    
    @keyword d: distance between harmonic oscillators (default: 1.0 * units.nanometer)
    @keyword mass: particle mass (default: 39.948 * units.amu)

    @return system: the system
    @type system: a System object

    @return coordinates: initial coordinates for the system
    @type coordinates: a Coordinates object

    EXAMPLES

    Create a constraint-coupled 3D harmonic oscillator with default parameters.

    >>> [system, coordinates] = ConstraintCoupledHarmonicOscillator()

    Create a constraint-coupled harmonic oscillator with specified mass, distance, and spring constant.

    >>> mass = 12.0 * units.amu
    >>> d = 5.0 * units.angstroms
    >>> K = 1.0 * units.kilocalories_per_mole / units.angstroms**2
    >>> [system, coordinates] = ConstraintCoupledHarmonicOscillator(K=K, d=d, mass=mass)

    Test the energy

    >>> # Create a Context.
    >>> temperature = 298.0 * units.kelvin
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy().in_units_of(units.kilocalories_per_mole)
    0.0 kcal/mol
    
    Integrate dynamics

    >>> nsteps = 100 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    # Use pyOpenMM by default.
    if mm is None:
        mm = simtk.chem.openmm

    # Default parameters
    if K       is None:  K           = 1.0 * units.kilojoules_per_mole / units.nanometer**2
    if d       is None:  d           = 1.0 * units.nanometer
    if mass    is None:  mass        = 39.948 * units.amu # arbitrary reference mass            

    # Create an empty system object.
    system = mm.System()

    # Add particles to the system.
    system.addParticle(mass)
    system.addParticle(mass)    

    # Set the coordinates.
    coordinates = units.Quantity(numpy.zeros([2,3], numpy.float32), units.angstroms)
    coordinates[1,0] = d

    # Add a restrining potential centered at the origin.
    force = mm.CustomExternalForce('(K/2.0) * ((x-d)^2 + y^2 + z^2)')
    force.addGlobalParameter('K', K)
    force.addPerParticleParameter('d')    
    force.addParticle(0, [0.0])
    force.addParticle(1, [d / units.nanometers])    
    system.addForce(force)

    # Add constraint between particles.
    system.addConstraint(0, 1, d)   

    # Add a harmonic bond force as well so minimization will roughly satisfy constraints.
    force = mm.HarmonicBondForce()
    K = 10.0 * units.kilocalories_per_mole / units.angstrom**2 # force constant
    force.addBond(0, 1, d, K)
    system.addForce(force)

    return (system, coordinates)

#=============================================================================================
# Array of harmonic oscillators
#=============================================================================================

def HarmonicOscillatorArray(K=None, d=None, mass=None, N=None, mm=None):
    """
    Create a 1D array of noninteracting particles in 3D harmonic oscillator wells.

    @keyword K: harmonic restraining potential (default: 1.0 * units.kilojoules_per_mole/units.nanometer**2)    
    @keyword d: distance between harmonic oscillators (default: 1.0 * units.nanometer)
    @keyword mass: particle mass (default: 39.948 * units.amu)

    @return system: the system
    @type system: a System object

    @return coordinates: initial coordinates for the system
    @type coordinates: a Coordinates object

    EXAMPLES

    Create a constraint-coupled 3D harmonic oscillator with default parameters.

    >>> [system, coordinates] = HarmonicOscillatorArray()

    Create a constraint-coupled harmonic oscillator with specified mass, distance, and spring constant.

    >>> mass = 12.0 * units.amu
    >>> d = 5.0 * units.angstroms
    >>> K = 1.0 * units.kilocalories_per_mole / units.angstroms**2
    >>> N = 10 # number of oscillators
    >>> [system, coordinates] = HarmonicOscillatorArray(K=K, d=d, mass=mass, N=N)

    Test the energy

    >>> # Create a Context.
    >>> temperature = 298.0 * units.kelvin
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy().in_units_of(units.kilocalories_per_mole)
    0.0 kcal/mol
    
    Integrate dynamics

    >>> nsteps = 100 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    # Use pyOpenMM by default.
    if mm is None:
        mm = simtk.chem.openmm

    # Default parameters
    if K       is None:  K           = 1.0 * units.kilojoules_per_mole / units.nanometer**2
    if d       is None:  d           = 1.0 * units.nanometer
    if mass    is None:  mass        = 39.948 * units.amu 
    if N       is None:  N           = 5 

    # Create an empty system object.
    system = mm.System()

    # Add particles to the system.
    for n in range(N):
        system.addParticle(mass)

    # Set the coordinates for a 1D array of particles spaced d apart along the x-axis.
    coordinates = units.Quantity(numpy.zeros([N,3], numpy.float32), units.angstroms)
    for n in range(N):
        coordinates[n,0] = n*d

    # Add a restrining potential for each oscillator.
    force = mm.CustomExternalForce('(K/2.0) * ((x-x0)^2 + y^2 + z^2)')
    force.addGlobalParameter('K', K)
    force.addPerParticleParameter('x0')
    for n in range(N):
        parameters = (d*n / units.nanometers, )
        force.addParticle(n, parameters)
    system.addForce(force)

    return (system, coordinates)

#=============================================================================================
# Salt crystal
#=============================================================================================

def SodiumChlorideCrystal(mm=None):
    """
    Create an FCC crystal of sodium chloride.

    Each atom is represented by a charged Lennard-Jones sphere in an Ewald lattice.

    @return system: the system
    @type system: a System object

    @return coordinates: initial coordinates for the system
    @type coordinates: a Coordinates object

    EXAMPLES

    Create sodium chloride crystal unit.
    
    >>> [system, coordinates] = SodiumChlorideCrystal()

    Test the energy

    >>> # Create a Context.
    >>> temperature = 298.0 * units.kelvin
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy().in_units_of(units.kilocalories_per_mole)
    -120.096618769 kcal/mol

    Integrate dynamics

    >>> nsteps = 10 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))
    
    TODO

    * Add nx, ny, nz arguments to allow user to specify replication of crystal unit in x,y,z.
    * Choose more appropriate lattice parameters and lattice spacing.

    """

    # Choose OpenMM package if not specified.
    if mm is None: mm = simtk.chem.openmm

    # Set default parameters (from Tinker).
    mass_Na     = 22.990 * units.amu
    mass_Cl     = 35.453 * units.amu
    q_Na        = 1.0 * units.elementary_charge 
    q_Cl        =-1.0 * units.elementary_charge 
    sigma_Na    = 3.330445 * units.angstrom
    sigma_Cl    = 4.41724 * units.angstrom
    epsilon_Na  = 0.002772 * units.kilocalorie_per_mole 
    epsilon_Cl  = 0.118 * units.kilocalorie_per_mole 

    # Create system
    system = mm.System()

    # Set box vectors.
    box_size = 5.628 * units.angstroms # box width
    a = units.Quantity(numpy.zeros([3]), units.nanometers); a[0] = box_size
    b = units.Quantity(numpy.zeros([3]), units.nanometers); b[1] = box_size
    c = units.Quantity(numpy.zeros([3]), units.nanometers); c[2] = box_size
    system.setPeriodicBoxVectors(a, b, c)

    # Create nonbonded force term.
    force = mm.NonbondedForce()

    # Set interactions to be periodic Ewald.
    force.setNonbondedMethod(mm.NonbondedForce.Ewald)

    # Set cutoff to be less than one half the box length.
    cutoff = box_size / 2.0 * 0.99
    force.setCutoffDistance(cutoff)
    
    # Allocate storage for coordinates.
    natoms = 2
    coordinates = units.Quantity(numpy.zeros([natoms,3], numpy.float32), units.angstroms)

    # Add sodium ion.
    system.addParticle(mass_Na)
    force.addParticle(q_Na, sigma_Na, epsilon_Na)
    coordinates[0,0] = 0.0 * units.angstrom
    coordinates[0,1] = 0.0 * units.angstrom
    coordinates[0,2] = 0.0 * units.angstrom
    
    # Add chloride atom.
    system.addParticle(mass_Cl)
    force.addParticle(q_Cl, sigma_Cl, epsilon_Cl)
    coordinates[1,0] = 2.814 * units.angstrom
    coordinates[1,1] = 2.814 * units.angstrom
    coordinates[1,2] = 2.814 * units.angstrom

    # Add nonbonded force term to the system.
    system.addForce(force)
       
    # Return system and coordinates.
    return (system, coordinates)

#=============================================================================================
# Cluster of Lennard-Jones particles
#=============================================================================================

def LennardJonesCluster(nx=3, ny=3, nz=3, K=None, mm=None):
    """
    Create a non-periodic rectilinear grid of Lennard-Jones particles in a harmonic restraining potential.

    @keyword nx: number of particles in x-dimension (default: 3)
    @keyword ny: number of particles in y-dimension (default: 3)
    @keyword nz: number of particles in z-dimension (default: 3)
    @keyword K: harmonic restraining potential (default: 1.0 * units.kilojoules_per_mole/units.nanometer**2)

    @return system: the system
    @type system: a System object

    @return coordinates: initial coordinates for the system
    @type coordinates: a Coordinates object

    EXAMPLES

    Create default 3x3x3 Lennard-Jones cluster in a harmonic restraining potential.

    >>> [system, coordinates] = LennardJonesCluster()

    Create a larger 10x10x10 grid of Lennard-Jones particles.

    >>> [system, coordinates] = LennardJonesCluster(nx=10, ny=10, nz=10)

    Test the energy

    >>> # Create a Context.
    >>> temperature = 298.0 * units.kelvin
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy().in_units_of(units.kilocalories_per_mole)
    340.714006989 kcal/mol
    
    Integrate dynamics

    >>> nsteps = 10 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    # Use pyOpenMM by default.
    if mm is None:
        mm = simtk.chem.openmm

    # Default parameters
    mass_Ar     = 39.9 * units.amu
    q_Ar        = 0.0 * units.elementary_charge
    sigma_Ar    = 3.350 * units.angstrom
    epsilon_Ar  = 0.001603 * units.kilojoule_per_mole

    # Set spring constant.
    if K is None:
        K = 1.0 * units.kilojoules_per_mole/units.nanometer**2

    scaleStepSizeX = 1.0
    scaleStepSizeY = 1.0
    scaleStepSizeZ = 1.0

    # Determine total number of atoms.
    natoms = nx * ny * nz

    # Create an empty system object.
    system = mm.System()

    # Create a NonbondedForce object with no cutoff.
    nb = mm.NonbondedForce()
    nb.setNonbondedMethod(mm.NonbondedForce.NoCutoff)

    coordinates = units.Quantity(numpy.zeros([natoms,3],numpy.float32), units.angstrom)

    maxX = 0.0 * units.angstrom
    maxY = 0.0 * units.angstrom
    maxZ = 0.0 * units.angstrom

    atom_index = 0
    for ii in range(nx):
        for jj in range(ny):
            for kk in range(nz):
                system.addParticle(mass_Ar)
                nb.addParticle(q_Ar, sigma_Ar, epsilon_Ar)
                x = sigma_Ar*scaleStepSizeX*(ii - nx/2.0)
                y = sigma_Ar*scaleStepSizeY*(jj - ny/2.0)
                z = sigma_Ar*scaleStepSizeZ*(kk - nz/2.0)

                coordinates[atom_index,0] = x
                coordinates[atom_index,1] = y
                coordinates[atom_index,2] = z
                atom_index += 1
                
                # Wrap coordinates as needed.
                if x>maxX: maxX = x
                if y>maxY: maxY = y
                if z>maxZ: maxZ = z

    # Add the nonbonded force.
    system.addForce(nb)

    # Add a restrining potential centered at the origin.
    force = mm.CustomExternalForce('(K/2.0) * (x^2 + y^2 + z^2)')
    force.addGlobalParameter('K', K)
    for particle_index in range(natoms):
        force.addParticle(particle_index, [])
    system.addForce(force)

    return (system, coordinates)

#=============================================================================================
# Periodic box of Lennard-Jones particles
#=============================================================================================

def LennardJonesFluid(nx=3, ny=3, nz=3, mm=None):
    """
    Create a periodic rectilinear grid of Lennard-Jones particles.

    @keyword nx: number of particles in x-dimension (default: 3)
    @keyword ny: number of particles in y-dimension (default: 3)
    @keyword nz: number of particles in z-dimension (default: 3)

    @return system: the system
    @type system: a System object

    @return coordinates: initial coordinates for the system
    @type coordinates: a Coordinates object

    EXAMPLES

    Create default 3x3x3 Lennard-Jones fluid.

    >>> [system, coordinates] = LennardJonesFluid()

    Create a larger 10x8x5 box of Lennard-Jones particles.

    >>> [system, coordinates] = LennardJonesFluid(nx=10, ny=8, nz=5)

    Test the energy

    >>> # Create a Context.
    >>> temperature = 298.0 * units.kelvin
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy().in_units_of(units.kilocalories_per_mole)
    -0.444987251016 kcal/mol
    
    Integrate dynamics

    >>> nsteps = 10 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    # Use pyOpenMM by default.
    if mm is None:
        mm = simtk.chem.openmm

    # Default parameters
    mass_Ar     = 39.9 * units.amu
    q_Ar        = 0.0 * units.elementary_charge
    sigma_Ar    = 3.350 * units.angstrom
    epsilon_Ar  = 0.001603 * units.kilojoule_per_mole

    scaleStepSizeX = 1.0
    scaleStepSizeY = 1.0
    scaleStepSizeZ = 1.0

    # Determine total number of atoms.
    natoms = nx * ny * nz

    # Create an empty system object.
    system = mm.System()

    # Set up periodic nonbonded interactions with a cutoff.
    nb = mm.NonbondedForce()
    nb.setNonbondedMethod(mm.NonbondedForce.CutoffPeriodic)

    coordinates = units.Quantity(numpy.zeros([natoms,3],numpy.float32), units.angstrom)

    maxX = 0.0 * units.angstrom
    maxY = 0.0 * units.angstrom
    maxZ = 0.0 * units.angstrom

    atom_index = 0
    for ii in range(nx):
        for jj in range(ny):
            for kk in range(nz):
                system.addParticle(mass_Ar)
                nb.addParticle(q_Ar, sigma_Ar, epsilon_Ar)
                x = sigma_Ar*scaleStepSizeX*ii
                y = sigma_Ar*scaleStepSizeY*jj
                z = sigma_Ar*scaleStepSizeZ*kk

                coordinates[atom_index,0] = x
                coordinates[atom_index,1] = y
                coordinates[atom_index,2] = z
                atom_index += 1
                
                # Wrap coordinates as needed.
                if x>maxX: maxX = x
                if y>maxY: maxY = y
                if z>maxZ: maxZ = z
                
    # Set periodic box vectors.
    x = maxX+2*sigma_Ar*scaleStepSizeX
    y = maxY+2*sigma_Ar*scaleStepSizeY
    z = maxZ+2*sigma_Ar*scaleStepSizeZ

    a = units.Quantity((x,                0*units.angstrom, 0*units.angstrom))
    b = units.Quantity((0*units.angstrom,                y, 0*units.angstrom))
    c = units.Quantity((0*units.angstrom, 0*units.angstrom, z))
    system.setPeriodicBoxVectors(a, b, c)

    # Set cutoff to be less than one-half the smallest box dimension.
    cutoff = min(a[0], b[1], c[2]) / 2.0 - sigma_Ar / 50.0
    nb.setCutoffDistance(cutoff)

    # Add the nonbonded force.
    system.addForce(nb)

    return (system, coordinates)

#=============================================================================================
# Periodic water box
#=============================================================================================

def WaterBox(waterbox='watbox216.pdb', mm=None, constrain=True, cutoff=None):
    """
    Create a test system containing a periodic box of TIP3P water.

    Flexible bonds and angles are always added, and constraints are optional (but on by default).
    Addition of flexible bond and angle terms doesn't affect constrained dynamics, but allows for minimization to work properly.

    OPTIONAL ARGUMENTS

    filename (string) - name of file containing water coordinates (default: 'watbox216.pdb')
    mm (OpenMM implementation) - name of simtk.chem.openmm implementation to use (default: simtk.chem.openmm)
    constraint (boolean) - if True, will also constrain OH and HH bonds in water (default: True)

    RETURNS

    system (System)
    coordinates (numpy array)

    EXAMPLES

    Create a 216-water system.
    
    >>> [system, coordinates] = WaterBox()

    Test the energy

    >>> # Create a Context.
    >>> kB = units.BOLTZMANN_CONSTANT_kB
    >>> temperature = 298.0 * units.kelvin
    >>> kT = kB * temperature
    >>> beta = 1.0 / kT
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy()
    -7288.90384945 kJ/mol
    
    Integrate dynamics

    >>> nsteps = 100 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    if mm is None:
        mm = simtk.chem.openmm

    # Construct filename
    filename = waterbox

    # Partial atomic charges for water
    massO  = 16.0 * units.amu
    massH  =  1.0 * units.amu

    # Partial atomic charges for TIP3P water
    qO  = -0.8340 * units.elementary_charge
    qH  =  0.4170 * units.elementary_charge
    
    # Lennard-Jones parameters for oxygen-oxygen interactions
    sigma   = 3.15061 * units.angstrom
    epsilon = 0.6364  * units.kilojoule_per_mole

    # Water bond and angle values
    rOH   = 0.9572  * units.angstrom
    aHOH  = 104.52  * units.degree

    # Water bond and angle spring constants.
    kOH   = 553.0 * units.kilocalories_per_mole / units.angstrom**2 # from AMBER parm96
    kHOH  = 100.0 * units.kilocalories_per_mole / units.radian**2 # from AMBER parm96

    # Distance between the two H atoms in water
    rHH = 2*rOH*units.cos(aHOH/2.0)

    def loadCoordsHOH(infile):
        """
        Load water coordinates from a PDB file.
        
        """
        pdbData = []
        atomNum = 0
        resNum = 0
        for line in infile:
            if line.find('HETATM')==0 or line.find('ATOM')==0:
                resName=line[17:20]
                if resName=='HOH' or resName=='WAT':
                    atomNum+=1
                    atomName=line[12:16].strip()
                    if atomName=='O':
                        resNum+=1
                    try:
                        if atomName==pdbData[-1][1] or atomName==pdbData[-2][1]:
                            raise Exception("bad water molecule near %s..." % line[:27])
                    except IndexError:
                        pass
                    x = float(line[30:38])
                    y = float(line[38:46])
                    z = float(line[46:54])
                    pdbData.append( (atomNum, atomName, resName, resNum, ((x, y, z) * units.angstrom) ) )
                    
        return pdbData

    # Load waters from input pdb file
    infile = open(filename, 'r')
    pdbData = loadCoordsHOH(infile)
    infile.close()

    # Determine number of atoms.
    natoms = len(pdbData)
    
    # Create water system, which will inlcude force field
    # and general system parameters
    system = mm.System()

    # Create nonbonded, harmonic bond, and harmonic angle forces.
    nb = mm.NonbondedForce()
    bond = mm.HarmonicBondForce()
    angle = mm.HarmonicAngleForce()

    # Add water molecules to system
    # Note that no bond forces are used. Bond lenths are rigid
    count = 0
    coordinates = units.Quantity(numpy.zeros([natoms,3], numpy.float32), units.nanometer)
    for atomNum, atomName, resName, resNum, xyz in pdbData:
        if atomName=='O':
            # Add an oxygen atom
            system.addParticle(massO)
            nb.addParticle(qO, sigma, epsilon)
            lastOxygen = count
        elif atomName.startswith('H'):            
            # Add an hydrogen atom
            system.addParticle(massH)
            zero_chargeprod = 0.0 * units.elementary_charge**2
            unit_sigma = 1.0 * units.angstroms
            zero_epsilon = 0.0 * units.kilocalories_per_mole            
            nb.addParticle(qH, unit_sigma, zero_epsilon)
            if (count == lastOxygen+1):
                # For the last oxygen and hydrogen number 1:
                if constrain: system.addConstraint(lastOxygen, count, rOH)        #O-H1
                # Add harmonic bond force for this bond.
                bond.addBond(lastOxygen, count, rOH, kOH)
                # Exception: chargeProd=0.0, sigma=1.0, epsilon=0.0
                nb.addException(lastOxygen, count, zero_chargeprod, unit_sigma, zero_epsilon)   #O-H1
            elif (count == lastOxygen+2):
                # For the last oxygen and hydrogen number 2:
                if constrain: system.addConstraint(lastOxygen, count, rOH)        #O-H2
                # Add harmonic bond force for this bond.
                bond.addBond(lastOxygen, count, rOH, kOH)
                # Exception: chargeProd=0.0, sigma=1.0, epsilon=0.0
                nb.addException(lastOxygen, count, zero_chargeprod, unit_sigma, zero_epsilon)   #O-H2

                # For hydrogen number 1 and hydrogen number 2
                if constrain: system.addConstraint(count-1, count, rHH)           #H1-H2
                # Add harmonic angle constraint.
                angle.addAngle(count-1, lastOxygen, count, aHOH, kHOH)
                # DEBUG
                #bond.addBond(count-1, count, rHH, kOH)
                # Exception: chargeProd=0.0, sigma=1.0, epsilon=0.0
                nb.addException(count-1, count, zero_chargeprod, unit_sigma, zero_epsilon)      #H1-H2
            else:
                s = "too many hydrogens:"
                s += " atomNum=%d, resNum=%d, resName=%s, atomName=%s" % (atomNum, resNum, resName, atomName)
                raise Exception(s)
                sys.exit(1)
        else:
            raise Exception("bad atom : %s" % atomName)
        for k in range(3):
            coordinates[count,k] = xyz[k]        
        count += 1

    # Determine box size from maximum extent.
    box_extents = units.Quantity(numpy.zeros([3]), units.nanometers)
    for k in range(3):
        box_extents[k] = (coordinates[:,k] / units.nanometers).max() * units.nanometers - (coordinates[:,k] / units.nanometers).min() * units.nanometers
    box_size = (box_extents / units.nanometers).max() * units.nanometers 
    
    # Set box vectors.
    a = units.Quantity(numpy.zeros([3]), units.nanometers); a[0] = box_size
    b = units.Quantity(numpy.zeros([3]), units.nanometers); b[1] = box_size
    c = units.Quantity(numpy.zeros([3]), units.nanometers); c[2] = box_size
    system.setPeriodicBoxVectors(a, b, c)

    # Set nonbonded cutoff.
    nb.setNonbondedMethod(mm.NonbondedForce.CutoffPeriodic)
    if (cutoff is None) or (cutoff >= box_size / 2.0):
        cutoff = box_size / 2.0 * 0.999 # cutoff should be smaller than half the box length
        print "Setting cutoff to %s" % str(cutoff)
    nb.setCutoffDistance(cutoff)    

    # Add force terms to system.
    system.addForce(nb)
    system.addForce(bond)
    system.addForce(angle)

    return (system, coordinates)

#=============================================================================================
# Alanine dipeptide in implicit solvent
#=============================================================================================

def AlanineDipeptideImplicit(mm=None):
    """
    Alanine dipeptide ff96 in OBC GBSA implicit solvent.
    
    EXAMPLES
    
    >>> [system, coordinates] = AlanineDipeptideImplicit()

    Test the energy

    >>> # Create a Context.
    >>> kB = units.BOLTZMANN_CONSTANT_kB
    >>> temperature = 298.0 * units.kelvin
    >>> kT = kB * temperature
    >>> beta = 1.0 / kT
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy()
    -88.0885657668 kJ/mol
    
    Integrate dynamics

    >>> nsteps = 10 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    # Determine prmtop and crd filenames in test directory.
    prmtop_filename = os.path.join(os.getenv('YANK_INSTALL_DIR'), 'test', 'systems', 'alanine-dipeptide-gbsa', 'alanine-dipeptide.prmtop')
    crd_filename = os.path.join(os.getenv('YANK_INSTALL_DIR'), 'test', 'systems', 'alanine-dipeptide-gbsa', 'alanine-dipeptide.crd')

    # Initialize system.
    import amber
    system = amber.readAmberSystem(prmtop_filename, mm=mm)

    # Read coordinates.
    coordinates = amber.readAmberCoordinates(crd_filename)

    return (system, coordinates)

#=============================================================================================
# Alanine dipeptide in explicit solvent
#=============================================================================================

def AlanineDipeptideExplicit(mm=None):
    """
    Alanine dipeptide ff96 in OBC GBSA implicit solvent.
    
    EXAMPLES
    
    >>> [system, coordinates] = AlanineDipeptideExplicit()

    Test the energy

    >>> # Create a Context.
    >>> kB = units.BOLTZMANN_CONSTANT_kB
    >>> temperature = 298.0 * units.kelvin
    >>> kT = kB * temperature
    >>> beta = 1.0 / kT
    >>> collision_rate = 90.0 / units.picosecond
    >>> timestep = 1.0 * units.femtosecond    
    >>> integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    >>> context = openmm.Context(system, integrator, platform)
    >>> # Set positions
    >>> context.setPositions(coordinates)
    >>> # Evaluate the potential energy.
    >>> state = context.getState(getEnergy=True)
    >>> print state.getPotentialEnergy()
    -24520.0375793 kJ/mol
    
    Integrate dynamics

    >>> nsteps = 10 # number of steps to integrate
    >>> integrator.step(nsteps)
    >>> # Retrieve configuration to make sure no coordinates are nan
    >>> state = context.getState(getPositions=True)
    >>> coordinates = state.getPositions(asNumpy=True)
    >>> if numpy.any(numpy.isnan(coordinates / units.nanometers)): raise Exception('some coordinates are nan after integration: %s' % str(coordinates))

    """

    # Determine prmtop and crd filenames in test directory.
    prmtop_filename = os.path.join(os.getenv('YANK_INSTALL_DIR'), 'test', 'systems', 'alanine-dipeptide-explicit', 'alanine-dipeptide.prmtop')
    crd_filename = os.path.join(os.getenv('YANK_INSTALL_DIR'), 'test', 'systems', 'alanine-dipeptide-explicit', 'alanine-dipeptide.crd')

    # Initialize system.
    import amber
    system = amber.readAmberSystem(prmtop_filename, mm=mm)

    # Read coordinates.
    [coordinates, box_vectors] = amber.readAmberCoordinates(crd_filename, read_box=True)

    # Set box vectors.
    system.setPeriodicBoxVectors(box_vectors[0], box_vectors[1], box_vectors[2])
    
    return (system, coordinates)

#=============================================================================================
# MAIN AND TESTS
#=============================================================================================

if __name__ == "__main__":
    import doctest

    # Test all systems on Reference platform.
    platform = openmm.Platform.getPlatformByName("Reference")
    doctest.testmod()