Geometry Optimization

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.io import read
from ase.optimize import FIRE
from weas_widget import WeasWidget

from data_tutorials.data import get_data

from janus_core.calculations.single_point import SinglePoint
from janus_core.calculations.geom_opt import GeomOpt

Use data_tutorials to get the data required for this tutorial:

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

Prepare for optimization of a deformed salt structure

NaCl = read("data/NaCl-deformed.cif")
NaCl.wrap()
v=WeasWidget()
v.from_ase(NaCl)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v

sp_mace = SinglePoint(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    properties="energy",
)

init_energy = sp_mace.run()["energy"]

To optimize only the atomic positions and not the cell, set filter_func = None:

optimized_NaCl = GeomOpt(
    struct=sp_mace.struct,
    fmax=0.001,
    filter_func=None,
)

optimized_NaCl.run()
v=WeasWidget()
v.from_ase(optimized_NaCl.struct)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v

Check energy has been lowered, and cell is unchanged:

print(f"Initial cell: {NaCl.cell.cellpar()}")
print(f"Initial energy: {init_energy}")

print(f"Final cell: {optimized_NaCl.struct.cell.cellpar()}")
print(f"Final energy: {optimized_NaCl.struct.get_potential_energy()}")

Optimizing cell vectors and atomic positions

Setting filter_kwargs = {"hydrostatic_strain": True} allows the cell lengths to be changed, in addition to atomic positions, but cell angles remain fixed:

optimized_NaCl_lengths = GeomOpt(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    fmax=0.001,
    filter_kwargs={"hydrostatic_strain": True},
)
optimized_NaCl_lengths.run()

Check energy has been lowered, and cell lengths have been updated, but angles remain unchanged:

print(f"Initial cell: {NaCl.cell.cellpar()}")
print(f"Initial energy: {init_energy}")

print(f"Final cell: {optimized_NaCl_lengths.struct.cell.cellpar()}")
print(f"Final energy: {optimized_NaCl_lengths.struct.get_potential_energy()}")

Optimizing at constant pressure and volume

Calculations can also be run at a fixed pressure and volume, by setting filter_kwargs = {"scalar_pressure": x, "constant_volume": True}

By default, both the cell lengths and angles will be optimized, in addition to the atomic positions.

We can also set the optimizer function and filter function used, either by passing the function itself (e.g. FIRE) or passing the name of the ASE function (e.g. "ExpCellFilter"):

optimized_NaCl_pressure = GeomOpt(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    fmax=0.01,
    filter_kwargs={"scalar_pressure": 0.05, "constant_volume": True},
    optimizer=FIRE,
    filter_func="ExpCellFilter",
)
optimized_NaCl_pressure.run()

Check cell lengths and angles have both been updated:

print(f"Initial cell: {NaCl.cell.cellpar()}")
print(f"Initial energy: {init_energy}")

print(f"Final cell: {optimized_NaCl_pressure.struct.cell.cellpar()}")
print(f"Final energy: {optimized_NaCl_pressure.struct.get_potential_energy()}")

Comparing MACE to CHGNet and SevenNet

optimized_NaCl_mace = GeomOpt(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    fmax=0.01,
)
optimized_NaCl_mace.run()

optimized_NaCl_chgnet = GeomOpt(
    struct=NaCl.copy(),
    arch="chgnet",
    device="cpu",
    calc_kwargs={"default_dtype": "float64"},
    fmax=0.01,
)
optimized_NaCl_chgnet.run()

optimized_NaCl_sevennet = GeomOpt(
    struct=NaCl.copy(),
    arch="sevennet",
    device="cpu",
    calc_kwargs={"default_dtype": "float64"},
    fmax=0.01,
)
optimized_NaCl_sevennet.run()

print(f"Initial energy: {init_energy}")

print(f"Final energy (MACE): {optimized_NaCl_mace.struct.get_potential_energy()}")
print(f"Final energy (CHGNet): {optimized_NaCl_chgnet.struct.get_potential_energy()}")
print(f"Final energy (SevenNet): {optimized_NaCl_sevennet.struct.get_potential_energy()}")