Fix force constants

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[str | None] = Field(
         None, description="The directories where the displacement jobs were run."
     )
-    static_run_job_dir: str = Field(
+    static_run_job_dir: str | None = Field(
         None, description="Directory where static run was performed."
     )
-    born_run_job_dir: str = Field(
+    born_run_job_dir: str | None = Field(
         None, description="Directory where born run was performed."
     )
-    optimization_run_job_dir: str = Field(
+    optimization_run_job_dir: str | None = 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: Optional[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")
 
@@ -308,7 +318,10 @@ def from_forces_born(
 
         # TODO: potentially add kwargs to avoid computation of eigenvectors
         phonon.run_band_structure(
-            qpoints, path_connections=connections, with_eigenvectors=kwargs.get("band_structure_eigenvectors", True), is_band_connection=kwargs.get("band_structure_eigenvectors", True),
+            qpoints,
+            path_connections=connections,
+            with_eigenvectors=kwargs.get("band_structure_eigenvectors", True),
+            is_band_connection=kwargs.get("band_structure_eigenvectors", True),
         )
         phonon.write_yaml_band_structure(filename=filename_band_yaml)
         bs_symm_line = get_ph_bs_symm_line(
@@ -441,7 +454,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,

src/atomate2/forcefields/flows/phonons.py

@@ -8,11 +8,8 @@
 from jobflow import Flow, Maker
 
 from atomate2.common.jobs.utils import structure_to_conventional, structure_to_primitive
-from atomate2.forcefields.jobs import CHGNetStaticMaker, CHGNetRelaxMaker
-from atomate2.vasp.flows.core import DoubleRelaxMaker
-from atomate2.vasp.jobs.core import DielectricMaker, StaticMaker, TightRelaxMaker
+from atomate2.forcefields.jobs import CHGNetRelaxMaker, CHGNetStaticMaker
 from atomate2.vasp.jobs.phonons import (
-    PhononDisplacementMaker,
     generate_frequencies_eigenvectors,
     generate_phonon_displacements,
     get_supercell_size,
@@ -21,8 +18,6 @@
 )
 
 if TYPE_CHECKING:
-    from pathlib import Path
-
     from emmet.core.math import Matrix3D
     from pymatgen.core.structure import Structure
 
@@ -134,13 +129,9 @@ class PhononMaker(Maker):
     prefer_90_degrees: bool = True
     get_supercell_size_kwargs: dict = field(default_factory=dict)
     use_symmetrized_structure: str | None = None
-    bulk_relax_maker: BaseVaspMaker | None = field(
-        default_factory= CHGNetRelaxMaker
-    )
+    bulk_relax_maker: BaseVaspMaker | None = field(default_factory=CHGNetRelaxMaker)
     static_energy_maker: BaseVaspMaker | None = field(default_factory=CHGNetStaticMaker)
-    phonon_displacement_maker: BaseVaspMaker = field(
-        default_factory=CHGNetStaticMaker
-    )
+    phonon_displacement_maker: BaseVaspMaker = field(default_factory=CHGNetStaticMaker)
     create_thermal_displacements: bool = True
     generate_frequencies_eigenvectors_kwargs: dict = field(default_factory=dict)
     kpath_scheme: str = "seekpath"
@@ -229,8 +220,8 @@ def make(
             bulk = self.bulk_relax_maker.make(structure)
             jobs.append(bulk)
             structure = bulk.output.structure
-            optimization_run_job_dir = None # TODO: should we save something?
-            #bulk.output.dir_name
+            optimization_run_job_dir = None  # TODO: should we save something?
+            # bulk.output.dir_name
             optimization_run_uuid = bulk.output.uuid
         else:
             optimization_run_job_dir = None
@@ -278,7 +269,7 @@ def make(
             static_job = self.static_energy_maker.make(structure=structure)
             jobs.append(static_job)
             total_dft_energy = static_job.output.output.energy
-            #static_run_job_dir = static_job.output.dir_name
+            # static_run_job_dir = static_job.output.dir_name
             static_run_job_dir = None
             static_run_uuid = static_job.output.uuid
         else:
@@ -294,8 +285,6 @@ def make(
             static_run_job_dir = None
             static_run_uuid = None
 
-
-
         phonon_collect = generate_frequencies_eigenvectors(
             supercell_matrix=supercell_matrix,
             displacement=self.displacement,

src/atomate2/forcefields/schemas.py

@@ -96,6 +96,10 @@ class ForceFieldTaskDocument(StructureMetadata):
         description="version of the interatomic potential used for relaxation.",
     )
 
+    dir_name: str | None = Field(
+        None, description="Directory where the force field calculations are performed."
+    )
+
     @classmethod
     def from_ase_compatible_result(
         cls,

src/atomate2/vasp/jobs/phonons.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
@@ -16,8 +17,8 @@
     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:
@@ -191,7 +192,10 @@ def generate_phonon_displacements(
     return displacements
 
 
-@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,
@@ -301,19 +305,16 @@ def run_phonon_displacements(
             "supercell_matrix": supercell_matrix,
             "displaced_structure": displacement,
         }
-        try:
+        with contextlib.suppress(Exception):
             phonon_job.update_maker_kwargs(
-                {"_set": {"write_additional_data->phonon_info:json": info}}, dict_mod=True
+                {"_set": {"write_additional_data->phonon_info:json": info}},
+                dict_mod=True,
             )
-        except:
-            pass
+
         phonon_jobs.append(phonon_job)
         outputs["displacement_number"].append(i)
         outputs["uuids"].append(phonon_job.output.uuid)
-        try:
-            outputs["dirs"].append(phonon_job.output.dir_name)
-        except:
-            outputs["dirs"]=None
+        outputs["dirs"].append(phonon_job.output.dir_name)
         outputs["forces"].append(phonon_job.output.output.forces)
         outputs["displaced_structures"].append(displacement)
 

tests/vasp/flows/test_phonon.py

@@ -4,14 +4,14 @@
 from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
 from pymatgen.phonon.dos import PhononDos
 
-from atomate2.vasp.flows.phonons import PhononMaker
-from atomate2.vasp.schemas.phonons import (
+from atomate2.common.schemas.phonons import (
     PhononBSDOSDoc,
     PhononComputationalSettings,
     PhononJobDirs,
     PhononUUIDs,
     ThermalDisplacementData,
 )
+from atomate2.vasp.flows.phonons import PhononMaker
 
 
 def test_phonon_wf_only_displacements3(mock_vasp, clean_dir):