Electrostatics in Julia

Molly.jl now has PME code, compatible with CPU and GPU backends via KernelAbstractions.jl: Molly.jl/src/interactions/ewald.jl at master · JuliaMolSim/Molly.jl · GitHub.

The implementation is based on OpenMM and seems to be decently performant, taking < 1 ms per step for 15k atoms on a RTX 4090 GPU. It is differentiable with Enzyme.

It is rather tied into the Molly infrastructure since the neighbour list is used to calculate the short-range contributions. Here is an example:

using Molly # Requires v0.23.0
using GLMakie

n_atoms = 100
atoms = [Atom(mass=10.0u"g/mol", charge=(i % 2 == 0 ? 0.3 : -0.3))
         for i in 1:n_atoms]

boundary = CubicBoundary(2.0u"nm")
coords = place_atoms(n_atoms, boundary; min_dist=0.3u"nm")

dist_cutoff = 0.9u"nm"
coul = CoulombEwald(dist_cutoff=dist_cutoff)
pme = PME(dist_cutoff, atoms, boundary)

sys = System(
    atoms=atoms,
    coords=coords,
    boundary=boundary,
    pairwise_inters=(coul,),
    general_inters=(pme,),
    loggers=(coords=CoordinatesLogger(10),),
)

temp = 1.0u"K"
simulator = VelocityVerlet(
    dt=0.001u"ps",
    coupling=AndersenThermostat(temp, 1.0u"ps"),
)

random_velocities!(sys, temp)
simulate!(sys, simulator, 2_000)

visualize(sys.loggers.coords, boundary, "sim_pme.mp4";
          color=[(i % 2 == 0 ? :red : :blue) for i in 1:n_atoms])

4 Likes