Merge pull request #398 from JaGeo/forcefield-phonons

Phonons for forcefields

pyproject.toml

@@ -43,7 +43,7 @@ mp = ["mp-api>=0.27.5"]
 phonons = ["phonopy>=1.10.8", "seekpath"]
 lobster = ["lobsterpy>=0.3.0"]
 defects = ["dscribe>=1.2.0", "pymatgen-analysis-defects>=2022.11.30"]
-forcefields = ["chgnet==0.2.0", "matgl==0.8.2"]
+forcefields = ["chgnet==0.2.0", "matgl==0.8.2", "quippy-ase==0.9.14"]
 docs = [
     "FireWorks==2.0.3",
     "autodoc_pydantic==1.9.0",
@@ -75,6 +75,7 @@ strict = [
     "pydantic==1.10.9",
     "pymatgen-analysis-defects==2023.7.31",
     "pymatgen==2023.8.10",
+    "quippy-ase==0.9.14",
     "seekpath==2.1.0",
     "typing-extensions==4.7.1",
 ]

src/atomate2/vasp/jobs/phonons.py → src/atomate2/common/jobs/phonons.py

@@ -2,31 +2,33 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 from jobflow import Flow, Response, job
 from phonopy import Phonopy
 from phonopy.units import VaspToTHz
+from pymatgen.core import Structure
 from pymatgen.io.phonopy import get_phonopy_structure, get_pmg_structure
 from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
 from pymatgen.phonon.dos import PhononDos
 from pymatgen.transformations.advanced_transformations import (
     CubicSupercellTransformation,
 )
 
+from atomate2.common.schemas.phonons import ForceConstants, PhononBSDOSDoc
 from atomate2.vasp.jobs.base import BaseVaspMaker
-from atomate2.vasp.schemas.phonons import PhononBSDOSDoc
 from atomate2.vasp.sets.core import StaticSetGenerator
 
 if TYPE_CHECKING:
     from pathlib import Path
 
     import numpy as np
     from emmet.core.math import Matrix3D
-    from pymatgen.core import Structure
 
+    from atomate2.forcefields.jobs import ForceFieldStaticMaker
     from atomate2.vasp.sets.base import VaspInputGenerator
 
 logger = logging.getLogger(__name__)
@@ -124,7 +126,7 @@ def get_supercell_size(
     return transformation.transformation_matrix.tolist()
 
 
-@job
+@job(data=[Structure])
 def generate_phonon_displacements(
     structure: Structure,
     supercell_matrix: np.array,
@@ -185,7 +187,10 @@ def generate_phonon_displacements(
     return [get_pmg_structure(cell) for cell in supercells]
 
 
-@job(output_schema=PhononBSDOSDoc, data=[PhononDos, PhononBandStructureSymmLine])
+@job(
+    output_schema=PhononBSDOSDoc,
+    data=[PhononDos, PhononBandStructureSymmLine, ForceConstants],
+)
 def generate_frequencies_eigenvectors(
     structure: Structure,
     supercell_matrix: np.array,
@@ -251,12 +256,12 @@ def generate_frequencies_eigenvectors(
     )
 
 
-@job
+@job(data=["forces", "displaced_structures"])
 def run_phonon_displacements(
     displacements,
     structure: Structure,
     supercell_matrix,
-    phonon_maker: BaseVaspMaker = None,
+    phonon_maker: BaseVaspMaker | ForceFieldStaticMaker = None,
     prev_vasp_dir: str | Path = None,
 ):
     """
@@ -268,7 +273,7 @@ def run_phonon_displacements(
     ----------
     displacements
     structure: Structure object
-        Fully optimized structure used for phonon computations
+        Fully optimized structure used for phonon computations.
     supercell_matrix: Matrix3D
         supercell matrix for meta data
     phonon_maker : .BaseVaspMaker
@@ -284,11 +289,13 @@ def run_phonon_displacements(
         "forces": [],
         "uuids": [],
         "dirs": [],
-        "displaced_structures": [],
     }
 
     for i, displacement in enumerate(displacements):
-        phonon_job = phonon_maker.make(displacement, prev_vasp_dir=prev_vasp_dir)
+        if prev_vasp_dir is not None:
+            phonon_job = phonon_maker.make(displacement, prev_vasp_dir=prev_vasp_dir)
+        else:
+            phonon_job = phonon_maker.make(displacement)
         phonon_job.append_name(f" {i + 1}/{len(displacements)}")
 
         # we will add some meta data
@@ -298,15 +305,17 @@ def run_phonon_displacements(
             "supercell_matrix": supercell_matrix,
             "displaced_structure": displacement,
         }
-        phonon_job.update_maker_kwargs(
-            {"_set": {"write_additional_data->phonon_info:json": info}}, dict_mod=True
-        )
+        with contextlib.suppress(Exception):
+            phonon_job.update_maker_kwargs(
+                {"_set": {"write_additional_data->phonon_info:json": info}},
+                dict_mod=True,
+            )
+
         phonon_jobs.append(phonon_job)
         outputs["displacement_number"].append(i)
         outputs["uuids"].append(phonon_job.output.uuid)
         outputs["dirs"].append(phonon_job.output.dir_name)
         outputs["forces"].append(phonon_job.output.output.forces)
-        outputs["displaced_structures"].append(displacement)
 
     displacement_flow = Flow(phonon_jobs, outputs)
     return Response(replace=displacement_flow)

src/atomate2/vasp/schemas/phonons.py → src/atomate2/common/schemas/phonons.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 from emmet.core.math import Matrix3D
+from monty.json import MSONable
 from phonopy import Phonopy
 from phonopy.phonon.band_structure import get_band_qpoints_and_path_connections
 from phonopy.structure.symmetry import symmetrize_borns_and_epsilon
@@ -80,19 +81,26 @@ class PhononUUIDs(BaseModel):
     born_run_uuid: str = Field(None, description="born run uuid")
 
 
+class ForceConstants(MSONable):
+    """A force constants class."""
+
+    def __init__(self, force_constants: List[List[Matrix3D]]):
+        self.force_constants = force_constants
+
+
 class PhononJobDirs(BaseModel):
     """Collection to save all job directories relevant for the phonon run."""
 
-    displacements_job_dirs: List[str] = Field(
+    displacements_job_dirs: List[Optional[str]] = Field(
         None, description="The directories where the displacement jobs were run."
     )
-    static_run_job_dir: str = Field(
+    static_run_job_dir: Optional[str] = Field(
         None, description="Directory where static run was performed."
     )
-    born_run_job_dir: str = Field(
+    born_run_job_dir: Optional[str] = Field(
         None, description="Directory where born run was performed."
     )
-    optimization_run_job_dir: str = Field(
+    optimization_run_job_dir: Optional[str] = Field(
         None, description="Directory where optimization run was performed."
     )
 
@@ -152,7 +160,7 @@ class PhononBSDOSDoc(BaseModel):
     )
 
     # needed, e.g. to compute Grueneisen parameter etc
-    force_constants: List[List[Matrix3D]] = Field(
+    force_constants: ForceConstants = Field(
         None, description="Force constants between every pair of atoms in the structure"
     )
 
@@ -181,7 +189,9 @@ class PhononBSDOSDoc(BaseModel):
         "Includes all data of the computation of the thermal displacements"
     )
 
-    jobdirs: PhononJobDirs = Field("Field including all relevant job directories")
+    jobdirs: Optional[PhononJobDirs] = Field(
+        "Field including all relevant job directories"
+    )
 
     uuids: PhononUUIDs = Field("Field including all relevant uuids")
 
@@ -223,7 +233,7 @@ def from_forces_born(
         code: str
             which code was used for computation
         displacement_data:
-            output of the VASP displacement data
+            output of the VASP displacement runs
         total_dft_energy: float
             total energy in eV per cell
         epsilon_static: Matrix3D
@@ -258,7 +268,6 @@ def from_forces_born(
         )
         phonon.generate_displacements(distance=displacement)
         set_of_forces = [np.array(forces) for forces in displacement_data["forces"]]
-
         if born is not None and epsilon_static is not None:
             if len(structure) == len(born):
                 borns, epsilon = symmetrize_borns_and_epsilon(
@@ -305,8 +314,13 @@ def from_forces_born(
 
         # phonon band structures will always be cmouted
         filename_band_yaml = "phonon_band_structure.yaml"
+
+        # TODO: potentially add kwargs to avoid computation of eigenvectors
         phonon.run_band_structure(
-            qpoints, path_connections=connections, with_eigenvectors=True
+            qpoints,
+            path_connections=connections,
+            with_eigenvectors=kwargs.get("band_structure_eigenvectors", False),
+            is_band_connection=kwargs.get("band_structure_eigenvectors", False),
         )
         phonon.write_yaml_band_structure(filename=filename_band_yaml)
         bs_symm_line = get_ph_bs_symm_line(
@@ -326,7 +340,8 @@ def from_forces_born(
         )
 
         # gets data for visualization on website - yaml is also enough
-        bs_symm_line.write_phononwebsite("phonon_website.json")
+        if kwargs.get("band_structure_eigenvectors", False):
+            bs_symm_line.write_phononwebsite("phonon_website.json")
 
         # get phonon density of states
         filename_dos_yaml = "phonon_dos.yaml"
@@ -381,7 +396,7 @@ def from_forces_born(
         # will compute thermal displacement matrices
         # for the primitive cell (phonon.primitive!)
         # only this is available in phonopy
-        if kwargs["create_thermal_displacements"]:
+        if kwargs.get("create_thermal_displacements", False):
             phonon.run_mesh(
                 kpoint.kpts[0], with_eigenvectors=True, is_mesh_symmetry=False
             )
@@ -438,7 +453,7 @@ def from_forces_born(
             temperatures=temperature_range.tolist(),
             total_dft_energy=total_dft_energy_per_formula_unit,
             has_imaginary_modes=imaginary_modes,
-            force_constants=phonon.force_constants.tolist()
+            force_constants={"force_constants": phonon.force_constants.tolist()}
             if kwargs["store_force_constants"]
             else None,
             born=borns.tolist() if borns is not None else None,
@@ -452,7 +467,7 @@ def from_forces_born(
                 "thermal_displacement_matrix": tdisp_mat,
                 "freq_min_thermal_displacements": freq_min_thermal_displacements,
             }
-            if kwargs["create_thermal_displacements"]
+            if kwargs.get("create_thermal_displacements", False)
             else None,
             jobdirs={
                 "displacements_job_dirs": displacement_data["dirs"],

src/atomate2/forcefields/__init__.py

@@ -0,0 +1 @@
+"""Tools and functions common to all forcefields."""

src/atomate2/forcefields/flows/__init__.py

@@ -0,0 +1,2 @@
+"""Flows for forcefields."""
+from .relax import CHGNetVaspRelaxMaker, M3GNetVaspRelaxMaker