ASE

FYI, you can open this documentation as a Google Colab notebook to follow along interactively

Any GraphPESModel can be used as an ase.calculators.calculator.Calculator by wrapping it in a GraphPESCalculator. This allows the model to be used in any ASE workflow, such as geometry optimization, molecular dynamics, or phonon calculations (see below).

class graph_pes.utils.calculator.GraphPESCalculator(model, device=None, skin=1.0, **kwargs)[source]

ASE calculator wrapping any GraphPESModel.

Implements a neighbour list caching scheme (see below) controlled by the skin parameter. Using skin > 0.0 will

  • accelerate MD and minimisations

  • slow down single point calculations

If you are predomintantly doing single point calculations, use skin=0, otherwise, tune the skin paramter for your use case (see below).

Parameters:

  • model (GraphPESModel) – The model to wrap

  • device (torch.device | str | None) – The device to use for the calculation, e.g. “cpu” or “cuda”. Defaults to None, in which case the model is not moved from its current device.

  • skin (float) – The additional skin to use for neighbour list calculations. If all atoms have moved less than half of this distance between calls to calculate, the neighbour list will be reused, saving (in some cases) significant computation time.

  • **kwargs – Properties passed to the ase.calculators.calculator.Calculator base class.

calculate(
atoms=None,
properties=None,
system_changes=['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'],
)[source]

Calculate the requested properties for the given structure, and store them to self.results, as per a normal ase.calculators.calculator.Calculator.

Underneath-the-hood, this uses a neighbour list cache to speed up repeated calculations on the similar structures (i.e. particularly effective for MD and relaxations).

property cache_hit_rate: float

The ratio of calls to calculate() for which the neighbour list was reused.

reset_cache_stats()[source]

Reset the cache_hit_rate statistic.

calculate_all(
structures,
properties=None,
batch_size=5,
)[source]

Semantically identical to:

[calc.calculate(structure, properties) for structure in structures]

but with significant acceleration due to internal batching.

Parameters:
Return type:

list[dict[PropertyKey, numpy.ndarray | float]]

1. Single point calculations

First we instantiate a calculator from a simple LennardJones model, but this can be replaced with any model that inherits from GraphPESModel:

[2]:

from graph_pes.models import LennardJones

# some-what arbitrary parameters for this model
model = LennardJones(epsilon=0.14, sigma=2.27)
model

[2]:

LennardJones(epsilon=0.14, sigma=2.27, cutoff=5.0)

[3]:

calculator = model.ase_calculator()  # equivalent to GraphPESCalculator(model)
calculator

[3]:

GraphPESCalculator(
  model=LennardJones(epsilon=0.14, sigma=2.27, cutoff=5.0),
  device=cpu,
  skin=1.0
)

[4]:

from ase.build import molecule

molecules = [molecule(s) for s in "CH4 H2O C2H6".split()]
molecules

[4]:

[Atoms(symbols='CH4', pbc=False),
 Atoms(symbols='OH2', pbc=False),
 Atoms(symbols='C2H6', pbc=False)]

GraphPESCalculator s behave like normal ase.calculators.calculator.Calculator objects:

[5]:

[calculator.get_potential_energy(m) for m in molecules]

[5]:

[14828.7490234375, 30633.576171875, 21518.35546875]

For accelerated predictions via batching, feel free to also use GraphPESCalculator.calculate_all:

[6]:

calculator.calculate_all(molecules, properties=["energy"])

[6]:

[{'energy': 14828.7490234375},
 {'energy': 30633.576171875},
 {'energy': 21518.35546875}]

Merging the results of these predictions is provided by merge_predictions:

graph_pes.utils.calculator.merge_predictions(
predictions: list[dict[PropertyKey, numpy.ndarray | float]],
) → dict[Literal['local_energies', 'forces', 'energy', 'stress', 'virial'], ndarray][source]
graph_pes.utils.calculator.merge_predictions(
predictions: list[dict[PropertyKey, torch.Tensor | float]],
) → dict[Literal['local_energies', 'forces', 'energy', 'stress', 'virial'], Tensor]

Take a list of property predictions and merge them in a sensible way. Implemented for both torch.Tensor and numpy.ndarray.

Parameters:

predictions – A list of property predictions each corresponding to a single structure.

Examples

>>> predictions = [
...     {"energy": np.array(1.0), "forces": np.array([[1, 2], [3, 4]])},
...     {"energy": np.array(2.0), "forces": np.array([[5, 6], [7, 8]])},
... ]
>>> merge_predictions(predictions)
{'energy': array([1., 2.]), 'forces': array([[1, 2], [3, 4], [5, 6], [7, 8]])}

[7]:

from graph_pes.utils.calculator import merge_predictions

single_predictions = calculator.calculate_all(molecules)
merge_predictions(single_predictions)

[7]:

{'energy': array([14828.74902344, 30633.57617188, 21518.35546875]),
 'forces': array([[-0.0000000e+00, -0.0000000e+00, -0.0000000e+00],
        [ 2.3727711e+04,  2.3727711e+04,  2.3727711e+04],
        [-2.3727711e+04, -2.3727711e+04,  2.3727711e+04],
        [ 2.3727711e+04, -2.3727711e+04, -2.3727711e+04],
        [-2.3727711e+04,  2.3727711e+04, -2.3727711e+04],
        [-0.0000000e+00, -0.0000000e+00,  2.3391900e+05],
        [-0.0000000e+00,  1.5019208e+05, -1.1695950e+05],
        [-0.0000000e+00, -1.5019208e+05, -1.1695950e+05],
        [-0.0000000e+00, -1.9921422e-01, -4.2235555e+04],
        [-0.0000000e+00,  1.9921875e-01,  4.2235555e+04],
        [-0.0000000e+00,  3.6874703e+04,  1.4250874e+04],
        [-3.1934238e+04, -1.8437250e+04,  1.4250783e+04],
        [ 3.1934238e+04, -1.8437250e+04,  1.4250783e+04],
        [-0.0000000e+00, -3.6874703e+04, -1.4250874e+04],
        [-3.1934238e+04,  1.8437250e+04, -1.4250783e+04],
        [ 3.1934238e+04,  1.8437250e+04, -1.4250783e+04]], dtype=float32)}

2. Phonon calculations

See the phonons section of the ASE documentation for more details.

[8]:

from ase.build import bulk

structure = bulk('Cu', cubic=True)
structure

[8]:

Atoms(symbols='Cu4', pbc=True, cell=[3.61, 3.61, 3.61])

[9]:

import matplotlib.pyplot as plt
from ase.phonons import Phonons

%config InlineBackend.figure_format = 'svg'


phonons = Phonons(structure, calculator, supercell=(7, 7, 7), delta=0.05)
phonons.run()
phonons.read(acoustic=True)
phonons.clean()

path = structure.cell.bandpath(npoints=200)
band_structure = phonons.get_band_structure(path)

plt.figure(figsize=(5, 3))
band_structure.plot(emin=0, emax=0.03, ax=plt.gca());

WARNING, 3 imaginary frequencies at q = ( 0.00,  0.00,  0.00) ; (omega_q = 1.204e-07*i)
WARNING, 3 imaginary frequencies at q = ( 0.00,  0.00,  0.00) ; (omega_q = 1.204e-07*i)
../_images/tools_ase_16_1.svg

3. Geometry optimization

[10]:

rattled = structure.copy()
rattled.rattle(0.2)
rattled.calc = calculator
rattled.get_potential_energy()

[10]:

-3.7550203800201416

[11]:

from ase.optimize import BFGS

# optimize the rattled structure
BFGS(rattled).run(fmax=0.01)
rattled.get_potential_energy()

      Step     Time          Energy          fmax
BFGS:    0 10:31:49       -3.755020        2.332890
BFGS:    1 10:31:49       -3.913639        1.616382
BFGS:    2 10:31:49       -4.147001        0.628381
BFGS:    3 10:31:49       -4.203425        0.399706
BFGS:    4 10:31:49       -4.214467        0.317304
BFGS:    5 10:31:49       -4.221203        0.283429
BFGS:    6 10:31:49       -4.229444        0.169319
BFGS:    7 10:31:49       -4.231928        0.066617
BFGS:    8 10:31:49       -4.232348        0.044161
BFGS:    9 10:31:49       -4.232435        0.035994
BFGS:   10 10:31:49       -4.232536        0.027330
BFGS:   11 10:31:49       -4.232615        0.018950
BFGS:   12 10:31:49       -4.232644        0.011048
BFGS:   13 10:31:49       -4.232650        0.008053

[11]:

-4.232650279998779

4. Molecular dynamics

See the molecular dynamics section of the ASE documentation for more details.

[12]:

from load_atoms import view

# create a supercell to run MD on
atoms = structure.repeat((2, 2, 2))
view(atoms, show_bonds=True)

[12]:

[13]:

from ase import units
from ase.md.langevin import Langevin

# melt copper by running an NVT simulation for 1ps at 2,000K
atoms.calc = calculator
dynamics = Langevin(
    atoms,
    timestep=5 * units.fs,
    temperature_K=2_000,
    friction=0.001 / units.fs,
)
dynamics.attach(
    lambda: atoms.write("trajectory.xyz", append=True),
    interval=10,
)
dynamics.run(1_000)

view(atoms, show_bonds=True)

[13]:

[14]:

import numpy as np
from ase.io import read

trajectory = read("trajectory.xyz", index=":")
all_positions = np.stack([atoms.positions for atoms in trajectory])  # type: ignore

# plot the motion of the first few atoms (in x and y)
for i in range(4):
    plt.plot(*all_positions[:, i, :2].T)
plt.gca().set_aspect("equal")
plt.axis("off");
../_images/tools_ase_23_0.svg

Neighbour list caching

GraphPESCalculator uses a neighbour list caching scheme to speed up repeated calculations. At a high level, we check if the atoms have moved by at most skin / 2 Å from their original positions, and only need to recompute the neighbour list if not. As such, tuning the skin parameter may give you increased performance in ASE MD and minimisations for free!

Details

Generating a neighbour list is an expensive operation (especially for larger systems). We therefore attempt to minimise the:

  • number of neighbour list generations: by using a neighbourlist with a larger cutoff (model.cutoff + skin) and checking if the atoms or cell have moved more than skin / 2 Å from their original positions. This check is relatively cheap relative to generating a new neighbour list, and so helps save on flops.

  • time that each neighbour list generation takes: we do this by relying on the fantastic [vesin](https://luthaf.fr/vesin/latest/index.html) library to generate neighbour lists by passing down this task to a C++ library.

[15]:

import time


def speed_test(calculator) -> float:
    """Time how long it takes to run 500 steps of MD."""
    structure = bulk("Cu", cubic=True).repeat(5) # a big unit cell!
    structure.calc = calculator
    dynamics = Langevin(
        structure,
        timestep=5 * units.fs,
        temperature_K=1_000,
        friction=0.001 / units.fs,
    )
    tick = time.time()
    dynamics.run(500)
    tock = time.time()
    return tock - tick


skins, durations, cache_hits = [], [], []

for skin in [0, 0.05, 0.1, 0.2, 0.4, 0.8]:
    calculator = model.ase_calculator(skin=skin)
    duration = speed_test(calculator)
    print(
        f"Skin: {skin:0.2f}",
        f"Duration: {duration:0.2f}s",
        f"Total time calculating NL: {int(calculator.total_nl_timing*1000):>4}ms",
        f"Cache hits: {calculator.cache_hit_rate:.1%}",
        sep=", ",
    )
    skins.append(skin)
    durations.append(duration)
    cache_hits.append(calculator.cache_hit_rate)


plt.figure(figsize=(4, 3))
plt.plot(skins, durations, marker="o")
plt.xlabel("Skin (Å)")
plt.ylabel("Duration (s)");

Skin: 0.00, Duration: 3.86s, Total time calculating NL: 1789ms, Cache hits: 0.0%
Skin: 0.05, Duration: 3.71s, Total time calculating NL: 1476ms, Cache hits: 3.2%
Skin: 0.10, Duration: 3.56s, Total time calculating NL: 1307ms, Cache hits: 15.4%
Skin: 0.20, Duration: 2.98s, Total time calculating NL:  740ms, Cache hits: 56.5%
Skin: 0.40, Duration: 2.73s, Total time calculating NL:  396ms, Cache hits: 76.2%
Skin: 0.80, Duration: 2.62s, Total time calculating NL:  223ms, Cache hits: 89.0%
../_images/tools_ase_26_1.svg