Molecular Dynamics

Set up environment (optional)

These steps are required for Google Colab, but may work on other systems too:

# import locale
# locale.getpreferredencoding = lambda: "UTF-8"

# ! pip uninstall torch torchaudio torchvision numpy -y
# ! uv pip install janus-core[all] data-tutorials torch==2.5.1 --system
# get_ipython().kernel.do_shutdown(restart=True)

from ase.build import bulk
from weas_widget import WeasWidget
from ase.io import read
from data_tutorials.data import get_data
import matplotlib.pyplot as plt
import numpy as np

from janus_core.calculations.md import NVE, NVT
from janus_core.helpers.stats import Stats
from janus_core.processing import post_process

Use data_tutorials to get the data required for this tutorial:

get_data(
    url="https://raw.githubusercontent.com/stfc/janus-tutorials/main/data/",
    filename=["precomputed_NaCl-traj.xyz"],
    folder="data",
)

Cooling

Build NaCl structure and attach the MACE calculator:

NaCl = bulk("NaCl", "rocksalt", a=5.63, cubic=True)
NaCl = NaCl * (2, 2, 2)

v=WeasWidget()
v.from_ase(NaCl)
v

Prepare a simulation, cooling the structure from 300K to 200K in stepx of 20K, with 5fs at each temperature:

cooling = NVT(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    temp_start=300.0,
    temp_end=200.0,
    temp_step=20,
    temp_time=5,
    stats_every=2,
)

Run cooling:

cooling.run()

All output files are saved in a results directory, janus_results.

Within this, the final structures at each temperature are saved in Cl32Na32-nvt-T300.0-final.xyz, Cl32Na32-nvt-T280.0-final.xyz, …, Cl32Na32-nvt-T200.0-final.xyz.

The statiscs from the simulation are also saved every 20 steps in Cl32Na32-nvt-T300.0-T200.0-stats.dat. This can then be analysed using the Stats module:

data = Stats("janus_results/Cl32Na32-nvt-T300.0-T200.0-stats.dat")

print(data)

plt.plot(data[0], data[5], label="Actual")
plt.plot(data[0], data[16], label="Target")
plt.legend()
plt.xlabel("Steps")
plt.ylabel("Temperature / K")
plt.show()

Heating, followed by MD

This will prepare an NVT MD simulation, initially increasing the temperature from 0K to 300K in 20K steps, with 1fs at each temperature, before a further 10 steps (10fs) at 300K.

The final structure at each temperature will be saved, e.g. Cl32Na32-nvt-T0-final.xyz, Cl32Na32-nvt-T0-final.xyz, …, Cl32Na32-nvt-T300-final.xyz.

heating = NVT(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    temp_start=0,  # Start of temperature ramp
    temp_end=300.0,  # End of temperature ramp
    temp_step=20,  # Temperature ramp increments
    temp_time=1,  # Time at each temperature in ramp
    temp=300,  # MD temperature
    steps=10,  # MD steps at 300K
)

heating.run()

The same structure can then be used to run an additional NVE MD similation for 50 steps (50fs), with post-processing performed to compute the RDF by setting post_process_kwargs = {"rdf_compute": True}, with the results saved to Cl32Na32-nve-rdf.dat:

md = NVE(
    struct=heating.struct,
    temp=300,
    stats_every=5,
    steps=50,
    post_process_kwargs={"rdf_compute": True, "rdf_rmax": 5, "rdf_bins": 50},
)

md.run()

# view trajectory

v=WeasWidget()
traj = read("data/precomputed_NaCl-traj.xyz", index=":")
v.from_ase(traj)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v

rdf = np.loadtxt("janus_results/Cl32Na32-nve-T300-rdf.dat")
bins, counts = zip(*rdf)

plt.plot(bins, counts)
plt.ylabel("RDF")
plt.xlabel("Distance / Å")
plt.show()

The same post-processing can also be performed separately after completing the simulation:

data = read("data/precomputed_NaCl-traj.xyz", index=":")

rdf = post_process.compute_rdf(traj, rmax=5.0, nbins=50)

plt.plot(rdf[0], rdf[1])
plt.ylabel("RDF")
plt.xlabel("Distance / Å")
plt.show()