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.
This will replace pre-installed versions of numpy and torch in Colab with versions that are known to be compatible with janus-core.
It may be possible to skip the steps that uninstall and reinstall torch, which will save a considerable amount of time.
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]:
# import locale
# locale.getpreferredencoding = lambda: "UTF-8"
# ! pip uninstall numpy -y # Uninstall pre-installed numpy
# ! pip uninstall torch torchaudio torchvision transformers -y # Uninstall pre-installed torch
# ! uv pip install torch==2.5.1 # Install pinned version of torch
# ! uv pip install janus-core[mace,orb,chgnet,visualise] data-tutorials --system # Install janus-core with MACE, Orb, CHGNet, and WeasWidget, and data-tutorials
# get_ipython().kernel.do_shutdown(restart=True) # Restart kernel to update libraries. This may warn that your session has crashed.
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.8.5
Phonons (periodic)¶
[3]:
from ase.build import bulk
from weas_widget import WeasWidget
from janus_core.calculations.phonons import Phonons
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/codecarbon/core/gpu.py:4: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml
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.
Using head Default out of ['Default']
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:197: 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)
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
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.
Using head Default out of ['Default']
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:197: 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 16:32:07      -27.030082        0.129748
LBFGS:    1 16:32:07      -27.030795        0.126795
LBFGS:    2 16:32:07      -27.046337        0.004990
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
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.
Using head Default out of ['Default']
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/mace/calculators/mace.py:197: 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 16:32:09      -26.918700        0.115825
LBFGS:    1 16:32:09      -26.919268        0.113180
LBFGS:    2 16:32:09      -26.931573        0.003954
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
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")
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
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 16:32:15      -29.326620        0.000000
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
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}.")
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
[20]:
phonons_orb.calc_bands(write_bands=True)
       Step     Time          Energy          fmax
LBFGS:    0 16:32:23      -27.077099        0.001187
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
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")
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
/home/runner/work/janus-core/janus-core/.venv/lib/python3.12/site-packages/seekpath/hpkot/__init__.py:156: DeprecationWarning: dict interface is deprecated. Use attribute interface instead
  conv_lattice = dataset['std_lattice']
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()