Phonons¶

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.

[1]:

# ! 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/

[2]:

from janus_core import __version__

print(__version__)

0.9.4

Phonons (periodic)¶

[3]:

from ase.build import bulk
from weas_widget import WeasWidget

from janus_core.calculations.phonons import Phonons

Prepare for phonon calculations on salt¶

[4]:

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

[4]:

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

[5]:

phonons_mace = Phonons(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model="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,
    hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_class": None},
    write_results=True,
)

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

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:

[6]:

phonons_mace.calc_force_constants()

Check cell parameters have not been changed by optimization:

[7]:

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

[ 5.63  5.63  5.63 90.   90.   90.  ]

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

[8]:

phonons_mace.calc_bands(write_bands=True)

../../_images/tutorials_python_phonons_17_0.png

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

[9]:

phonons_mace.calc_thermal_props(write_thermal=True)

Phonon calculations 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

[10]:

phonons_mace_lengths_only = Phonons(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model="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,
    hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_kwargs": {"hydrostatic_strain": True}},
    write_results=True,
)

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)

[11]:

phonons_mace_lengths_only.calc_bands(write_bands=True)

       Step     Time          Energy          fmax
LBFGS:    0 20:44:07      -27.030082        0.129748
LBFGS:    1 20:44:07      -27.030795        0.126795
LBFGS:    2 20:44:07      -27.046337        0.004990

../../_images/tutorials_python_phonons_23_1.png

Confirm changes to cell lengths:

[12]:

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

[ 5.68761079  5.68761079  5.68761079 90.         90.         90.        ]

Phonon calculations with pressure¶

Calculations can also be run at a fixed pressure, as well as optimising 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

[13]:

phonons_mace_pressure = Phonons(
    struct=NaCl.copy(),
    arch="mace_mp",
    device="cpu",
    model="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,
    hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_kwargs": {"scalar_pressure": 0.1}},
    write_results=True,
)

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)

[14]:

phonons_mace_pressure.calc_bands(write_bands=True)

       Step     Time          Energy          fmax
LBFGS:    0 20:44:09      -26.918700        0.115825
LBFGS:    1 20:44:09      -26.919268        0.113180
LBFGS:    2 20:44:09      -26.931573        0.003954

../../_images/tutorials_python_phonons_29_1.png

Confirm changes to cell:

[15]:

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

[ 5.68122692  5.68122692  5.68122692 90.         90.         90.        ]

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

[16]:

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

../../_images/tutorials_python_phonons_33_0.png

../../_images/tutorials_python_phonons_33_1.png

../../_images/tutorials_python_phonons_33_2.png

Comparing the band structure from MACE to CHGNet and Orb:¶

Calculate band structure using CHGNet¶

[17]:

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,
    hdf5=True,
    plot_to_file=True,
    symmetrize=False,
    write_full=True,
    minimize_kwargs={"filter_class": None},
    write_results=True,
)

CHGNet v0.3.0 initialized with 412,525 parameters
CHGNet will run on cpu

[18]:

phonons_chgnet.calc_bands(write_bands=True)

       Step     Time          Energy          fmax
LBFGS:    0 20:44:14      -29.326620        0.000000

/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/chgnet/model/model.py:898: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.
Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:835.)
  volumes = torch.tensor(volumes, dtype=TORCH_DTYPE, device=atomic_numbers.device)

../../_images/tutorials_python_phonons_37_2.png

Calculate band structure using Orb¶

[19]:

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

/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/orb_models/utils.py:30: UserWarning: Setting global torch default dtype to torch.float32.
  warnings.warn(f"Setting global torch default dtype to {torch_dtype}.")

[20]:

phonons_orb.calc_bands(write_bands=True)

       Step     Time          Energy          fmax
LBFGS:    0 20:44:28      -27.077099        0.001187

../../_images/tutorials_python_phonons_40_1.png

Compare and save plots for each MLIP:

[21]:

phonons_mace.write_bands(plot_file="MACE.svg")
phonons_chgnet.write_bands(plot_file="chgnet.svg")
phonons_orb.write_bands(plot_file="orb.svg")

../../_images/tutorials_python_phonons_42_0.png

../../_images/tutorials_python_phonons_42_1.png

../../_images/tutorials_python_phonons_42_2.png

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

[22]:

import torch
torch.set_default_dtype(torch.float64)

phonons_mace.calc_force_constants()