Molecular Dynamics

Set up environment (optional)

These steps are required to run this tutorial with Google Colab. To do so, uncomment and run the cell below.

These instructions but may work for other systems too, but it is typically preferable to prepare a virtual environment separately before running this notebook if possible.

# ! uv pip install janus-core[mace,orb,chgnet,visualise] data-tutorials --system # Install janus-core with MACE, Orb, CHGNet, and WeasWidget, and data-tutorials

To ensure you have the latest version of janus-core installed, compare the output of the following cell to the latest version available at https://pypi.org/project/janus-core/

from janus_core import __version__

print(__version__)

0.9.4

Prepare data and modules

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-core/main/docs/source/tutorials/data/",
    filename=["precomputed_NaCl-traj.xyz"],
    folder="data",
)

try to download precomputed_NaCl-traj.xyz from https://raw.githubusercontent.com/stfc/janus-core/main/docs/source/tutorials/data/ and save it in data/precomputed_NaCl-traj.xyz

saved in data/precomputed_NaCl-traj.xyz

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="small",
    calc_kwargs={"default_dtype": "float64"},
    temp_start=300.0,
    temp_end=200.0,
    temp_step=20,
    temp_time=5,
    stats_every=2,
)

/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))

cuequivariance or cuequivariance_torch is not available. Cuequivariance acceleration will be disabled.
Using Materials Project MACE for MACECalculator with /home/runner/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.

/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:226: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  torch.load(f=model_path, map_location=device)
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/ase/md/langevin.py:102: FutureWarning: The implementation of `fixcm=True` in `Langevin` does not strictly sample the correct NVT distributions. The deviations are typically small for large systems but can be more pronounced for small systems. Use `fixcm=False` together with `ase.constraints.FixCom`. `fixcm` is deprecated since ASE 3.28.0 and will be removed in a future release.
  warnings.warn(msg, FutureWarning)

Run cooling:

cooling.run()

/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/ase/io/extxyz.py:320: UserWarning: Skipping unhashable information real_time
  warnings.warn('Skipping unhashable information '

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 statistics 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)

contains 17 timeseries, each with 16 elements
index label units
0 # Step
1 Real_Time s
2 Time fs
3 Epot/N eV
4 EKin/N eV
5 T K
6 ETot/N eV
7 Density g/cm^3
8 Volume Ang^3
9 P GPa
10 Pxx GPa
11 Pyy GPa
12 Pzz GPa
13 Pyz GPa
14 Pxz GPa
15 Pxy GPa
16 Target_T K

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="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
)

Using Materials Project MACE for MACECalculator with /home/runner/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.

/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:226: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  torch.load(f=model_path, map_location=device)
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/ase/md/langevin.py:102: FutureWarning: The implementation of `fixcm=True` in `Langevin` does not strictly sample the correct NVT distributions. The deviations are typically small for large systems but can be more pronounced for small systems. Use `fixcm=False` together with `ase.constraints.FixCom`. `fixcm` is deprecated since ASE 3.28.0 and will be removed in a future release.
  warnings.warn(msg, FutureWarning)

heating.run()

/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/ase/io/extxyz.py:320: UserWarning: Skipping unhashable information real_time
  warnings.warn('Skipping unhashable information '

The same structure can then be used to run an additional NVE MD simulation 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()