Molecular Dynamics with Python¶

from scm.plams import *
import os
import numpy as np
from ase.visualize.plot import plot_atoms
import matplotlib.pyplot as plt

def show(mol, figsize=None, rotation=None, **kwargs):
    """ Show a molecule in a Jupyter notebook """
    plt.figure(figsize=figsize or (2,2))
    plt.axis('off')
    rotation = rotation or ('80x,10y,0z')
    plot_atoms(toASE(mol), rotation=rotation, **kwargs)

init()

single_molecule = from_smiles('C', forcefield='uff') # or use Molecule('your-own-file.xyz')
show(single_molecule)

box = packmol(
    single_molecule,
    n_molecules=8,
    region_names='methane',
    density=0.4, # g/cm^3
)
print("Lattice vectors of the box:\n{}".format(box.lattice))
show(box)

box = preoptimize(box, model='UFF', maxiterations=50)
show(box)

nvt_job = AMSNVTJob(molecule=box,
                    settings=s,
                    name='NVT-sim-1',
                    nsteps=1000,
                    samplingfreq=50,
                    timestep=1.0, # femtoseconds
                    temperature=300,
                    thermostat='Berendsen',
                    tau=100 # femtoseconds
                   )
nvt_job.run();

def plot_temperature(job, title=None):
    time = job.results.get_history_property('Time', 'MDHistory')
    temperature = job.results.get_history_property('Temperature', 'MDHistory')
    plt.figure(figsize=(5,3))
    plt.plot(time, temperature)
    plt.xlabel("Time (fs)")
    plt.ylabel("Temperature (K)")
    plt.title(title or job.name + " Temperature")
    plt.show()

plot_temperature(nvt_job)
plot_energies(nvt_job)

plot_energies(nvt_job, conserved_energy=True)

npt_job = AMSNPTJob.restart_from(nvt_job,
                                 name='npt',
                                 barostat='Berendsen',
                                 pressure=1e5, # Pa
                                 barostat_tau=1000, # fs
                                 equal='XYZ' # scale all lattice vectors equally
                                )

npt_job.run();

def moving_average(x, y, window:int):
    if not window:
        return x, y
    window = min(len(x)-1, window)
    if window <= 1:
        return x, y
    ret_x = np.convolve(x, np.ones(window)/window, mode='valid')
    ret_y = np.convolve(y, np.ones(window)/window, mode='valid')
    return ret_x, ret_y

def plot_pressure(job, title=None, window:int=None):
    time = job.results.get_history_property('Time', 'MDHistory')
    pressure = job.results.get_history_property('Pressure', 'MDHistory')
    pressure = Units.convert(pressure, 'au', 'bar') # convert from atomic units to bar
    time, pressure = moving_average(time, pressure, window)
    plt.figure(figsize=(5,3))
    plt.plot(time, pressure)
    plt.xlabel("Time (fs)")
    plt.ylabel("Pressure (bar)")
    plt.title(title or job.name + " Pressure")
    plt.show()

plot_pressure(npt_job)

plot_pressure(npt_job, window=10, title="npt Pressure moving average")

def plot_density(job, title=None, window:int=None):
    time = job.results.get_history_property('Time', 'MDHistory')
    density = job.results.get_history_property('Density', 'MDHistory')
    density = np.array(density) * Units.convert(1.0, 'au', 'kg') / Units.convert(1.0, 'bohr', 'm')**3
    time, density = moving_average(time, density, window)
    plt.figure(figsize=(5,3))
    plt.plot(time, density)
    plt.xlabel("Time (fs)")
    plt.ylabel("Density (kg/m^3)")
    plt.title(title or job.name + " Density")
    plt.show()

plot_density(npt_job)

nvt_eq_job = AMSNVTJob(
    molecule=box,
    settings=s,
    name='NVT-equilibration-1',
    temperature=300,
    thermostat='Berendsen',
    tau=200,
    timestep=1.0,
    nsteps=40000,
    samplingfreq=500,
)
nvt_eq_job.run()
plot_energies(nvt_eq_job)

nvt_prod_job = AMSNVTJob.restart_from(
    nvt_eq_job,
    name='NVT-production-1',
    nsteps=50000,
    thermostat='NHC',
    samplingfreq=100
)

nvt_prod_job.run();

plot_energies(nvt_prod_job)
plot_temperature(nvt_prod_job)

atom_indices = [i for i,at in enumerate(nvt_prod_job.results.get_main_molecule(), 1) if at.symbol == 'C']
print("Calculating MSD for carbon atoms with indices {}".format(atom_indices))
msd_job = AMSMSDJob(
    nvt_prod_job,
    name='msd-'+nvt_prod_job.name,
    atom_indices=atom_indices, # indices start with 1 for the first atom
    max_correlation_time_fs=8000,  # max correlation time must be set before running the job
    start_time_fit_fs=3000 # start_time_fit can also be changed later in the postanalysis
)
msd_job.run();

def plot_msd(job, start_time_fit_fs=None):
    """ job: an AMSMSDJob """
    time, msd = job.results.get_msd()
    fit_result, fit_x, fit_y = job.results.get_linear_fit(start_time_fit_fs=start_time_fit_fs)
    # the diffusion coefficient can also be calculated as fit_result.slope/6 (ang^2/fs)
    diffusion_coefficient = job.results.get_diffusion_coefficient(start_time_fit_fs=start_time_fit_fs) # m^2/s
    plt.figure(figsize=(5,3))
    plt.plot(time, msd, label='MSD')
    plt.plot(fit_x, fit_y, label='Linear fit slope={:.5f} ang^2/fs'.format(fit_result.slope))
    plt.legend()
    plt.xlabel("Correlation time (fs)")
    plt.ylabel("Mean square displacement (ang^2)")
    plt.title("MSD: Diffusion coefficient = {:.2e} m^2/s".format(diffusion_coefficient))
    plt.show()

plot_msd(msd_job)

plot_msd(msd_job, start_time_fit_fs=7000)

mol = nvt_prod_job.results.get_main_molecule()
atom_indices = [i for i, at in enumerate(mol, 1) if at.symbol == 'C']
atom_indices_to = [i for i, at in enumerate(mol, 1) if at.symbol == 'H']

# rmax = half the shortest box length for orthorhombic cells
rmax = min([mol.lattice[0][0], mol.lattice[1][1], mol.lattice[2][2]]) / 2

rdf_job = AMSRDFJob(
    nvt_prod_job,
    name='rdf-C-H-'+nvt_prod_job.name,
    atom_indices=atom_indices, # from C
    atom_indices_to=atom_indices_to, # to H
    rmin=0.5,
    rmax=rmax,
    rstep=0.1,
)
rdf_job.run();

def plot_rdf(job, title=""):
    plt.figure(figsize=(5,3))
    r, rdf = job.results.get_rdf()
    plt.plot(r, rdf)
    plt.xlabel("r (angstrom)")
    plt.ylabel("g(r)")
    plt.title(title or job.name + " RDF")
    plt.show()

plot_rdf(rdf_job)

print(nvt_prod_job.get_input())

print(msd_job.get_input())

nvespawner_job = AMSNVESpawnerJob(
    nvt_prod_job,
    name='nvespawner-'+nvt_prod_job.name,
    n_nve=3, # the number of NVE simulations to run
    samplingfreq=1, # write to disk every frame for velocity autocorrelation
    writevelocities=True,
    timestep=1.0, # ideally use smaller timestep
    nsteps=3000 # 3000 steps (here = 3000 fs) is very short, for demonstration purposes only
)
nvespawner_job.run();

vacf_jobs = []
for job in nvespawner_job.children.values():
    vacf_job = AMSVACFJob(
        job,
        name='vacf-'+job.name+'-'+nvespawner_job.name,
        max_correlation_time_fs=1000 # 1000 fs is very short, for demo purposes only
    )
    vacf_jobs.append(vacf_job)
    vacf_jobs[-1].run();

def plot_vacf(vacf_jobs):
    """vacf_jobs is a list of AMSVACFJob, plot the average velocity autocorrelation function"""
    x, y = zip(*[job.results.get_vacf() for job in vacf_jobs])
    x = np.mean(x, axis=0)
    y = np.mean(y, axis=0)
    plt.figure(figsize=(5,3))
    plt.plot(x, y)
    plt.xlabel("Correlation time (fs)")
    plt.ylabel("Normalized velocity autocorrelation function")
    plt.title("Average velocity autocorrelation function")

def plot_power_spectrum(vacf_jobs, newfigure=True):
    x, y = zip(*[job.results.get_power_spectrum() for job in vacf_jobs])
    x = np.mean(x, axis=0)
    y = np.mean(y, axis=0)
    if newfigure:
        plt.figure(figsize=(5,3))
    plt.plot(x, y)
    plt.xlabel("Frequency (cm^-1)")
    plt.ylabel("Intensity (arb. units)")
    plt.title("Average power spectrum")

plot_vacf(vacf_jobs)
plot_power_spectrum(vacf_jobs)

sfreq = Settings()
sfreq.input.ams.Task = 'GeometryOptimization'
sfreq.input.ams.Properties.NormalModes = 'Yes'
sfreq += s # UFF engine settings used before
freq_job = AMSJob(settings=sfreq, name='methane-optfreq', molecule=single_molecule) # single_molecule defined at the beginning
freq_job.run()
print("Normal modes stored in {}".format(freq_job.results.rkfpath(file='engine')))

freqs = freq_job.results.get_frequencies()
print("There are {} frequencies for gasphase methane".format(len(freqs)))
for f in freqs:
    print("{:.1f} cm^-1".format(f))

def gaussian(x, a, b, c):
    return a*np.exp(-(x-b)**2/(2*c**2))

def sumgaussians(x, a, b, c):
    # calculate the sum of gaussians centered at b for the different x values
    # use numpy array broadcasting by placing x in a row and b in a column
    x = np.reshape(x, (1,-1)) # row
    b = np.reshape(b, (-1,1)) # column
    return np.sum(gaussian(x, a, b, c), axis=0)

x = np.arange(5000) # x axis
y = sumgaussians(x, a=2, b=freqs, c=30)  # intensities, use suitable numbers for a and c

plot_power_spectrum(vacf_jobs)
plt.plot(x, y)
plt.legend(['power spectrum from MD', 'gasphase optimization'])
plt.show()

reaxff_s = Settings()
reaxff_s.input.ReaxFF.ForceField = 'Water2017.ff'
reaxff_s.runscript.nproc = 1
reaxff_s.runscript.preamble_lines = ['export OMP_NUM_THREADS=1']

low_density_box = packmol(from_smiles('O'), n_molecules=32, region_names='water', density=0.4)
low_density_box = preoptimize(low_density_box, maxiterations=50)
show(low_density_box)

# step 2 : initial NVT at low density

nvt_low_density_job = AMSNVTJob(
    molecule=low_density_box,
    name="nvt_low_density_job",
    settings=reaxff_s,
    nsteps=3000, # should be much larger for production purposes,
    temperature=300,
    tau=100,
    samplingfreq=500,
    timestep=0.5
)
nvt_low_density_job.run();

# step 3 : scan to higher density

scan_density_job = AMSMDScanDensityJob.restart_from(
    nvt_low_density_job,
    name="scan_density_job",
    scan_density_upper=1.4,
    nsteps=6000,
    samplingfreq=100
)
scan_density_job.run();

plot_energies(scan_density_job, kinetic_energy=False, potential_energy=False)
plot_density(scan_density_job)

history_index = scan_density_job.results.get_lowest_energy_index('TotalEnergy', 'MDHistory')
lowest_energy_molecule = scan_density_job.results.get_history_molecule(history_index)
lowest_energy_time = scan_density_job.results.get_property_at_step(history_index, 'Time', 'MDHistory')
print(f'Extracting frame {history_index}')
print(f'At MD-time: {lowest_energy_time:.1f} fs')
print(f'Density: {lowest_energy_molecule.get_density():.1f} kg/m^3')

# step 4 : short NVT equilibration at new density

nvt_equilibration_job = AMSNVTJob(
    molecule=lowest_energy_molecule,
    name="nvt_equilibration",
    settings=reaxff_s,
    nsteps=3000,
    timestep=0.5,
    temperature=300
)
nvt_equilibration_job.run();

# step 5 : long NPT equilibration

npt_equilibration = AMSNPTJob.restart_from(
    nvt_equilibration_job,
    name="npt_equilibration",
    nsteps=10000, # should be much longer!
    temperature=300,
    thermostat='NHC',
    tau=100,
    barostat='MTK',
    equal='XYZ',
    pressure=1e5,
    barostat_tau=1000
)

npt_equilibration.run();

plot_density(npt_equilibration)

equilibrated_box = npt_equilibration.results.get_equilibrated_molecule()
print("Density: {:.2f} kg/m^3".format(equilibrated_box.get_density()))
equilibrated_box.write("water_equilibrated_box.xyz")
show(equilibrated_box)

def get_lennard_jones_settings():
    s = Settings()
    s.input.lennardjones.eps = 3.785e-4
    s.input.lennardjones.rmin = 3.81637
    s.input.lennardjones.cutoff = 8.0
    s.runscript.nproc = 1    # to run in serial
    return s

sizes = [50, 150, 450] # number of atoms
inverted_L_list = [] # inverted cubic box length in angstrom
density = 1.0  # system density in g/cm^3

Ar = Molecule()
Ar.add_atom(Atom(symbol='Ar', coords=(0, 0, 0)))   # a single Ar atom

boxes = []

for size in sizes:
    mol = packmol(Ar, n_molecules=size, density=density)
    mol = preoptimize(mol, settings=get_lennard_jones_settings(), maxiterations=50)
    boxes.append(mol)
    L = mol.lattice[0][0]
    inverted_L_list.append(1.0/L)
    print(f'{size} Ar atoms, density = {density} g/cm^3')
    show(mol, figsize=(L/20, L/20))
    plt.show()

D_list = [] # diffusion coefficients
temperature = 400
max_correlation_time_fs = 6000
start_time_fit_fs = 3000

for size, mol in zip(sizes, boxes):
    eq_job = AMSNVTJob(name=f'eq_{size}', settings=get_lennard_jones_settings(), molecule=mol,
                       temperature=temperature, nsteps=10000, samplingfreq=500,
                       thermostat='Berendsen', tau=100, timestep=1.0, writebonds=False,
                       writecharges=False, writemolecules=False, writevelocities=False)
    eq_job.run()

    prod_job = AMSNVTJob.restart_from(eq_job, name=f'prod_{size}', nsteps=60000,
                                      samplingfreq=25, temperature=temperature, thermostat='NHC')
    prod_job.run()

    msd_job = AMSMSDJob(prod_job, name=f'msd_{size}', max_correlation_time_fs=max_correlation_time_fs)
    msd_job.run()

    D = msd_job.results.get_diffusion_coefficient(start_time_fit_fs=start_time_fit_fs)
    D_list.append(D)
    log(f"{size} atoms: D = {D:.2e} m^2 s^-1")

from scipy.stats import linregress
def linear_fit_extrapolate_to_0(x, y):
    """
        linear regression on x and y.
        Returns 3-tuple fit_x, fit_y, slope.

        Note: fit_x contains an extra x-value at 0 (the final value). This ensures that the point at 0 is included when plotting.
    """
    result = linregress(x, y)
    fit_x = list(x) + [0]    # append a 0
    fit_x = np.array(fit_x)
    fit_y = result.slope * fit_x + result.intercept

    return fit_x, fit_y, result.slope, result.intercept

fit_x, fit_y, slope, intercept = linear_fit_extrapolate_to_0(inverted_L_list, D_list)

plt.figure(figsize=(6,3))
plt.plot(inverted_L_list, D_list, '.', label="Simulation results")
plt.plot(fit_x, fit_y, label="Linear fit")
dilute_D = intercept  # at 0 inverted box length, what is the diffusion coefficient?
s = f"Extrapolated (dilute) D: {dilute_D:.2e} m² s⁻¹"
log(s)
plt.title(s)
plt.legend()
plt.xlabel("1/L (Å⁻¹)")
plt.ylabel("Diffusion coefficient D (m² s⁻¹)");

T_list = [50, 100, 200, 300]
mol = packmol(Ar, n_molecules=50, density=density)
mol = preoptimize(mol, settings=get_lennard_jones_settings(), maxiterations=50)
show(mol)

maxjobs = 4
config.default_jobrunner = JobRunner(parallel=True, maxjobs=maxjobs)

msd_jobs = []
DT_list = []
for T in T_list:
    eq_job = AMSNVTJob(name=f'eq_T_{T}', settings=get_lennard_jones_settings(), molecule=mol,
                       temperature=T, nsteps=10000, samplingfreq=500,
                       thermostat='Berendsen', tau=100, timestep=1.0, writebonds=False,
                       writecharges=False, writemolecules=False, writevelocities=False)
    eq_job.run()

    prod_job = AMSNVTJob.restart_from(
        eq_job,
        name=f'prod_T_{T}',
        nsteps=50000,
        samplingfreq=25,
        temperature=T,
        thermostat='NHC',
        use_prerun=True  # only access the equilibration job results when the production job is about to start
    )
    prod_job.run()

    msd_job = AMSMSDJob(prod_job, name=f'msd_T_{T}', max_correlation_time_fs=max_correlation_time_fs)
    msd_job.run()
    msd_jobs.append(msd_job)  # do not access the msd results yet, wait until all jobs have been started

# now we can access the results
for T, msd_job in zip(T_list, msd_jobs):
    D = msd_job.results.get_diffusion_coefficient(start_time_fit_fs=start_time_fit_fs)
    DT_list.append(D)
    log(f"Temperature {T} K: D = {D:.2e} m^2 s^-1")

# fit a straight line when plotting ln(d) vs. 1/T
inverted_T_list = 1.0/np.array(T_list)
ln_D_list = np.log(DT_list)
fit_x, fit_y, slope, intercept = linear_fit_extrapolate_to_0(inverted_T_list, ln_D_list)

D0 = np.exp(intercept)
boltzmann_eV_per_K = 8.617e-5 # eV/K
barrier = -boltzmann_eV_per_K * slope
title = f"D0: {D0:.2e} m² s⁻¹. Barrier: {barrier:.3f} eV"
log(title)

# multiply x-axis by 1000 for friendlier numbers
plt.figure(figsize=(6,3))
plt.plot(inverted_T_list*1000, ln_D_list, '.', label="Simulation results")
plt.plot(fit_x*1000, fit_y, label="Linear fit")
plt.title(title)
plt.legend()
plt.xlabel("1000/T (K⁻¹)")
plt.ylabel("ln D");