<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Truncated potentials for long-range electrostatics
Björn Stenqvist and Mikael Lund, Div. of Theoretical Chemistry, Lund University, Sweden, 2016
In this notebook we test the validity of a variety of Wolf-inspired potentials and compare with a default Ewald summation on a sinple box of water.
Install prerequisites using conda
Step2:  Create initial box of water
First create a single water molecule pdb file, then use this to fill up a predefined box using the command line tool packmol. In this example, $N$ and $L$ are taken form the Gromacs spc216.gro file.
Step4:  Create OpenMM System classes for a variety of long-range correction schemes
Here the idea is to create all the different setups and add then to the systemlist dictionary. Ewald summation (and PME) can be setup using the default NonbondedForce class, while CustomNonbondedForce must be used for custom pair potentials as used in Wolf, Stenqvist etc.
For more information on how to control force objects and simulation, check,
OpenMM python API.
Lennard-Jones + Ewald at http
Step5:  Here we perform the actual MD simulation, incl. minimization of all defined systems
Step6:  Oxygen-oxygen radial distribution function
Step7:  Average potential energy
Step8:  Static dielectric constant
The exact formula for $\epsilon_r$ is not given in the mdtraj documentation, but can be found in the source code on github,
https
Step9:  Gromacs analysis
We now generate a topology (top/tpr) file for Gromacs in order to use use their analysis tools, in particular for dipol correlations. These files are identitical for all electrostatic schemes.
NOTE
Step10:  Convert trajectory to XTC format and calculate dipolar correlations
See description of the -g option in gmx dipoles, relating to http
<ASSISTANT_TASK:>
Python Code:
%matplotlib inline
from simtk.openmm.app import *
from simtk.openmm import *
from simtk.unit import *
from sys import stdout
import matplotlib, matplotlib.pyplot as plt
import pandas as pd
from io import StringIO
import numpy as np, os
matplotlib.rcParams.update({'font.size': 12})
# simulation parameters
nsteps     = 50000          # number of MD steps each of 2 fs
nwater     = 1000           # number of water molecules
cutoff     = 1.2*nanometers # pair potential cutoff
platform   = 'OpenCL'       # 'OpenCL', 'CUDA', 'CPU'
rho        = 216/18.6206**3 # density from gromacs spc216.gro
boxlen     = (nwater/rho)**(1/3)
print('box length = ', boxlen*angstrom)
print('half box   = ', boxlen*angstrom/2)
Explanation: Truncated potentials for long-range electrostatics
Björn Stenqvist and Mikael Lund, Div. of Theoretical Chemistry, Lund University, Sweden, 2016
In this notebook we test the validity of a variety of Wolf-inspired potentials and compare with a default Ewald summation on a sinple box of water.
Install prerequisites using conda:
bash
    $ conda config --add channels omnia
    $ conda install -c omnia openmm mdtraj packmol
Also, Gromacs analysis tools (gmx) must be installed for parts of the analysis.
End of explanation
%%writefile hoh.pdb
CRYST1   30.000   30.000   30.000  90.00  90.00  90.00 P 1           1
ATOM      1  OW  HOH A   1      27.552  11.051   7.172  1.00  0.00          O
ATOM      2  HW1 HOH A   1      27.900  10.721   8.050  1.00  0.00          H
ATOM      3  HW2 HOH A   1      26.606  11.355   7.281  1.00  0.00          H 
END
# write input file for packmol
PACKMOL_INPUT =  
tolerance %f
filetype pdb
output %s
# hoh will be put in a box
# defined by the minimum coordinates x, y and z = 0. 0. 0. and maximum
# coordinates box_size box_size box_size That is, they will be put in a cube of side
# box_size (the keyword "inside cube 0. 0. 0. box_size") could be used as well.
structure %s
  number %d
  inside box 0. 0. 0. %f %f %f 
  add_box_sides 0.0
end structure
 % (2.,'water.pdb','hoh.pdb', nwater, boxlen, boxlen, boxlen)
!echo '$PACKMOL_INPUT' > packmol_input.txt
!packmol < packmol_input.txt > /dev/null
Explanation: Create initial box of water
First create a single water molecule pdb file, then use this to fill up a predefined box using the command line tool packmol. In this example, $N$ and $L$ are taken form the Gromacs spc216.gro file.
End of explanation
elec_to_kJmol = (
    constants.elementary_charge**2 *
    AVOGADRO_CONSTANT_NA / (4*np.pi*1.0*8.854187817e-12
                            * constants.farad/constants.meter)).value_in_unit(kilojoule_per_mole*nanometer)
def findForce(system, forcetype, add=True):
     Finds a specific force in the system force list - added if not found.
    for force in system.getForces():
        if isinstance(force, forcetype):
            return force
    if add==True:
        system.addForce(forcetype())
        return findForce(system, forcetype)
    return None
def make_qpotential_system(topology, Rc=0.9*nanometers, moments=1000):
    ''' return a q-potential system (TO '''
    
    ff = ForceField('spce-custom.xml') # this will create a CustomNonbondedForce
    system = ff.createSystem(
        topology, nonbondedMethod=CutoffPeriodic,
        nonbondedCutoff=Rc, constraints=HBonds, rigidWater=True)
    
    def qPochhammerSymbol( Rc, moments ):
        if isinstance(Rc, Quantity):
            Rc = Rc / nanometer # strip unit
        qP = 1.0
        r = np.linspace(0, Rc, 5000)
        for i in range( moments ):
            qP *= (1 - (r/Rc)**(i+1) )
        return qP
 
    qP = Continuous1DFunction( qPochhammerSymbol(Rc, moments), 0*nanometers, Rc)
    nonbonded = findForce(system, CustomNonbondedForce)
    nonbonded.addTabulatedFunction( 'qP', qP )         # 'qP(r)' can now be used in energy function
    nonbonded.addGlobalParameter( 'f', elec_to_kJmol ) # convert to kJ/mol
    
    nonbonded.setEnergyFunction(
        'f * charge1 * charge2 * qP(r)/r' \
        ' + 4 * epsilon * ( (sigma/r)^12 - (sigma/r)^6 )' \
        ' ; sigma = 0.5 * ( sigma1+sigma2 ); epsilon = sqrt( epsilon1*epsilon2 )'
    )
    print('qpot')
    print('    periodic boundaries:  ', nonbonded.usesPeriodicBoundaryConditions())
    print('    switching function:   ', nonbonded.getUseSwitchingFunction())
    print('    long-range correction:', nonbonded.getUseLongRangeCorrection())
    print('    cutoff distance:      ', nonbonded.getCutoffDistance())
    print('    energy function:      ', nonbonded.getEnergyFunction())
    
    return system
def make_grekiss_system(topology, Rc=0.9*nanometers):
    ''' return a grekis system '''
    
    ff = ForceField('spce-custom.xml') # this will create a CustomNonbondedForce
    system = ff.createSystem(
        topology, nonbondedMethod=CutoffPeriodic,
        nonbondedCutoff=Rc, constraints=HBonds, rigidWater=True)
    
    nonbonded = findForce(system, CustomNonbondedForce)
    nonbonded.addGlobalParameter( 'Rc', Rc )           # 'Rc' can now be used in energy function
    nonbonded.addGlobalParameter( 'f', elec_to_kJmol ) # 'lB' bjerrum length in nm and kJ/mol
    
    nonbonded.setEnergyFunction(
        'f * charge1 * charge2 * 1/r * ( 1 - 1.75*r/Rc + 5.25*(r/Rc)^5 - 7*(r/Rc)^6 + 2.5*(r/Rc)^7 )' \
        ' + 4 * epsilon * ( (sigma/r)^12 - (sigma/r)^6 )' \
        ' ; sigma = 0.5 * (sigma1+sigma2); epsilon = sqrt( epsilon1*epsilon2 )'
    )
    print('fanourgakis')
    print('    periodic boundaries:  ', nonbonded.usesPeriodicBoundaryConditions())
    print('    switching function:   ', nonbonded.getUseSwitchingFunction())
    print('    long-range correction:', nonbonded.getUseLongRangeCorrection())
    print('    cutoff distance:      ', nonbonded.getCutoffDistance())
    print('    energy function:      ', nonbonded.getEnergyFunction())
    
    return system
def make_yonezawa_system(topology, Rc=0.9*nanometers):
    ''' return a yonezawa system '''
    
    ff = ForceField('spce-custom.xml') # this will create a CustomNonbondedForce
    system = ff.createSystem(
        topology, nonbondedMethod=CutoffPeriodic,
        nonbondedCutoff=Rc, constraints=HBonds, rigidWater=True)
    nonbonded = findForce(system, CustomNonbondedForce)
    nonbonded.addGlobalParameter( 'Rc', Rc )           # 'Rc' can now be used in energy function
    nonbonded.addGlobalParameter( 'f', elec_to_kJmol ) # 'lB' bjerrum length in nm and kJ/mol
    
    nonbonded.setEnergyFunction(
        'f * charge1 * charge2 * ( 1/r - 1/Rc + 1/Rc^2 * (r-Rc) )' \
        ' + 4 * epsilon * ( (sigma/r)^12 - (sigma/r)^6 )' \
        ' ; sigma = 0.5 * ( sigma1+sigma2 ); epsilon = sqrt( epsilon1 * epsilon2 )'
    )
    print('yonezawa')
    print('    periodic boundaries:  ', nonbonded.usesPeriodicBoundaryConditions())
    print('    switching function:   ', nonbonded.getUseSwitchingFunction())
    print('    long-range correction:', nonbonded.getUseLongRangeCorrection())
    print('    cutoff distance:      ', nonbonded.getCutoffDistance())
    print('    energy function:      ', nonbonded.getEnergyFunction())
    
    return system
def make_ewald_system(topology, Rc=0.9*nanometers, method=Ewald, ljcorr=False):
    ''' returns an Ewald system '''
    ff = ForceField('spce.xml') # this will create a NonbondedForce
    
    system = ff.createSystem(
        topology, nonbondedMethod=method,
        nonbondedCutoff=Rc, constraints=HBonds, rigidWater=True)
    
    nonbonded = findForce(system, NonbondedForce)
    nonbonded.setUseDispersionCorrection( ljcorr )
    print('ewald')
    print('    err. tolerance:       ', nonbonded.getEwaldErrorTolerance())
    print('    LJ switching function:', nonbonded.getUseSwitchingFunction())
    print('    LJ correction:        ', nonbonded.getUseDispersionCorrection())
    print('    cutoff distance:      ', nonbonded.getCutoffDistance())
 
    return system
pdb = PDBFile('water.pdb')
# here we define our list of system incl. plotting properties. All files from simulation
# and analysis will be prefixed with 'qpot.', 'ewald.', etc.
systemlist = {
    'yonezawa': {
        'system': make_yonezawa_system( pdb.topology, Rc=cutoff ), 'npt': False,
        'run': False, 'color': 'blue', 'label': 'Yonezawa'
    },
    'grekiss': {
        'system': make_grekiss_system( pdb.topology, Rc=cutoff ), 'npt': False,
        'run': False, 'color': 'green', 'label': 'Fanourgakis'
    },
    'qpot' : {
        'system': make_qpotential_system( pdb.topology, Rc=cutoff ), 'npt': False,
        'run': False, 'color': 'red', 'label': '$q$ ($n=\infty$)'
    },
    'qpot2' : {
        'system': make_qpotential_system( pdb.topology, Rc=cutoff, moments=2 ), 'npt': False,
        'run': False, 'color': 'orange', 'label': '$q$ ($n=2$)'
    },
    'qpot3' : {
        'system': make_qpotential_system( pdb.topology, Rc=cutoff, moments=3 ), 'npt': False,
        'run': False, 'color': 'magenta', 'label': '$q$ ($n=3$)'
    },
    'ewald': {
        'system': make_ewald_system( pdb.topology, Rc=cutoff ), 'npt': False,
        'run': False, 'color': 'black', 'label': 'Ewald'
    }
}
Explanation: Create OpenMM System classes for a variety of long-range correction schemes
Here the idea is to create all the different setups and add then to the systemlist dictionary. Ewald summation (and PME) can be setup using the default NonbondedForce class, while CustomNonbondedForce must be used for custom pair potentials as used in Wolf, Stenqvist etc.
For more information on how to control force objects and simulation, check,
OpenMM python API.
Lennard-Jones + Ewald at http://docs.openmm.org/7.0.0/userguide/theory.html#nonbondedforce
Long range correction beyond cut-off: http://docs.openmm.org/7.0.0/userguide/theory.html#customnonbondedforce
TODO: Smoothen LJ part at $R_c$. Use switching function?
End of explanation
for name, prop in systemlist.items(): # loop over systems
    
    if prop['run']:
        print(name)
        if prop['npt']: # NVT -> NPT ensemble ?
            print('    adding barostat.')
            barostat = MonteCarloBarostat(1.0*bar, 298.15*kelvin, 25) 
            system.addForce(barostat)
        integrator = LangevinIntegrator( 298.15*kelvin, 1.0/picoseconds, 2*femtoseconds )
        integrator.setConstraintTolerance(0.00001)
        _platform = Platform.getPlatformByName( platform ) # CPU, OpenCL
        sim = Simulation(pdb.topology, prop['system'], integrator, _platform)
        sim.context.setPositions(pdb.positions) # set particle positions
        if os.path.isfile( name+'.chk' ):
            with open( name+'.chk', 'rb') as f:
                print('    loading restart file.')
                sim.context.loadCheckpoint( f.read() )
        else:
            print('    minimizing energy...')
            sim.reporters.clear()
            %time sim.minimizeEnergy( tolerance=50*kilojoule/mole, maxIterations=1000 )
            sim.context.setVelocitiesToTemperature( 298.15*kelvin ) # initial random velocities
        print('    running Production...')
        sim.reporters.clear()
        sim.reporters.append( DCDReporter( name+'.dcd', 1000) )
        sim.reporters.append( StateDataReporter(name+'.energy', 1000, step=True, potentialEnergy=True,
                                                temperature=True, density=True) )
        sim.reporters.append( StateDataReporter(stdout, 1000, step=True, potentialEnergy=True,
                                                temperature=True, density=True, separator='\t',
                                                progress=True,
                                                totalSteps = nsteps) )
        %time sim.step( nsteps )
        with open( name+'.chk', 'wb') as f:
            print('    saving restart file.')
            f.write( sim.context.createCheckpoint() )
        # save final configuration to PDB file
        # todo: something is fishy here... water molecules are outside the box
        #       in pdb although the trajectory looks fine. Å->nm problem?
        positions = sim.context.getState( getPositions=True ).getPositions()
        PDBFile.writeFile(sim.topology, positions, open( name+'.pdb', 'w'))
        print()
Explanation: Here we perform the actual MD simulation, incl. minimization of all defined systems
End of explanation
import mdtraj as md
for name, prop in systemlist.items():
    print(name)
    
    traj = md.load(name+'.dcd', top=name+'.pdb')
    sel = traj.top.select('name O')
    OOpairs = traj.top.select_pairs('name O', 'name O')
    rdf = md.compute_rdf( traj, pairs=OOpairs, bin_width=0.005, r_range=[0.2, boxlen/2/10] )
    
    plt.plot( rdf[0], rdf[1], label=prop['label'], color=prop['color'], lw=1)
plt.xlim(0.2, 0.8)
plt.xlabel(r'$r$/nm')
plt.ylabel(r'$g(r)$')
plt.legend(loc=0, frameon=False, fontsize='medium')
Explanation: Oxygen-oxygen radial distribution function
End of explanation
for name, prop in systemlist.items():
    step, Upot, temp, density = np.loadtxt(name+'.energy', skiprows=1, delimiter=',', unpack=True)
    Upot = Upot / pdb.topology.getNumResidues() # -> energy per water molecule
    plt.plot( step, Upot, label=prop['label'], color=prop['color'], lw=0.5 )
    plt.plot( [0,step.max()], [Upot.mean(), Upot.mean()], color=prop['color'], ls='--', lw=0.5 )
    prop['energy'] = Upot.mean()
plt.xlabel(r'steps')
plt.ylabel(r'$U_{pot}/N$ (kJ/mol)')
plt.legend(loc=1, frameon=True, fontsize='small')    
Explanation: Average potential energy
End of explanation
def getChargeVector( force, stripunit=True ):
    ''' Extract particle charge array from force object if available
    
    Note that this function assumes that charge is located at parameter
    position 0 which is default in NonbondedForce while it can be anywhere
    in CustomNonbondedForce, following the order in which
    parameters were added in python/xml.
    '''
    if isinstance(force, (NonbondedForce, CustomNonbondedForce) ):
        if 'getParticleParameters' in dir(force):
            chargelist = []
            for i in range( force.getNumParticles() ):
                charge = force.getParticleParameters( i )[ 0 ]
                if stripunit:
                    if isinstance(charge, Quantity):
                        charge = charge / elementary_charge # strip unit
                chargelist.append( charge )
            return chargelist
    return None
    
for name, prop in systemlist.items():
    for force in prop['system'].getForces():
        if isinstance(force, (NonbondedForce, CustomNonbondedForce) ):
            charges = getChargeVector( force )
            traj = md.load(name+'.dcd', top=name+'.pdb')
            l = len(traj)
            l2= int(l/2)
            
            eps12= md.static_dielectric( traj[1:l], charges, temperature=298.15)  # all
            eps1 = md.static_dielectric( traj[1:l2], charges, temperature=298.15) # 1st half
            eps2 = md.static_dielectric( traj[l2:l], charges, temperature=298.15) # 2nd half
            
            prop['diel'] = eps12
            
            print(name, ': diel. const.', eps12, eps1, eps2)
pd.DataFrame(systemlist).T
Explanation: Static dielectric constant
The exact formula for $\epsilon_r$ is not given in the mdtraj documentation, but can be found in the source code on github,
https://github.com/mdtraj/mdtraj/blob/master/mdtraj/geometry/thermodynamic_properties.py#L90
End of explanation
%%bash
rm -fR 
echo -n "6\n" | gmx -quiet -nobackup pdb2gmx -f water.pdb -o water.gro -p water.top -n water.ndx -water spce
gmx -quiet -nobackup grompp -f water.mdp -c water.gro -p water.top -o water.tpr -maxwarn 10
Explanation: Gromacs analysis
We now generate a topology (top/tpr) file for Gromacs in order to use use their analysis tools, in particular for dipol correlations. These files are identitical for all electrostatic schemes.
NOTE: Depending on the size of the water box, grompp may complain that the cut-off is too short. Fix this by increasing the verlet tolerance. It will not matter for the analysis, but consider making a slightly larger system.
End of explanation
for name, prop in systemlist.items():
    print(name)
    
    traj = md.load(name+'.dcd', top=name+'.pdb')
    traj.save_xtc(name+'.xtc')
    !echo -n "0\n" | gmx -quiet -nobackup dipoles -f $name'.xtc' -s water.tpr -temp 298.15 -g $name'.gkr.xvg'
    print()
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(nrows=2, ncols=2, figsize=(8, 7))
for name, prop in systemlist.items():
    r, G, cos, hOO, gOO, energy = np.loadtxt(name+'.gkr.xvg', skiprows=27, unpack=True)
    
    ax1.plot(r, G,   label=prop['label'], color=prop['color'], lw=1 )
    ax1.set_xlabel('$r$/nm')
    ax1.set_ylabel('$G_k(r)$')
    ax2.plot(r, cos, label=prop['label'], color=prop['color'], lw=1 )
    ax2.set_xlabel('$r$/nm')
    ax2.set_ylabel('cosine something...')
    ax2.set_xlim(0.2, 0.8)
    ax2.set_ylim(-0.1, 0.55)
    
    ax3.plot(r, hOO, label=prop['label'], color=prop['color'], lw=1 )
    ax3.set_xlabel('$r$/nm')
    ax3.set_ylabel('hOO')
    ax3.set_xlim(0.2, 0.8)
    
    ax4.plot(r, gOO, label=prop['label'], color=prop['color'], lw=1 )
    ax4.set_xlabel('$r$/nm')
    ax4.set_ylabel('gOO')
    ax4.set_xlim(0.2, 0.8)
    
    ax2.legend(loc=0, frameon=False, fontsize='medium')
plt.tight_layout()
Explanation: Convert trajectory to XTC format and calculate dipolar correlations
See description of the -g option in gmx dipoles, relating to http://manual.gromacs.org/programs/gmx-dipoles.html (Nymand/Linse).
End of explanation
<END_TASK>