Add MLP phonon workflow (#1642)

```python
from quacc.recipes.mlp.phonons import phonon_flow
from ase.io import read

atoms = read("structure.cif")
output = phonon_flow(atoms, method="mace")
print(output)
```

.github/workflows/tests.yaml

@@ -413,6 +413,7 @@ jobs:
  pip install -r tests/requirements.txt
  pip install -r tests/requirements-mlp.txt
  pip install -r tests/requirements-newtonnet.txt
+ pip install -r tests/requirements-phonopy.txt
  pip install .[dev]
 
  - name: Run tests with pytest
@@ -424,7 +425,7 @@ jobs:
  name: ${{ github.job }} coverage report
  path: "coverage.xml"
  retention-days: 1
- 
+
  codecov:
  needs:
  - tests-core

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Added
 
+- Added phonon job for MLPs
 - Added an ASE relax job recipe for ONETEP
 - Added a non-SCF job for Quantum Espresso
 - Added a DOS job for Quantum Espresso

docs/user/recipes/recipes_list.md

@@ -89,10 +89,11 @@ The list of available quacc recipes is shown below. The "Req'd Extras" column sp
 
 <center>
 
-| Name | Decorator | Documentation | Req'd Extras |
-| ---------- | --------------- | ------------------------------------- | ------------ |
-| MLP Static | `#!Python @job` | [quacc.recipes.mlp.core.static_job][] | `quacc[mlp]` |
-| MLP Relax | `#!Python @job` | [quacc.recipes.mlp.core.relax_job][] | `quacc[mlp]` |
+| Name | Decorator | Documentation | Req'd Extras |
+| ----------- | ---------------- | ----------------------------------------- | -------------------- |
+| MLP Static | `#!Python @job` | [quacc.recipes.mlp.core.static_job][] | `quacc[mlp]` |
+| MLP Relax | `#!Python @job` | [quacc.recipes.mlp.core.relax_job][] | `quacc[mlp]` |
+| MLP Phonons | `#!Python @flow` | [quacc.recipes.mlp.phonons.phonon_flow][] | `quacc[mlp,phonons]` |
 
 </center>
 

src/quacc/recipes/mlp/_base.py

@@ -14,15 +14,15 @@
 
 
 @lru_cache
-def _pick_calculator(
- calculator: Literal["mace", "m3gnet", "chgnet"], **kwargs
+def pick_calculator(
+ method: Literal["mace", "m3gnet", "chgnet"], **kwargs
 ) -> Calculator:
  """
  Adapted from `matcalc.util.get_universal_calculator`.
 
  Parameters
  ----------
- calculator
+ method
  Name of the calculator to use
  **kwargs
  Custom kwargs for the underlying calculator. Set a value to
@@ -40,7 +40,7 @@ def _pick_calculator(
  if not torch.cuda.is_available():
  logger.warning("CUDA is not available to PyTorch. Calculations will be slow.")
 
- if calculator.lower().startswith("m3gnet"):
+ if method.lower().startswith("m3gnet"):
  import matgl
  from matgl import __version__
  from matgl.ext.ase import M3GNetCalculator
@@ -49,20 +49,22 @@ def _pick_calculator(
  kwargs.setdefault("stress_weight", 1.0 / 160.21766208)
  calc = M3GNetCalculator(potential=model, **kwargs)
 
- elif calculator.lower() == "chgnet":
+ elif method.lower() == "chgnet":
  from chgnet import __version__
  from chgnet.model.dynamics import CHGNetCalculator
 
  calc = CHGNetCalculator(**kwargs)
 
- elif calculator.lower() == "mace":
+ elif method.lower() == "mace":
  from mace import __version__
  from mace.calculators import mace_mp
 
+ if "default_dtype" not in kwargs:
+ kwargs["default_dtype"] = "float64"
  calc = mace_mp(**kwargs)
 
  else:
- raise ValueError(f"Unrecognized {calculator=}.")
+ raise ValueError(f"Unrecognized {method=}.")
 
  calc.parameters["version"] = __version__
 

src/quacc/recipes/mlp/core.py

@@ -6,7 +6,7 @@
 from ase.optimize import FIRE
 
 from quacc import job
-from quacc.recipes.mlp._base import _pick_calculator
+from quacc.recipes.mlp._base import pick_calculator
 from quacc.runners.ase import run_calc, run_opt
 from quacc.schemas.ase import summarize_opt_run, summarize_run
 from quacc.utils.dicts import recursive_dict_merge
@@ -47,7 +47,7 @@ def static_job(
  calc_defaults = {"default_dtype": "float64"} if method == "mace" else {}
  calc_flags = recursive_dict_merge(calc_defaults, calc_kwargs)
 
- atoms.calc = _pick_calculator(method, **calc_flags)
+ atoms.calc = pick_calculator(method, **calc_flags)
  final_atoms = run_calc(atoms)
  return summarize_run(
  final_atoms, atoms, additional_fields={"name": f"{method} Static"}
@@ -90,13 +90,10 @@ def relax_job(
  See the type-hint for the data structure.
  """
 
- calc_defaults = {"default_dtype": "float64"} if method == "mace" else {}
- calc_flags = recursive_dict_merge(calc_defaults, calc_kwargs)
-
  opt_defaults = {"fmax": 0.05, "max_steps": 1000, "optimizer": FIRE}
  opt_flags = recursive_dict_merge(opt_defaults, opt_params)
 
- atoms.calc = _pick_calculator(method, **calc_flags)
+ atoms.calc = pick_calculator(method, **calc_kwargs)
 
  dyn = run_opt(atoms, relax_cell=relax_cell, **opt_flags)
 

src/quacc/recipes/mlp/phonons.py

@@ -0,0 +1,104 @@
+"""Phonon recipes for MLPs."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from quacc import flow
+from quacc.recipes.common.phonons import phonon_flow as common_phonon_flow
+from quacc.recipes.mlp.core import relax_job, static_job
+from quacc.utils.dicts import recursive_dict_merge
+from quacc.wflow_tools.customizers import customize_funcs
+
+if TYPE_CHECKING:
+ from typing import Any, Callable, Literal
+
+ from ase.atoms import Atoms
+
+ from quacc.schemas._aliases.phonons import PhononSchema
+
+
+@flow
+def phonon_flow(
+ atoms: Atoms,
+ method: Literal["mace", "m3gnet", "chgnet"],
+ symprec: float = 1e-4,
+ min_length: float | None = 15.0,
+ displacement: float = 0.01,
+ t_step: float = 10,
+ t_min: float = 0,
+ t_max: float = 1000,
+ run_relax: bool = True,
+ job_params: dict[str, dict[str, Any]] | None = None,
+ job_decorators: dict[str, Callable | None] | None = None,
+) -> PhononSchema:
+ """
+ Carry out a phonon workflow, consisting of:
+
+ 1. Optional relaxation.
+ - name: "relax_job"
+ - job: [quacc.recipes.mlp.core.relax_job][]
+
+ 2. Generation of supercells.
+
+ 3. Static calculations on supercells
+ - name: "static_job"
+ - job: [quacc.recipes.mlp.core.static_job][]
+
+ 4. Calculation of thermodynamic properties.
+
+ Parameters
+ ----------
+ atoms
+ Atoms object
+ method
+ Universal ML interatomic potential method to use
+ symprec
+ Precision for symmetry detection.
+ min_length
+ Minimum length of each lattice dimension (A).
+ displacement
+ Atomic displacement (A).
+ t_step
+ Temperature step (K).
+ t_min
+ Min temperature (K).
+ t_max
+ Max temperature (K).
+ job_params
+ Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+ the keys are the names of the jobs and the values are dictionaries of parameters.
+ job_decorators
+ Custom decorators to apply to each Job in the Flow. This is a dictionary where
+ the keys are the names of the jobs and the values are decorators.
+
+ Returns
+ -------
+ PhononSchema
+ Dictionary of results from [quacc.schemas.phonons.summarize_phonopy][].
+ See the type-hint for the data structure.
+ """
+ calc_defaults = {
+ "relax_job": {"method": method, "opt_params": {"fmax": 1e-3}},
+ "static_job": {"method": method},
+ }
+ job_params = recursive_dict_merge(calc_defaults, job_params)
+
+ relax_job_, static_job_ = customize_funcs(
+ ["relax_job", "static_job"],
+ [relax_job, static_job],
+ parameters=job_params,
+ decorators=job_decorators,
+ )
+
+ return common_phonon_flow(
+ atoms,
+ static_job_,
+ relax_job=relax_job_ if run_relax else None,
+ symprec=symprec,
+ min_length=min_length,
+ displacement=displacement,
+ t_step=t_step,
+ t_min=t_min,
+ t_max=t_max,
+ additional_fields={"name": f"{method} Phonons"},
+ )

tests/core/recipes/mlp_recipes/test_phonon_recipes.py

@@ -0,0 +1,24 @@
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("mace")
+pytest.importorskip("phonopy")
+import numpy as np
+from ase.build import bulk
+
+from quacc.recipes.mlp.phonons import phonon_flow
+
+
+def _set_dtype(size, type_="float"):
+ globals()[f"{type_}_th"] = getattr(torch, f"{type_}{size}")
+ globals()[f"{type_}_np"] = getattr(np, f"{type_}{size}")
+ torch.set_default_dtype(getattr(torch, f"float{size}"))
+
+
+def test_phonon_flow(tmp_path, monkeypatch):
+ monkeypatch.chdir(tmp_path)
+ _set_dtype(64)
+ atoms = bulk("Cu")
+ output = phonon_flow(atoms, method="mace", min_length=5.0)
+ assert output["results"]["force_constants"].shape == (8, 8, 3, 3)
+ assert len(output["results"]["thermal_properties"]["temperatures"]) == 101