Phonons

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)

Phonons (periodic)

from ase.build import bulk
from weas_widget import WeasWidget

from janus_core.calculations.phonons import Phonons

Prepare for phonon calculations on salt

NaCl = bulk("NaCl", "rocksalt", a=5.63, cubic=True)
v=WeasWidget()
v.from_ase(NaCl)
v.avr.model_style = 1
v.avr.show_hydrogen_bonds = True
v

Note: Set filter_func = None for geometry optimization via minimize_kwargs, so cell is fixed

phonons_mace = Phonons(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    supercell=[2, 2, 2],
    displacement=0.01,
    temp_step=10.0,
    temp_min=0.0,
    temp_max=1000.0,
    minimize=False,
    force_consts_to_hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_func": None},
    write_results=True,
)

Optimize structure and calculate force constants using phonopy.

This will save phonopy to Cl4Na4-phonopy.yml, and additionally save force constants to Cl4Na4-force_constants.hdf5:

phonons_mace.calc_force_constants()

Check cell parameters have not been changed by optimization:

print(phonons_mace.struct.cell.cellpar())

Calculate and plot band structure, writing results to Cl4Na4-auto_bands.yml, and saving the figure as Cl4Na4-auto_bands.svg:

phonons_mace.calc_bands(write_bands=True)

Calculate thermal properties, saving the heat capacity, enthalpy, and entropy, to Cl4Na4-thermal.dat:

phonons_mace.calc_thermal_props(write_thermal=True)

Phonon calcualtions with optimization of cell

The same calculations can be run with cell lengths, but not angles, optimized.

Note: Set "filter_kwargs" = {"hydrostatic_strain": True} for geometry optimization via minimize_kwargs, so cell angles are fixed, but lengths can change

phonons_mace_lengths_only = Phonons(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    supercell=[2, 2, 2],
    displacement=0.01,
    temp_step=10.0,
    temp_min=0.0,
    temp_max=1000.0,
    minimize=True,
    force_consts_to_hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_kwargs": {"hydrostatic_strain": True}},
    write_results=True,
)

phonons_mace_lengths_only.calc_bands(write_bands=True)

Confirm changes to cell lengths:

print(phonons_mace_lengths_only.struct.cell.cellpar())

Phonon calculations with pressure

Calculations can also be run at a fixed pressure, as well as optmising both the cell lengths and angles.

Note: Set "filter_kwargs" = {"scalar_pressure": x} for geometry optimization via minimize_kwargs to set the pressure. Without setting hydrostatic_strain =  True, both the cell lengths and angles will be optimized

phonons_mace_pressure = Phonons(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model_path="small",
    calc_kwargs={"default_dtype": "float64"},
    supercell=[2, 2, 2],
    displacement=0.01,
    temp_step=10.0,
    temp_min=0.0,
    temp_max=1000.0,
    minimize=True,
    force_consts_to_hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_kwargs": {"scalar_pressure": 0.1}},
    write_results=True,
)

phonons_mace_pressure.calc_bands(write_bands=True)

Confirm changes to cell:

print(phonons_mace_pressure.struct.cell.cellpar())

Compare band structures for different optimization options and save to files:

phonons_mace.write_bands(plot_file="NaCl_mace.svg")
phonons_mace_lengths_only.write_bands(plot_file="NaCl_lengths_only.svg")
phonons_mace_pressure.write_bands(plot_file="NaCl_pressure.svg")

Comparing the band structure from MACE to CHGNet and SevenNet:

Calculate band structure using CHGNet

phonons_chgnet = Phonons(
    struct=NaCl.copy(),
    arch="chgnet",
    device="cpu",
    supercell=[2, 2, 2],
    displacement=0.01,
    temp_step=10.0,
    temp_min=0.0,
    temp_max=1000.0,
    minimize=True,
    force_consts_to_hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_func": None},
    write_results=True,
)

phonons_chgnet.calc_bands(write_bands=True)

Calculate band structure using SevenNet

phonons_sevennet = Phonons(
    struct=NaCl.copy(),
    arch="sevennet",
    device="cpu",
    supercell=[2, 2, 2],
    displacement=0.01,
    temp_step=10.0,
    temp_min=0.0,
    temp_max=1000.0,
    minimize=True,
    force_consts_to_hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_func": None},
    write_results=True,
)

phonons_sevennet.calc_bands(write_bands=True)

Compare and save plots for each MLIP:

phonons_mace.write_bands(plot_file="MACE.svg")
phonons_chgnet.write_bands(plot_file="chgnet.svg")
phonons_sevennet.write_bands(plot_file="sevennet.svg")

Note: It may be necessary to reset the default PyTorch dtype if different calculators have been set:

import torch
torch.set_default_dtype(torch.float64)

phonons_mace.calc_force_constants()