<html><head><meta name="color-scheme" content="light dark"></head><body><pre style="word-wrap: break-word; white-space: pre-wrap;">#!/usr/local/bin/env python

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

"""
Compare platforms using different tolerances.

DESCRIPTION

COPYRIGHT

@author John D. Chodera &lt;jchodera@gmail.com&gt;

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 &lt;http://www.gnu.org/licenses/&gt;.

TODO

"""

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

import os
import os.path
import sys
import math
import copy
import numpy
import time

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

#=============================================================================================
# SUBROUTINES
#=============================================================================================

def setIonCharge(system, atom_index, new_charge):
    """
    Set ion charge to specified value.

    """
    for force_index in range(system.getNumForces()):
        force = system.getForce(force_index)
        if hasattr(force, 'getParticleParameters'):
            [old_charge, sigma, epsilon] = force.getParticleParameters(atom_index)
            force.setParticleParameters(atom_index, new_charge, sigma, epsilon)
            break
        
    return        

def equilibrate(system, coordinates, platform):
    print "Equilibrating..."

    # Create integrator and context.
    integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    context = openmm.Context(system, integrator, platform)

    # Set coordinates.
    context.setPositions(coordinates)

    # Equilibrate.
    integrator.step(10000)

    # Compute energy
    print "Computing energy."
    state = context.getState(getEnergy=True, getPositions=True)
    potential_energy = state.getPotentialEnergy()

    # Get coordinates.
    coordinates = state.getPositions(asNumpy=True)
    print "Done"
    
    return [coordinates, potential_energy]

def compute_potential(system, coordinates, platform):
    # Create integrator and context.
    integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    context = openmm.Context(system, integrator, platform)

    # Set coordinates.
    context.setPositions(coordinates)

    # Compute energy
    state = context.getState(getEnergy=True)
    potential_energy = state.getPotentialEnergy()

    return potential_energy

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

if __name__ == "__main__":

    # Parameters
    prmtop_filename = 'system.prmtop'
    crd_filename = 'system.crd'

    temperature = 298.0 * units.kelvin
    collision_rate = 5.0 / units.picoseconds
    timestep = 1.0 * units.femtoseconds

    # List all available platforms
    print "Available platforms:"
    for platform_index in range(openmm.Platform.getNumPlatforms()):
        platform = openmm.Platform.getPlatform(platform_index)
        print "%5d %s" % (platform_index, platform.getName())
    print ""

    # Select platform.
    equilibration_platform = openmm.Platform.getPlatformByName("Cuda")
    platform = openmm.Platform.getPlatformByName("Cuda")    

    # Create system.
    print "Reading system..."
    cutoff = 9.0 * units.angstroms
    solvent_dielectric = 78.5
    #system = amber.readAmberSystem(prmtop_filename, nonbondedMethod='reaction-field', nonbondedCutoff=cutoff, shake='h-bonds')
    system = amber.readAmberSystem(prmtop_filename, nonbondedMethod='reaction-field', nonbondedCutoff=cutoff, shake='none')    
    [coordinates, box_vectors] = amber.readAmberCoordinates(crd_filename, read_box=True)
    # make box isotropic
    box_length = 0.0 * units.angstroms
    for k in range(3):
        if box_vectors[k][k] &gt; box_length:
            box_length = box_vectors[k][k]
    for k in range(3):
        box_vectors[k][k] = box_length
    system.setDefaultPeriodicBoxVectors(box_vectors[0], box_vectors[1], box_vectors[2])

    # Test energy.
    print "potential energy of original system: %.3f kcal/mol" % (compute_potential(system, coordinates, platform) / units.kilocalories_per_mole)
    
    # Zero ion charge.
    #setIonCharge(system, 0, 1.0 * units.elementary_charge)

    # Find NonbondedForce.
    nonbonded_force = None
    for force_index in range(system.getNumForces()):
        force = system.getForce(force_index)
        if hasattr(force, 'getParticleParameters'):
            nonbonded_force = force
            break

    # Add charge interaction back in with CustomBondForce.
    #energy_expression = '(C*ion_charge*charge/r)*step(cutoff-r)' # Coulomb interaction with cutoff
    energy_expression = 'k_e*ion_charge*charge*(r^(-1) + k_rf*r^2 - c_rf)*step(cutoff-r); k_rf = cutoff^(-3)*(epsilon_solvent-1.0)/(2.0*epsilon_solvent+1.0); c_rf = cutoff^(-1)*(3.0*epsilon_solvent)/(2.0*epsilon_solvent+1.0)' # reaction-field with cutoff
    epsilon_solvent = nonbonded_force.getReactionFieldDielectric()
    force = openmm.CustomBondForce(energy_expression)
    force.addGlobalParameter('k_e', 332.063711 * units.kilocalories_per_mole * units.angstroms / units.elementary_charge**2) # 1 / (4 pi epsilon0)
    force.addGlobalParameter('ion_charge', 1.0 * units.elementary_charge)
    force.addGlobalParameter('epsilon_solvent', epsilon_solvent)
    force.addGlobalParameter('cutoff', cutoff)
    force.addGlobalParameter('epsilon_solvent', solvent_dielectric)
    force.addPerBondParameter('charge')
    for atom_index in range(1, system.getNumParticles()):
        [charge, sigma, epsilon] = nonbonded_force.getParticleParameters(atom_index)
        force.addBond(0, atom_index, [charge])
    system.addForce(force)

    # Test energy.
    print "potential energy of modified system: %.3f kcal/mol" % (compute_potential(system, coordinates, platform) / units.kilocalories_per_mole)
    
    # Generate work samples.
    outfile = open('work.out', 'w') # file for writing work values
    niterations = 1000 # number of work samples to collect
    nsteps_to_try = [1, 5, 10, 50, 100, 500, 1000, 5000] # number of steps per switch    
    for iteration in range(niterations):
        print "iteration %5d / %5d" % (iteration, niterations)

        # Equilibrate.
        [coordinates, potential] = equilibrate(system, coordinates, equilibration_platform)
        outfile.write("%12.3f" % (potential / units.kilocalories_per_mole))        

        for nsteps in nsteps_to_try:
            initial_time = time.time()
            
            # Create context.
            integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)    
            context = openmm.Context(system, integrator, platform)
            context.setPositions(coordinates)

            # Accumulate switching work.
            work = 0.0 * units.kilocalories_per_mole
            for step in range(nsteps):
                # Compute old energy.
                context.setParameter('ion_charge', (1.0 - 2.0*float(step)/float(nsteps)) * units.elementary_charge)
                old_state = context.getState(getEnergy=True)
                old_energy = old_state.getPotentialEnergy()
                
                # Compute new energy.
                context.setParameter('ion_charge', (1.0 - 2.0*float(step+1)/float(nsteps)) * units.elementary_charge)
                new_state = context.getState(getEnergy=True)
                new_energy = new_state.getPotentialEnergy()
                
                # Evolve dynamics.
                integrator.step(1)
                
                # Accumulate work.
                delta_energy = new_energy - old_energy
                work += delta_energy
                #print "step %5d / %5d : energy = %8.1f kcal/mol, delta_energy = %8.1f kcal/mol, accumulated work = %8.1f kcal/mol" % (step+1, nsteps, new_energy / units.kilocalories_per_mole, delta_energy / units.kilocalories_per_mole, work / units.kilocalories_per_mole)
                pass
                
            outfile.write("%12.3f" % (work / units.kilocalories_per_mole))
            final_time = time.time()
            elapsed_time = final_time - initial_time
            time_per_step = elapsed_time / nsteps
            print "%d steps took %.3f s (%.1f ms per step)" % (nsteps, elapsed_time, time_per_step * 1000)
            
        outfile.write("\n")
        outfile.flush()


    outfile.close()
</pre></body></html>