MACE Tutorial¶

Introduction to Atomic Simulation Environment (ASE)¶

The Atomic Simulation Environment (ASE) is a set of tools and Python modules for setting up, manipulating, running, visualizing and analyzing atomistic simulations. The code is freely available under the GNU LGPL license.

Installation¶

Here we create a virtual environment named "mace".

mkdir mace
python -m venv mace

Then we activate this environment

source mace/bin/activate

Then install PyTorch from their website. Choose the correct installation command based on your computer hardware and OS.

pip install torch torchvision torchaudio

Then you can install a specific branch of MACE

pip install git+https://github.com/ACEsuit/mace.git@multi-head-interface

You might need some other packages such as pymatgen and jupyter.

pip install ase jupyter pymatgen

Run Jupyter on VS Code¶

We prefer run jupyter through VS Code. Alternatively, you can also start jupyter by typing jupyter notebook.

Install Python and Jupyter extension on VS Code Extension Store.
Then start a new Jupyer notebook by Ctrl/⌘ + shift + P
Select Create: New Jupyter Notebook
Choose the correct kernel from the top right side

Calculate total energy¶

You can change device for different architecture:

cpu: CPU only
cuda: NVIDIA GPU (CUDA)
mps: Apple Metal (you need to install specific version of Pytorch on machine with Apple Silicon, see https://developer.apple.com/metal/pytorch/)

In [1]:

Copied!





from mace.calculators import mace_mp
from ase import build

atoms = build.molecule('H2O')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
atoms.calc = calc
print(atoms.get_potential_energy())
from mace.calculators import mace_mp
from ase import build

atoms = build.molecule('H2O')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
atoms.calc = calc
print(atoms.get_potential_energy())

Using Materials Project MACE for MACECalculator with /home/jerry/.cache/mace/5yyxdm76
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
-14.15936617752067

Visualize structure using ASE¶

Get the cif file of of Al(HCOO)₃ here: https://files.matsci.dev/Al_empty.cif
We can visualize the structure of Al(HCOO)₃ using ASE

In [2]:

Copied!

from ase.visualize import view
from ase.io import read

atoms = read('Al_empty.cif')
view(atoms, viewer='x3d')
from ase.visualize import view
from ase.io import read

atoms = read('Al_empty.cif')
view(atoms, viewer='x3d')

Out[2]:

ASE atomic visualization

Geometry optimization¶

What is geometry optimization?¶

In computational chemistry, energy minimization (also known as energy or geometry optimization) refers to the process of determining the spatial arrangement of a set of atoms where, based on a computational model of chemical bonding, the net force acting on each atom is nearly zero.

Why geomtery optimization?¶

Optimized structures often correspond to a structure can be used in a variety of experimental and theoretical investigations in the fields of chemical structure, thermodynamics, chemical kinetics, spectroscopy and others.
Self-consistency: our data is fitted through DFT which comes from first-principles simulation. So we want everything to be consistent.
Experiment also has errors, so we cannot always rely on experimental data, although experimental data is usually a good starting point.

LBFGS: limited memory BFGS, a multi variable optimization algorithm.
By default, only ion positions are optimized. https://wiki.fysik.dtu.dk/ase/ase/optimize.html Symmetry is also not kept by default.

In [3]:

Copied!





from mace.calculators import mace_mp
from ase.optimize import LBFGS
from ase.io import read

initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
initial_structure.calc = calc
print('LATTICE before optimization:', initial_structure.get_cell())
optimization = LBFGS(initial_structure)
optimization.run(fmax=0.01)
cell = optimization.atoms.get_cell()
print('LATTICE after optimization:', cell)
from mace.calculators import mace_mp
from ase.optimize import LBFGS
from ase.io import read

initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
initial_structure.calc = calc
print('LATTICE before optimization:', initial_structure.get_cell())
optimization = LBFGS(initial_structure)
optimization.run(fmax=0.01)
cell = optimization.atoms.get_cell()
print('LATTICE after optimization:', cell)

Using Materials Project MACE for MACECalculator with /home/jerry/.cache/mace/5yyxdm76
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
LATTICE before optimization: Cell([11.4477, 11.4477, 11.4477])
       Step     Time          Energy          fmax
LBFGS:    0 23:36:38     -687.985801        8.985016
LBFGS:    1 23:36:39     -698.894783        5.045000
LBFGS:    2 23:36:40     -703.518858        1.314353
LBFGS:    3 23:36:40     -704.565872        0.566692
LBFGS:    4 23:36:41     -704.741253        0.205089
LBFGS:    5 23:36:41     -704.772760        0.102207
LBFGS:    6 23:36:41     -704.783226        0.062383
LBFGS:    7 23:36:41     -704.786430        0.080856
LBFGS:    8 23:36:41     -704.797567        0.139131
LBFGS:    9 23:36:41     -704.817199        0.202122
LBFGS:   10 23:36:41     -704.852366        0.247443
LBFGS:   11 23:36:41     -704.892357        0.209105
LBFGS:   12 23:36:41     -704.923101        0.079777
LBFGS:   13 23:36:42     -704.930583        0.039786
LBFGS:   14 23:36:42     -704.933032        0.049596
LBFGS:   15 23:36:42     -704.941557        0.100894
LBFGS:   16 23:36:42     -704.954391        0.137067
LBFGS:   17 23:36:42     -704.969628        0.124299
LBFGS:   18 23:36:42     -704.977193        0.069619
LBFGS:   19 23:36:42     -704.980221        0.030677
LBFGS:   20 23:36:42     -704.981990        0.036245
LBFGS:   21 23:36:42     -704.985073        0.051129
LBFGS:   22 23:36:42     -704.992450        0.072587
LBFGS:   23 23:36:42     -705.005228        0.103108
LBFGS:   24 23:36:42     -705.020642        0.105787
LBFGS:   25 23:36:42     -705.034254        0.061800
LBFGS:   26 23:36:43     -705.040569        0.022868
LBFGS:   27 23:36:43     -705.041223        0.006139
LATTICE after optimization: Cell([11.4477, 11.4477, 11.4477])

Geometry optimization with variable volume and cell shape¶

You need to use filter to relax cell shape, see here: https://wiki.fysik.dtu.dk/ase/ase/filters.html

In [4]:

Copied!





from mace.calculators import mace_mp
from ase.optimize import LBFGS
from ase.filters import UnitCellFilter
from ase.io import Trajectory

initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
initial_structure.calc = calc
print('LATTICE before optimization:', initial_structure.get_cell())
filter = UnitCellFilter(atoms= initial_structure )
traj = Trajectory('nasicon_rlx.traj', 'w', initial_structure)
optimization = LBFGS(filter)
optimization.attach(traj)
optimization.run(fmax=0.01)

cell = optimization.atoms.get_cell()
print('LATTICE after optimization:', cell)
from mace.calculators import mace_mp
from ase.optimize import LBFGS
from ase.filters import UnitCellFilter
from ase.io import Trajectory

initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
initial_structure.calc = calc
print('LATTICE before optimization:', initial_structure.get_cell())
filter = UnitCellFilter(atoms= initial_structure )
traj = Trajectory('nasicon_rlx.traj', 'w', initial_structure)
optimization = LBFGS(filter)
optimization.attach(traj)
optimization.run(fmax=0.01)

cell = optimization.atoms.get_cell()
print('LATTICE after optimization:', cell)

Using Materials Project MACE for MACECalculator with /home/jerry/.cache/mace/5yyxdm76
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
LATTICE before optimization: Cell([11.4477, 11.4477, 11.4477])
       Step     Time          Energy          fmax
LBFGS:    0 23:36:43     -687.985801        8.985016
LBFGS:    1 23:36:43     -698.914440        5.038004
LBFGS:    2 23:36:43     -703.522645        1.314023
LBFGS:    3 23:36:43     -704.567930        0.565156
LBFGS:    4 23:36:43     -704.741735        0.204268
LBFGS:    5 23:36:43     -704.772839        0.102224
LBFGS:    6 23:36:43     -704.783130        0.061987
LBFGS:    7 23:36:43     -704.786238        0.079769
LBFGS:    8 23:36:43     -704.797120        0.137034
LBFGS:    9 23:36:43     -704.816307        0.199312
LBFGS:   10 23:36:44     -704.850851        0.245009
LBFGS:   11 23:36:44     -704.890516        0.208752
LBFGS:   12 23:36:44     -704.921875        0.080856
LBFGS:   13 23:36:44     -704.929964        0.058823
LBFGS:   14 23:36:44     -704.932704        0.054160
LBFGS:   15 23:36:44     -704.942119        0.106394
LBFGS:   16 23:36:44     -704.957030        0.147527
LBFGS:   17 23:36:44     -704.975576        0.137814
LBFGS:   18 23:36:44     -704.985317        0.080052
LBFGS:   19 23:36:44     -704.989747        0.054724
LBFGS:   20 23:36:44     -704.991979        0.057910
LBFGS:   21 23:36:44     -704.995254        0.059925
LBFGS:   22 23:36:44     -705.003731        0.073990
LBFGS:   23 23:36:45     -705.017741        0.105520
LBFGS:   24 23:36:45     -705.033517        0.103482
LBFGS:   25 23:36:45     -705.045536        0.055653
LBFGS:   26 23:36:45     -705.050274        0.018427
LBFGS:   27 23:36:45     -705.050665        0.009880
LATTICE after optimization: Cell([[11.427339442399717, 2.460701429600511e-16, 2.9154499928436267e-15], [2.460701429600521e-16, 11.427339442399704, 1.5007287177603883e-15], [2.9154499928435967e-15, 1.5007287177603883e-15, 11.427339442399704]])

Geometry optimization with symmetry fixed¶

https://docs.matlantis.com/atomistic-simulation-tutorial/en/2_2_opt_symmetry.html

In [12]:

Copied!





from ase import Atoms
from ase.calculators.calculator import Calculator
from ase.filters import FrechetCellFilter
from ase.optimize import LBFGS
from ase.constraints import FixSymmetry
from mace.calculators import mace_mp
from ase.io import read

def opt_with_symmetry(
    atoms_in: Atoms,
    calculator: Calculator,
    fix_symmetry: bool = False,
    hydrostatic_strain: bool = False,
) -> Atoms:
    atoms = atoms_in.copy()
    atoms.calc = calculator
    if fix_symmetry:
        atoms.set_constraint([FixSymmetry(atoms)])
    ecf = FrechetCellFilter(atoms, hydrostatic_strain=hydrostatic_strain)
    opt = LBFGS(ecf)
    opt.run(fmax=0.01)

    cell_diff = (atoms.cell.cellpar() / atoms_in.cell.cellpar() - 1.0) * 100
    print("Optimized Cell         :", atoms.cell.cellpar())
    print("Optimized Cell diff (%):", cell_diff)
    # print("Scaled positions       :\n", atoms.get_scaled_positions())
    print(f"Epot after opt: {atoms.get_potential_energy()} eV")
    return atoms

atoms = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
atoms2 = opt_with_symmetry(atoms, calc, fix_symmetry=True, hydrostatic_strain=False)
from ase import Atoms
from ase.calculators.calculator import Calculator
from ase.filters import FrechetCellFilter
from ase.optimize import LBFGS
from ase.constraints import FixSymmetry
from mace.calculators import mace_mp
from ase.io import read

def opt_with_symmetry(
    atoms_in: Atoms,
    calculator: Calculator,
    fix_symmetry: bool = False,
    hydrostatic_strain: bool = False,
) -> Atoms:
    atoms = atoms_in.copy()
    atoms.calc = calculator
    if fix_symmetry:
        atoms.set_constraint([FixSymmetry(atoms)])
    ecf = FrechetCellFilter(atoms, hydrostatic_strain=hydrostatic_strain)
    opt = LBFGS(ecf)
    opt.run(fmax=0.01)

    cell_diff = (atoms.cell.cellpar() / atoms_in.cell.cellpar() - 1.0) * 100
    print("Optimized Cell         :", atoms.cell.cellpar())
    print("Optimized Cell diff (%):", cell_diff)
    # print("Scaled positions       :\n", atoms.get_scaled_positions())
    print(f"Epot after opt: {atoms.get_potential_energy()} eV")
    return atoms

atoms = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
atoms2 = opt_with_symmetry(atoms, calc, fix_symmetry=True, hydrostatic_strain=False)

Using Materials Project MACE for MACECalculator with /home/jerry/.cache/mace/5yyxdm76
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
       Step     Time          Energy          fmax
LBFGS:    0 23:42:56     -687.985801        8.985016
LBFGS:    1 23:42:56     -698.914442        5.038003
LBFGS:    2 23:42:57     -703.522645        1.314024
LBFGS:    3 23:42:57     -704.567930        0.565156
LBFGS:    4 23:42:57     -704.741735        0.204268
LBFGS:    5 23:42:57     -704.772839        0.102224
LBFGS:    6 23:42:57     -704.783130        0.061987
LBFGS:    7 23:42:57     -704.786238        0.079769
LBFGS:    8 23:42:57     -704.797120        0.137034
LBFGS:    9 23:42:57     -704.816307        0.199312
LBFGS:   10 23:42:57     -704.850851        0.245009
LBFGS:   11 23:42:57     -704.890516        0.208752
LBFGS:   12 23:42:58     -704.921875        0.080856
LBFGS:   13 23:42:58     -704.929964        0.058833
LBFGS:   14 23:42:58     -704.932704        0.054168
LBFGS:   15 23:42:58     -704.942120        0.106395
LBFGS:   16 23:42:58     -704.957031        0.147528
LBFGS:   17 23:42:58     -704.975577        0.137812
LBFGS:   18 23:42:58     -704.985318        0.080050
LBFGS:   19 23:42:58     -704.989747        0.054702
LBFGS:   20 23:42:58     -704.991979        0.057885
LBFGS:   21 23:42:58     -704.995254        0.059898
LBFGS:   22 23:42:59     -705.003729        0.073978
LBFGS:   23 23:42:59     -705.017736        0.105502
LBFGS:   24 23:42:59     -705.033510        0.103467
LBFGS:   25 23:42:59     -705.045529        0.055652
LBFGS:   26 23:42:59     -705.050268        0.018427
LBFGS:   27 23:42:59     -705.050660        0.009874
Optimized Cell         : [11.42735741 11.42735741 11.42735741 90.         90.         90.        ]
Optimized Cell diff (%): [-0.17770025 -0.17770025 -0.17770025  0.          0.          0.        ]
Epot after opt: -705.050659741988 eV

Benchmark CPU vs GPU¶

On this machine (Intel i9-13900KF + Nvidia RTX4090), GPU is 10 times faster than CPU.

In [3]:

Copied!





from mace.calculators import mace_mp
from ase.optimize import LBFGS
import time
from ase.io import read

# CPU
initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cpu')
initial_structure.calc = calc
optimization = LBFGS(atoms = initial_structure)
start_time_stamp  = time.time()
optimization.run(steps=30,fmax=1e-9)
end_time_stamp  = time.time()
cell = optimization.atoms.get_cell()


print('CPU time: ', end_time_stamp - start_time_stamp)

# GPU without cueq
initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=False)
initial_structure.calc = calc
optimization = LBFGS(atoms = initial_structure)
start_time_stamp  = time.time()
optimization.run(steps=30,fmax=1e-9)
end_time_stamp  = time.time()

cell = optimization.atoms.get_cell()

print('GPU time without cueq: ', end_time_stamp - start_time_stamp)


# GPU
initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
initial_structure.calc = calc
optimization = LBFGS(atoms = initial_structure)
start_time_stamp  = time.time()
optimization.run(steps=30,fmax=1e-9)
end_time_stamp  = time.time()

cell = optimization.atoms.get_cell()

print('GPU time with cueq: ', end_time_stamp - start_time_stamp)
from mace.calculators import mace_mp
from ase.optimize import LBFGS
import time
from ase.io import read

# CPU
initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cpu')
initial_structure.calc = calc
optimization = LBFGS(atoms = initial_structure)
start_time_stamp  = time.time()
optimization.run(steps=30,fmax=1e-9)
end_time_stamp  = time.time()
cell = optimization.atoms.get_cell()


print('CPU time: ', end_time_stamp - start_time_stamp)

# GPU without cueq
initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=False)
initial_structure.calc = calc
optimization = LBFGS(atoms = initial_structure)
start_time_stamp  = time.time()
optimization.run(steps=30,fmax=1e-9)
end_time_stamp  = time.time()

cell = optimization.atoms.get_cell()

print('GPU time without cueq: ', end_time_stamp - start_time_stamp)


# GPU
initial_structure = read('Al_empty.cif')
calc = mace_mp(model="medium", dispersion=False, default_dtype="float64", device='cuda', enable_cueq=True)
initial_structure.calc = calc
optimization = LBFGS(atoms = initial_structure)
start_time_stamp  = time.time()
optimization.run(steps=30,fmax=1e-9)
end_time_stamp  = time.time()

cell = optimization.atoms.get_cell()

print('GPU time with cueq: ', end_time_stamp - start_time_stamp)

Using Materials Project MACE for MACECalculator with /home/jerry/.cache/mace/20231203mace128L1_epoch199model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.

/home/jerry/app/mace/lib64/python3.12/site-packages/mace/calculators/mace.py:134: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  torch.load(f=model_path, map_location=device)

       Step     Time          Energy          fmax
LBFGS:    0 11:47:52     -687.985801        8.985016
LBFGS:    1 11:47:53     -698.894783        5.045000
LBFGS:    2 11:47:54     -703.518858        1.314353
LBFGS:    3 11:47:55     -704.565872        0.566692
LBFGS:    4 11:47:56     -704.741253        0.205089
LBFGS:    5 11:47:57     -704.772760        0.102207
LBFGS:    6 11:47:58     -704.783226        0.062383
LBFGS:    7 11:47:59     -704.786430        0.080856
LBFGS:    8 11:48:00     -704.797567        0.139131
LBFGS:    9 11:48:01     -704.817199        0.202122
LBFGS:   10 11:48:02     -704.852366        0.247443
LBFGS:   11 11:48:03     -704.892357        0.209105
LBFGS:   12 11:48:04     -704.923101        0.079777
LBFGS:   13 11:48:05     -704.930583        0.039786
LBFGS:   14 11:48:06     -704.933032        0.049596
LBFGS:   15 11:48:06     -704.941557        0.100894
LBFGS:   16 11:48:07     -704.954391        0.137067
LBFGS:   17 11:48:08     -704.969628        0.124299
LBFGS:   18 11:48:09     -704.977193        0.069619
LBFGS:   19 11:48:10     -704.980221        0.030677
LBFGS:   20 11:48:11     -704.981990        0.036245
LBFGS:   21 11:48:12     -704.985073        0.051129
LBFGS:   22 11:48:13     -704.992450        0.072587
LBFGS:   23 11:48:14     -705.005228        0.103108
LBFGS:   24 11:48:15     -705.020642        0.105787
LBFGS:   25 11:48:16     -705.034254        0.061800
LBFGS:   26 11:48:17     -705.040569        0.022868
LBFGS:   27 11:48:18     -705.041223        0.006139
LBFGS:   28 11:48:19     -705.041245        0.002933
LBFGS:   29 11:48:20     -705.041252        0.001097
LBFGS:   30 11:48:21     -705.041255        0.000241
CPU time:  29.81839919090271
Using Materials Project MACE for MACECalculator with /home/jerry/.cache/mace/20231203mace128L1_epoch199model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.

/home/jerry/app/mace/lib64/python3.12/site-packages/mace/calculators/mace.py:134: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  torch.load(f=model_path, map_location=device)

       Step     Time          Energy          fmax
LBFGS:    0 11:48:21     -687.985801        8.985016
LBFGS:    1 11:48:21     -698.894783        5.045000
LBFGS:    2 11:48:21     -703.518858        1.314353
LBFGS:    3 11:48:21     -704.565872        0.566692
LBFGS:    4 11:48:21     -704.741253        0.205089
LBFGS:    5 11:48:22     -704.772760        0.102207
LBFGS:    6 11:48:22     -704.783226        0.062383
LBFGS:    7 11:48:22     -704.786430        0.080856
LBFGS:    8 11:48:22     -704.797567        0.139131
LBFGS:    9 11:48:22     -704.817199        0.202122
LBFGS:   10 11:48:22     -704.852366        0.247443
LBFGS:   11 11:48:22     -704.892357        0.209105
LBFGS:   12 11:48:23     -704.923101        0.079777
LBFGS:   13 11:48:23     -704.930583        0.039786
LBFGS:   14 11:48:23     -704.933032        0.049596
LBFGS:   15 11:48:23     -704.941557        0.100894
LBFGS:   16 11:48:23     -704.954391        0.137067
LBFGS:   17 11:48:23     -704.969628        0.124299
LBFGS:   18 11:48:23     -704.977193        0.069619
LBFGS:   19 11:48:23     -704.980221        0.030677
LBFGS:   20 11:48:23     -704.981990        0.036245
LBFGS:   21 11:48:23     -704.985073        0.051129
LBFGS:   22 11:48:23     -704.992450        0.072587
LBFGS:   23 11:48:23     -705.005228        0.103108
LBFGS:   24 11:48:23     -705.020642        0.105787
LBFGS:   25 11:48:24     -705.034254        0.061800
LBFGS:   26 11:48:24     -705.040569        0.022868
LBFGS:   27 11:48:24     -705.041223        0.006139
LBFGS:   28 11:48:24     -705.041245        0.002933
LBFGS:   29 11:48:24     -705.041252        0.001097
LBFGS:   30 11:48:24     -705.041255        0.000241
GPU time without cueq:  3.1842451095581055
Using Materials Project MACE for MACECalculator with /home/jerry/.cache/mace/20231203mace128L1_epoch199model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
Converting models to CuEq for acceleration

/home/jerry/app/mace/lib64/python3.12/site-packages/mace/calculators/mace.py:134: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  torch.load(f=model_path, map_location=device)
/home/jerry/app/mace/lib64/python3.12/site-packages/mace/modules/models.py:69: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  "atomic_numbers", torch.tensor(atomic_numbers, dtype=torch.int64)

       Step     Time          Energy          fmax
LBFGS:    0 11:48:28     -687.985801        8.985016
LBFGS:    1 11:48:28     -698.894783        5.045000
LBFGS:    2 11:48:28     -703.518858        1.314353
LBFGS:    3 11:48:28     -704.565872        0.566692
LBFGS:    4 11:48:28     -704.741253        0.205089
LBFGS:    5 11:48:28     -704.772760        0.102207
LBFGS:    6 11:48:28     -704.783226        0.062383
LBFGS:    7 11:48:28     -704.786430        0.080856
LBFGS:    8 11:48:28     -704.797567        0.139131
LBFGS:    9 11:48:28     -704.817199        0.202122
LBFGS:   10 11:48:28     -704.852366        0.247443
LBFGS:   11 11:48:28     -704.892357        0.209105
LBFGS:   12 11:48:29     -704.923101        0.079777
LBFGS:   13 11:48:29     -704.930583        0.039786
LBFGS:   14 11:48:29     -704.933032        0.049596
LBFGS:   15 11:48:29     -704.941557        0.100894
LBFGS:   16 11:48:29     -704.954391        0.137067
LBFGS:   17 11:48:29     -704.969628        0.124299
LBFGS:   18 11:48:29     -704.977193        0.069619
LBFGS:   19 11:48:29     -704.980221        0.030677
LBFGS:   20 11:48:29     -704.981990        0.036245
LBFGS:   21 11:48:29     -704.985073        0.051129
LBFGS:   22 11:48:29     -704.992450        0.072587
LBFGS:   23 11:48:29     -705.005228        0.103108
LBFGS:   24 11:48:29     -705.020642        0.105787
LBFGS:   25 11:48:29     -705.034254        0.061800
LBFGS:   26 11:48:29     -705.040569        0.022868
LBFGS:   27 11:48:29     -705.041223        0.006139
LBFGS:   28 11:48:29     -705.041245        0.002933
LBFGS:   29 11:48:29     -705.041252        0.001097
LBFGS:   30 11:48:29     -705.041255        0.000241
GPU time with cueq:  0.7055063247680664

Molecular Dynamics¶

Check more information in Molecular Dynamics
More details about functions in ASE can be found here: https://wiki.fysik.dtu.dk/ase/ase/md.html

In [ ]:

Copied!





from ase import units
from ase.md import MDLogger
from time import perf_counter
from mace.calculators import MACECalculator
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution,Stationary
from ase.md.nvtberendsen import NVTBerendsen
from ase import Atoms
from mace.calculators import mace_mp
import os

def run_md(atoms:Atoms,
           calc:MACECalculator,
           i:int=0,
           output_dir:str='./md/',
           time_step:float = 1.0,
           temperature:int=800,
           num_md_steps:int = 15000,
           num_interval:int = 100,
           taut:float=1.0)->None:
    """
    Run MD simulation for the given atoms object with the given calculator.
    Parameters:
        atoms (Atoms): The atoms object to perform MD simulation on.
        calc (MACECalculator): The MACE calculator to use for the MD simulation.
        i (int, optional): The index of the MD simulation. Default is 0.
        output_dir (str, optional): The directory to save the output files. Default is './md/'.
        time_step (float, optional): The time step for the MD simulation. Default is 1.0.
        temperature (int, optional): The temperature for the MD simulation. Default is 800.
        num_md_steps (int, optional): The number of MD steps to run. Default is 15000.
        num_interval (int, optional): The interval at which to print the MD simulation information. Default is 100.
        taut (float, optional): The relaxation time for the thermostat. Default is 1.0.
    Returns:
        None
    """
    # Print statements
    def print_dyn()->None:
        imd = dyn.get_number_of_steps()
        etot  = atoms.get_total_energy()
        temp_K = atoms.get_temperature()
        stress = atoms.get_stress(include_ideal_gas=True)/units.GPa
        stress_ave = (stress[0]+stress[1]+stress[2])/3.0
        elapsed_time = perf_counter() - start_time
        print(f"  {imd: >3}   {etot:.3f}    {temp_K:.2f}    {stress_ave:.2f}  {stress[0]:.2f}  {stress[1]:.2f}  {stress[2]:.2f}  {stress[3]:.2f}  {stress[4]:.2f}  {stress[5]:.2f}    {elapsed_time:.3f}")

    output_filename = "{}_NVT_{}K_{}".format(str(atoms.get_chemical_formula()),temperature,i)
    log_filename = output_dir + output_filename + ".log"
    traj_filename = output_dir + output_filename + ".traj"

    atoms.calc = calc
    try:
        os.mkdir('md')
    except IOError:
        print(f'folder md exist!')
    # Set the momenta corresponding to the given "temperature"
    MaxwellBoltzmannDistribution(atoms, temperature_K=temperature,force_temp=True)
    Stationary(atoms)  # Set zero total momentum to avoid drifting

    dyn = NVTBerendsen(atoms, time_step*units.fs, temperature_K = temperature, taut=taut*units.fs, loginterval=num_interval, trajectory=traj_filename)
    # dyn = NPTBerendsen(atoms = atoms, 
    #                   timestep = time_step*units.fs, 
    #                   temperature_K = temperature, 
    #                   taut=taut*units.fs,
    #                   pressure_au = 1.01325 * units.bar,
    #                   taup=1000 * units.fs, 
    #                   compressibility_au=4.57e-5 / units.bar,
    #                   loginterval=num_interval, 
    #                   trajectory=traj_filename)
    dyn.attach(print_dyn, interval=num_interval)
    dyn.attach(MDLogger(dyn, atoms, log_filename, header=True, stress=True, peratom=True, mode="w"), interval=num_interval)

    # run MD
    start_time = perf_counter()
    print(f"    imd     Etot(eV)    T(K)    stress(mean,xx,yy,zz,yz,xz,xy)(GPa)  elapsed_time(sec)")
    dyn.run(num_md_steps)  # take 10000 steps


calc = mace_mp(model="medium", dispersion=False, default_dtype="float32", device='cuda',enable_cueq=True)

atoms = read('Al_empty.cif')

run_md(atoms=atoms,calc=calc,num_md_steps=200,num_interval=1,temperature=300,output_dir='./md/')
from ase import units
from ase.md import MDLogger
from time import perf_counter
from mace.calculators import MACECalculator
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution,Stationary
from ase.md.nvtberendsen import NVTBerendsen
from ase import Atoms
from mace.calculators import mace_mp
import os

def run_md(atoms:Atoms,
           calc:MACECalculator,
           i:int=0,
           output_dir:str='./md/',
           time_step:float = 1.0,
           temperature:int=800,
           num_md_steps:int = 15000,
           num_interval:int = 100,
           taut:float=1.0)->None:
    """
    Run MD simulation for the given atoms object with the given calculator.
    Parameters:
        atoms (Atoms): The atoms object to perform MD simulation on.
        calc (MACECalculator): The MACE calculator to use for the MD simulation.
        i (int, optional): The index of the MD simulation. Default is 0.
        output_dir (str, optional): The directory to save the output files. Default is './md/'.
        time_step (float, optional): The time step for the MD simulation. Default is 1.0.
        temperature (int, optional): The temperature for the MD simulation. Default is 800.
        num_md_steps (int, optional): The number of MD steps to run. Default is 15000.
        num_interval (int, optional): The interval at which to print the MD simulation information. Default is 100.
        taut (float, optional): The relaxation time for the thermostat. Default is 1.0.
    Returns:
        None
    """
    # Print statements
    def print_dyn()->None:
        imd = dyn.get_number_of_steps()
        etot  = atoms.get_total_energy()
        temp_K = atoms.get_temperature()
        stress = atoms.get_stress(include_ideal_gas=True)/units.GPa
        stress_ave = (stress[0]+stress[1]+stress[2])/3.0
        elapsed_time = perf_counter() - start_time
        print(f"  {imd: >3}   {etot:.3f}    {temp_K:.2f}    {stress_ave:.2f}  {stress[0]:.2f}  {stress[1]:.2f}  {stress[2]:.2f}  {stress[3]:.2f}  {stress[4]:.2f}  {stress[5]:.2f}    {elapsed_time:.3f}")

    output_filename = "{}_NVT_{}K_{}".format(str(atoms.get_chemical_formula()),temperature,i)
    log_filename = output_dir + output_filename + ".log"
    traj_filename = output_dir + output_filename + ".traj"

    atoms.calc = calc
    try:
        os.mkdir('md')
    except IOError:
        print(f'folder md exist!')
    # Set the momenta corresponding to the given "temperature"
    MaxwellBoltzmannDistribution(atoms, temperature_K=temperature,force_temp=True)
    Stationary(atoms)  # Set zero total momentum to avoid drifting

    dyn = NVTBerendsen(atoms, time_step*units.fs, temperature_K = temperature, taut=taut*units.fs, loginterval=num_interval, trajectory=traj_filename)
    # dyn = NPTBerendsen(atoms = atoms, 
    #                   timestep = time_step*units.fs, 
    #                   temperature_K = temperature, 
    #                   taut=taut*units.fs,
    #                   pressure_au = 1.01325 * units.bar,
    #                   taup=1000 * units.fs, 
    #                   compressibility_au=4.57e-5 / units.bar,
    #                   loginterval=num_interval, 
    #                   trajectory=traj_filename)
    dyn.attach(print_dyn, interval=num_interval)
    dyn.attach(MDLogger(dyn, atoms, log_filename, header=True, stress=True, peratom=True, mode="w"), interval=num_interval)

    # run MD
    start_time = perf_counter()
    print(f"    imd     Etot(eV)    T(K)    stress(mean,xx,yy,zz,yz,xz,xy)(GPa)  elapsed_time(sec)")
    dyn.run(num_md_steps)  # take 10000 steps


calc = mace_mp(model="medium", dispersion=False, default_dtype="float32", device='cuda',enable_cueq=True)

atoms = read('Al_empty.cif')

run_md(atoms=atoms,calc=calc,num_md_steps=200,num_interval=1,temperature=300,output_dir='./md/')