Python API to evaluate ewald Coulomb force

[JZ Most]

Hi guys, I am a new user of OpenMM, I am trying to use the OpenMM via the python API to evaluate the long range coulomb force by PME or ewald summation. Below is my code, I am evaluating a system of 60 Na and 50 Cl. However I found this code is even much slower than Pymatgen (written in Python). May I ask how I can fix my code?


import numpy as np
from simtk import openmm as mm
from simtk.openmm import app
from simtk.unit import nanometer, picosecond

positions = np.random.rand(110, 3) * 1000

# define the box size
box_size = 1000.0

# define a new system
system = mm.System()
system.setDefaultPeriodicBoxVectors((box_size, 0, 0), (0, box_size, 0), (0, 0, box_size))

# define the ion number and charge
num_Na = 60
num_Cl = 50
charge_Na = +1.0 # +
charge_Cl = -1.0 # -

# charge array
charges = np.concatenate([np.full(num_Na, charge_Na), np.full(num_Cl, charge_Cl)])

# add ions to the system
for _ in range(num_Na + num_Cl):
system.addParticle(1.0) # assume the mass is 1 for simplicity

# NonbondedForce, Particle Mesh Ewald
nonbonded = mm.NonbondedForce()
nonbonded.setNonbondedMethod(mm.NonbondedForce.PME)
nonbonded.setEwaldErrorTolerance(1.0e-1)
nonbonded.setCutoffDistance(5.0)

# add the charge to the system and ignore VDW
for charge in charges:
nonbonded.addParticle(charge, 0, 0)

# add the force to the system
system.addForce(nonbonded)

# create simulation
dummy_integrator = mm.CustomIntegrator(0)
simulation = app.Simulation(topology=None, system=system, integrator=dummy_integrator)
simulation.context.setPositions(positions)

# evaluate the force
state = simulation.context.getState(getForces=True)
forces = state.getForces(asNumpy=True)

print(forces)

[Peter Eastman]

What are you timing? The whole script, or just the call to getState()?

All the setup is much more expensive than the force evaluation. Creating a Simulation involves a lot of work, but it only has to be done once. The very first call to getState() often also has to do additional initialization. That makes it take longer than all subsequent calls.

The cost of these things depends on which platform you're using. Try all the different platforms and see which is best for your application. The Reference platform does minimal setup, so creating a Simulation is very fast, but it's the slowest at computing forces. The CPU platform does more initial work but then is faster at computing forces. The CUDA and OpenCL platforms have to compile GPU kernels, so the setup is quite expensive, but then they're potentially much faster, at least for a large system and a good GPU.

You don't want to set the Ewald error tolerance to 0.1. That's a huge error. It won't give meaningful results.

[JZ Most]

Hi Peter,

Thanks for replying. I am timing the whole script but certainly the simulation.context.getState(getForces=True) is the only time consuming line. I use a normal PC to run this code, of which cost 2.5 second to complete a Pymatgen ewald evaluation. Openmm is written in C++, it should be much faster than Pymatgen. I also know that 0.1 tolerance make no sense but if I use default tolerance the code will crush and will never get a result. So I think this is not an issue of platform, I can find where the coding mistake is.

[Peter Eastman]

Another issue is your box size. 1000 nm is huge. The cost of PME scales with the size of the box. Do you really intend it to be that big?

[JZ Most]

My system size is 1000 angstrom actually, which is equivalent to 100 nm. My code can be changed to that size but no difference showed.
I changed my code alittle bit, and it can run very fastly, only <0.1s, but I am not sure it make sense or not. I changed the size to 100A, cutoff to 10A, tolerance to 1.0e-4.

ositions = np.random.rand(110, 3) * 100*angstrom

# define the box size
box_size = 100*angstrom

# define a new system
system = mm.System()
system.setDefaultPeriodicBoxVectors((box_size, 0, 0), (0, box_size, 0), (0, 0, box_size))

# define the ion number and charge
num_Na = 60
num_Cl = 50
charge_Na = +1.0 # +
charge_Cl = -1.0 # -

# charge array
charges = np.concatenate([np.full(num_Na, charge_Na), np.full(num_Cl, charge_Cl)])

# add ions to the system
for _ in range(num_Na + num_Cl):
system.addParticle(1.0) # assume the mass is 1 for simplicity

# NonbondedForce, Particle Mesh Ewald
nonbonded = mm.NonbondedForce()
nonbonded.setNonbondedMethod(mm.NonbondedForce.PME)
nonbonded.setEwaldErrorTolerance(1.0e-4)
nonbonded.setCutoffDistance(10*angstrom)

# add the charge to the system and ignore VDW
for charge in charges:
nonbonded.addParticle(charge, 0, 0)

# add the force to the system
system.addForce(nonbonded)

# create simulation
dummy_integrator = mm.CustomIntegrator(0)
simulation = app.Simulation(topology=None, system=system, integrator=dummy_integrator)
simulation.context.setPositions(positions)
state = simulation.context.getState(getForces=True)
forces = state.getForces(asNumpy=True)

[Peter Eastman]

I modified your script to loop over all platforms and time the force evaluation for each one.

Code: Select all

import numpy as np
import openmm as mm
from openmm import app
from openmm.unit import angstrom
import time

positions = np.random.rand(110, 3) * 100*angstrom

# define the box size
box_size = 100*angstrom

# define a new system
system = mm.System()
system.setDefaultPeriodicBoxVectors((box_size, 0, 0), (0, box_size, 0), (0, 0, box_size))

# define the ion number and charge
num_Na = 60
num_Cl = 50
charge_Na = +1.0 # +
charge_Cl = -1.0 # -

# charge array
charges = np.concatenate([np.full(num_Na, charge_Na), np.full(num_Cl, charge_Cl)])

# add ions to the system
for _ in range(num_Na + num_Cl):
    system.addParticle(1.0) # assume the mass is 1 for simplicity

# NonbondedForce, Particle Mesh Ewald
nonbonded = mm.NonbondedForce()
nonbonded.setNonbondedMethod(mm.NonbondedForce.PME)
nonbonded.setEwaldErrorTolerance(1.0e-4)
nonbonded.setCutoffDistance(10*angstrom)

# add the charge to the system and ignore VDW
for charge in charges:
    nonbonded.addParticle(charge, 0, 0)

# add the force to the system
system.addForce(nonbonded)

for i in range(mm.Platform.getNumPlatforms()):
    # create simulation
    platform = mm.Platform.getPlatform(i)
    dummy_integrator = mm.CustomIntegrator(0)
    simulation = app.Simulation(topology=None, system=system, integrator=dummy_integrator, platform=platform)
    simulation.context.setPositions(positions)
    state = simulation.context.getState(getForces=True)
    t1 = time.time()
    state = simulation.context.getState(getForces=True)
    forces = state.getForces(asNumpy=True)
    t2 = time.time()
    print(platform.getName(), t2-t1)

Here is the output on my laptop.

Code: Select all

Reference 0.02939319610595703
CPU 0.005501747131347656
OpenCL 0.0022039413452148438

All platforms take only a tiny fraction of a second.

[JZ Most]

Hi Peter,

Thanks for your replying, now the code can run very fastly. However, the results did not match with my PDE simulation, I am not sure which part was wrong. So I am trying to implement all the Brownian dynamic simulation (including LJ force and random diffusion) via Openmm. My problem setting is 50 free Cl- and 50 Na+ in a 1000 angstrom box, with extra 10 Na+ fixed in the x = 0 surface, in total would be 60 Na+ and 50 Cl-. I have some problems with the codes, for example I cant fix the 10 Na+ in the x = 0 surface.

Code: Select all

# initial
rho_perm = np.random.rand(10, 3) * 1e-7
rho_perm[:,0] = 0
positions = np.random.rand(50*2, 3) * 1e-7
positions = np.concatenate((positions, rho_perm), axis=0)

# system para
num_Na_fixed = 10
num_Na_free = 50
num_Cl = 50
num_Na = num_Na_fixed + num_Na_free 

# build the system
system = mm.System()

box_size = 100.0 * nanometer

system.setDefaultPeriodicBoxVectors((box_size, 0, 0), (0, box_size, 0), (0, 0, box_size))

# Cl and Na mass
for _ in range(num_Cl):
    system.addParticle(35.45)  

for _ in range(num_Na):
    system.addParticle(22.99) 

# NonbondedForce
nonbonded = mm.NonbondedForce()
nonbonded.setNonbondedMethod(mm.NonbondedForce.PME)
nonbonded.setCutoffDistance(30.0 * nanometer)

# LJ parameter
for _ in range(num_Cl):
    nonbonded.addParticle(-1.0, 0.227 * nanometer, 0.15*4.184 * kilojoule_per_mole)

for _ in range(num_Na):
    nonbonded.addParticle(1.0, 0.141 * nanometer, 0.0469*4.184 * kilojoule_per_mole)

system.addForce(nonbonded)

# CustomExternalForce to fix the 10 Na
force = mm.CustomExternalForce('k*periodicdistance(x, y, z, x0, y0, z0)^2')
force.addGlobalParameter('k', 1000000000.0 * kilocalorie_per_mole/angstrom**2) 
force.addPerParticleParameter('x0')
force.addPerParticleParameter('y0')
force.addPerParticleParameter('z0')
for i in range(num_Na_fixed):
    force.addParticle(i + num_Cl + num_Na_free, positions[i + num_Cl + num_Na_free])
system.addForce(force)

integrator = mm.LangevinIntegrator(300 * kelvin, 1 / picosecond, 0.1 * picosecond)
simulation = app.Simulation(topology=None, system=system, integrator=integrator)
simulation.context.setPositions(positions*1e+10*angstrom)

for i in range(1000000):
    simulation.step(1)
    state = simulation.context.getState(getPositions=True)
    positions_array = state.getPositions()
    positions_array = np.array(positions_array.value_in_unit(nanometer))
    positions_array%=100
    
    for i in range(50):
        index1 = int(positions_array[i,0]/(5))
        index2 = int(positions_array[i,1]/(5))
        index3 = int(positions_array[i,2]/(5))
        pos_total[index1,index2,index3] += 1

    for i in range(50):
        index1 = int(positions_array[i+50,0]/(5))
        index2 = int(positions_array[i+50,1]/(5))
        index3 = int(positions_array[i+50,2]/(5))
        neg_total[index1,index2,index3] += 1  

[Peter Eastman]

If you want the particles to be completely immobile, you can just set their masses to zero. That will prevent them from moving at all.

[JZ Most]

Hi Peter,

How are you, I have now looked back on this problem. And recently I found some problems about the computation.

Code: Select all

box_size = 100*angstrom

# define a new system
system = mm.System()

system.setDefaultPeriodicBoxVectors((box_size, 0, 0), (0, box_size, 0), (0, 0, box_size))

# define the ion number and charge
num_Na = num_particles+36
num_Cl = num_particles+36
charge_Na = +1.0 # +
charge_Cl = -1.0 # -
mass_Na = 22.98976
mass_Cl = 35.453

# charge array
charges = np.concatenate([np.full(num_Cl, charge_Cl), np.full(num_Na, charge_Na)])

# add ions to the system
for _ in range(num_Cl):
    system.addParticle(mass_Cl) # assume the mass is 1 for simplicity
for _ in range(num_Na):
    system.addParticle(mass_Na)
    
# NonbondedForce, Particle Mesh Ewald
nonbonded = mm.NonbondedForce()
nonbonded.setNonbondedMethod(mm.NonbondedForce.PME)
nonbonded.setEwaldErrorTolerance(1.0e-4)
nonbonded.setCutoffDistance(50*angstrom)

# Set PME parameters (alpha, reciprocal space cutoff)
alpha = 0.052/nanometer
grid_size_x = 17  # Example FFT grid size for x dimension
grid_size_y = 17  # Example FFT grid size for y dimension
grid_size_z = 17  # Example FFT grid size for z dimension

# Set PME parameters
nonbonded.setPMEParameters(alpha, grid_size_x, grid_size_y, grid_size_z)

# add the charge to the system and consider lennerd-jones force
for _ in range(num_Cl):
    nonbonded.addParticle(-1, 0.227 * nanometer, 0.15*4.184 * kilojoule_per_mole)
for _ in range(num_Na):
    nonbonded.addParticle(+1, 0.141 * nanometer, 0.0469*4.184 * kilojoule_per_mole)

# add the force to the system
system.addForce(nonbonded)


# create simulation

dummy_integrator = mm.CustomIntegrator(0)
simulation = app.Simulation(topology=None, system=system, integrator=dummy_integrator, platform=platform)
simulation.context.setPositions(positions*1e+10*angstrom)


state = simulation.context.getState(getForces=True)
forces = state.getForces(asNumpy=True)
t2 = time.time()
print(platform.getName(), t2-t1)
    

I defined a cubic system (L = 1e-8m), with 2 layer of fixed charges at x = 0 and x = L/2. I tried to compute the corresponding PME Coulombic and LJ force at each particles at every timestep. However, I found that if I try to use state.getPositions(asNumpy=True) to see the positions, the return value is error or very nonsense value (e.g. 1e-300), does that mean the force I got is also wrong?
If you have time could you please have a look?

Best wishes,
Jinyuan

[Peter Eastman]

simulation.context.setPositions(positions*1e+10*angstrom)

Is that really what you mean? I don't know what the values in positions are, but unless they're all tiny values on the order of 1e-10, you're setting the positions to huge values.