Harmonic lattice dynamics with hiPhive and a machine-learned potential

This notebook runs the hiPhive phonon workflow end to end using MACE-OMAT-0-medium for the forces. No DFT is needed, so it runs on a laptop.

Background

hiPhive extracts interatomic force constants by fitting a cluster expansion of the force-constant potential to a set of randomly displaced supercells. It plays the same role as pheasy: both replace phonopy’s symmetry-reduced finite differences with a regression over far fewer, randomly displaced configurations.

The workflow is identical to the pheasy one apart from the fitting step. ALM counts the irreducible force-constant elements and sizes the displacement set, phonopy generates the supercells, and phonopy post-processes the fitted force constants into band structures, densities of states and thermodynamic properties. Only the fit itself is hiPhive.

Second-order force constants are physically sparse, so the fit uses LASSO by default, matching pheasy. The Huang and Born-Huang rotational sum rules are enforced on the fitted parameters.

The same workflow is available for VASP as atomate2.vasp.flows.hiphive.PhononMaker.

Installation

Two extras are needed, phonons and hiphive:

pip install 'atomate2[phonons,hiphive]'
pip install 'mace-torch>=0.3.16' mp-api

The hiphive extra pulls in hiPhive, trainstation and ALM. ALM is compiled from source and needs a C++ toolchain with OpenMP, plus the Eigen, Boost and spglib headers. The macOS system clang provides none of them, so install them first (this is the same set atomate2’s CI uses):

conda install -c conda-forge compilers llvm-openmp "eigen=3.3" boost cmake spglib

On macOS, ALM’s build also applies its C++ flags to a C source file, which clang rejects. Prefix the atomate2 install with CFLAGS="-x c++" to force that file to compile as C++.

The potential

MACE-OMAT-0-medium downloads once and caches under ~/.cache/mace. It is released under the Academic Software License.

float64 is worth the cost here. The fit reads forces from 0.01 A displacements, where float32 noise is not negligible.

import os
import warnings

# macOS: conda's llvm-openmp and torch's bundled libomp both load and the
# duplicate aborts the process. This must be set before torch is imported,
# so the mace import below is deliberately not at the top of the cell.
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("OMP_NUM_THREADS", "1")
warnings.filterwarnings("ignore")

from mace.calculators import mace_mp  # noqa: E402

MODEL = "medium-omat-0"
CALC_KWARGS = {"model": MODEL, "device": "cpu", "default_dtype": "float64"}

_ = mace_mp(**CALC_KWARGS)  # downloads and caches on first use
MODEL

The materials

These 14 entries are two per crystal system, all of which already have pheasy phonons in the Materials Project, so the result can be compared against DFT. They were selected with energy_above_hull < 0.05, at least two elements, and spglib agreeing with the crystal system MP reports.

QUICK is one material per crystal system, chosen so the whole loop finishes in about five minutes. The two omitted from it, HfO2 and K2ZnSi3O8, each take roughly half an hour: HfO2 has a 432-atom supercell and K2ZnSi3O8 is triclinic with over 5000 free parameters.

MATERIALS = {
    "mp-861479": "PrYMg2",  # cubic,        spg 225
    "mp-861950": "Ca2SbAu",  # cubic,        spg 225
    "mp-867515": "NaCoO2",  # hexagonal,    spg 194
    "mp-867841": "GaTc",  # hexagonal,    spg 187
    "mp-985829": "HfS2",  # trigonal,     spg 164
    "mp-998561": "CsSrCl3",  # trigonal,     spg 161
    "mp-979115": "Ti2Ag",  # tetragonal,   spg 139
    "mp-985278": "AcF3",  # tetragonal,   spg 139
    "mp-784630": "Ni2Mo",  # orthorhombic, spg 71
    "mp-775757": "HfO2",  # orthorhombic, spg 61,  slow
    "mp-997089": "NaAuO2",  # monoclinic,   spg 12
    "mp-863754": "KYSiS4",  # monoclinic,   spg 4
    "mp-997036": "CdAuO2",  # triclinic,    spg 2
    "mp-1224204": "K2ZnSi3O8",  # triclinic,    spg 1,   slow
}

# one per crystal system, about five minutes in total
QUICK = [
    "mp-861479",
    "mp-867841",
    "mp-985829",
    "mp-985278",
    "mp-784630",
    "mp-997089",
    "mp-997036",
]

len(MATERIALS), len(QUICK)

Pulling the structures

Each pheasy entry stores the supercell it used. Reusing that supercell keeps supercell size from becoming a variable when comparing against DFT.

Two things to know about this endpoint. Its primary key is identifier, and material_id is a filter rather than a returnable field. And it is delta-backed, so a search without an explicit num_chunks downloads the whole collection to disk.

Set your key first with export MP_API_KEY=....

from mp_api.client import MPRester
from pymatgen.core import Structure

if not os.environ.get("MP_API_KEY"):
    raise OSError("export MP_API_KEY before running this cell")


def fetch(mp_id: str) -> tuple[Structure, list[list[int]]]:
    """Primitive cell and supercell matrix from the pheasy phonon entry."""
    with MPRester() as mpr:
        hits = mpr.materials.phonon.search(
            material_ids=[mp_id],
            phonon_method="pheasy",
            num_chunks=1,  # bounded: avoids pulling the whole delta table
            chunk_size=4,
            fields=["identifier", "phonon_method", "structure", "supercell_matrix"],
        )
    hits = [h for h in hits if str(h.phonon_method) == "pheasy"]
    if not hits:
        raise ValueError(f"no pheasy phonon entry for {mp_id}")
    struct = hits[0].structure
    if isinstance(struct, dict):
        struct = Structure.from_dict(struct)
    sc_matrix = [[round(x) for x in row] for row in hits[0].supercell_matrix]
    return struct, sc_matrix


structure, supercell = fetch("mp-861479")
structure.composition.reduced_formula, supercell

Building the workflow

The forcefield PhononMaker needs its relaxation and static makers built explicitly, because the defaults use a different MACE model. Everything else stays at the workflow defaults: LASSO fitting, 0.01 A displacements, and the ALM-based configuration count.

from atomate2.forcefields.flows.hiphive import PhononMaker
from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker

# any MACE member routes through mace_mp; the model name in calculator_kwargs
# selects which weights are used
FF_NAME = "MACE-MPA-0"


def make_maker() -> PhononMaker:
    """HiPhive PhononMaker wired to MACE-OMAT-0-medium."""
    return PhononMaker(
        bulk_relax_maker=ForceFieldRelaxMaker(
            force_field_name=FF_NAME,
            calculator_kwargs=CALC_KWARGS,
            relax_kwargs={"fmax": 1e-5},
        ),
        phonon_displacement_maker=ForceFieldStaticMaker(
            force_field_name=FF_NAME, calculator_kwargs=CALC_KWARGS
        ),
        static_energy_maker=None,
        born_maker=None,
        create_thermal_displacements=False,
    )


flow = make_maker().make(structure=structure, supercell_matrix=supercell)
flow.draw_graph().show()

Running one material

The flow relaxes the cell, generates randomly displaced supercells, evaluates their forces with MACE, fits the force constants with hiPhive, and produces the band structure and DOS. PrYMg2 takes well under a minute.

from jobflow import run_locally

responses = run_locally(flow, create_folders=True, ensure_success=True)
doc = responses[flow.jobs[-1].uuid][1].output

# the output document keeps frequencies as plain lists; to_pmg returns the
# pymatgen objects that the plotters and the analysis helpers expect
band_structure = doc.phonon_bandstructure.to_pmg
round(float(band_structure.bands.max()), 3), doc.has_imaginary_modes

from pymatgen.phonon.plotter import PhononBSPlotter, PhononDosPlotter

dos_plot = PhononDosPlotter()
dos_plot.add_dos(label=MATERIALS["mp-861479"], dos=doc.phonon_dos.to_pmg)
dos_plot.get_plot()

bs_plot = PhononBSPlotter(bs=band_structure)
bs_plot.get_plot()

Running the quick set

One material per crystal system. Results are written after each one, so the loop can be interrupted and resumed. Swap QUICK for MATERIALS to run all 14, which takes about an hour and a half.

Ti2Ag is in MATERIALS but not in QUICK. Its 2x2x2 supercell of a body-centred cell gives a pair cutoff that reaches no neighbour shell, and the workflow raises a clear error saying so.

import json
import time
from pathlib import Path

OUT = Path("hiphive_mace_results.json")
results = json.loads(OUT.read_text()) if OUT.exists() else {}

for mp_id in QUICK:
    if mp_id in results:
        continue
    formula = MATERIALS[mp_id]
    started = time.time()
    try:
        struct, sc_matrix = fetch(mp_id)
        job = make_maker().make(structure=struct, supercell_matrix=sc_matrix)
        out = run_locally(job, create_folders=True, ensure_success=True)[
            job.jobs[-1].uuid
        ][1].output
        freqs = out.phonon_bandstructure.to_pmg.bands
        results[mp_id] = {
            "formula": formula,
            "min_freq_THz": round(float(freqs.min()), 4),
            "max_freq_THz": round(float(freqs.max()), 4),
            "imaginary": bool(out.has_imaginary_modes),
            "minutes": round((time.time() - started) / 60, 1),
        }
    except Exception as exc:  # noqa: BLE001
        results[mp_id] = {"formula": formula, "error": f"{type(exc).__name__}: {exc}"}
    print(mp_id, results[mp_id])  # noqa: T201
    OUT.write_text(json.dumps(results, indent=1))

Comparing against DFT

The same entries have DFT force constants in the Materials Project. phonon_method has to be passed explicitly: the convenience accessors on MPRester default to "dfpt" and will not find a pheasy entry.

import numpy as np
import pandas as pd

rows = []
with MPRester() as mpr:
    for mp_id, row in results.items():
        if "error" in row:
            rows.append(
                {
                    "material_id": mp_id,
                    "formula": row["formula"],
                    "note": row["error"][:40],
                }
            )
            continue
        band_struct = mpr.materials.phonon.get_bandstructure_from_material_id(
            mp_id, phonon_method="pheasy"
        )
        dft_max = float(np.asarray(band_struct.frequencies).max())
        rows.append(
            {
                "material_id": mp_id,
                "formula": row["formula"],
                "dft_max_THz": round(dft_max, 3),
                "mace_max_THz": row["max_freq_THz"],
                "diff_pct": round(100 * (row["max_freq_THz"] - dft_max) / dft_max, 1),
                "note": "",
            }
        )

pd.DataFrame(rows)

Reading the results

Across this set the MLIP maximum frequency lands within about 10 percent of DFT. Most entries come out a few percent low. HfS2 comes out slightly high. That gap is the potential, not the fit. To separate the two, run the same material through the phonopy workflow with the identical potential and supercell (atomate2.forcefields.flows.phonons.PhononMaker). Where phonopy and hiPhive agree with each other, the fit is working and what remains is the error in MACE.

Known limitations

The workflow has only been tested on the 14 materials listed above so far, not at high-throughput scale. What follows is what that test surfaced.

Low symmetry. On genuinely triclinic cells the cluster space grows into the thousands of parameters while the ALM-sized displacement set does not grow with it, and the LASSO fit can produce spurious soft modes. K2ZnSi3O8 is the example here. Raising the number of configurations does not fix it, so compare against phonopy before trusting a P1 result.

Small or skewed supercells. A diagonal repeat of a centred primitive cell can be skewed enough that the largest usable pair cutoff reaches no neighbour shell, leaving the cluster space empty. Ti2Ag is the example. Use a larger or non-diagonal supercell_matrix for such cells.

Few displacements. When the symmetry-reduced finite-displacement count is three or fewer, the workflow uses those displacements with an ordinary least squares fit instead of random displacements with LASSO. This is inherited from pheasy and affects the most symmetric materials here.