bug fix for linear strain passing in qha (#1061)

* bug fix for linear strain qha maker

* refactor to only rely on one method to generate supercells

* fix tests

* add test

* add test 2

* fix type annotation

* fix type annotation

* fix docs

* fix docs

* add test for supercell

* fix supercell tests

* fix supercell tests

Author: JaGeo

src/atomate2/common/flows/phonons.py

@@ -69,6 +69,8 @@ class BasePhononMaker(Maker, ABC):
         displacement distance for phonons
     min_length: float
         min length of the supercell that will be built
+    max_length: float
+        max length of the supercell that will be built
     prefer_90_degrees: bool
         if set to True, supercell algorithm will first try to find a supercell
         with 3 90 degree angles

src/atomate2/common/flows/qha.py

@@ -10,17 +10,21 @@
 from jobflow import Flow, Maker
 
 from atomate2.common.flows.eos import CommonEosMaker
-from atomate2.common.jobs.qha import analyze_free_energy, get_phonon_jobs
+from atomate2.common.jobs.qha import (
+    analyze_free_energy,
+    get_phonon_jobs,
+    get_supercell_size,
+)
 
 if TYPE_CHECKING:
     from pathlib import Path
 
+    from emmet.core.math import Matrix3D
     from pymatgen.core import Structure
 
     from atomate2.common.flows.phonons import BasePhononMaker
     from atomate2.forcefields.jobs import ForceFieldRelaxMaker
     from atomate2.vasp.jobs.core import BaseVaspMaker
-
 supported_eos = frozenset(("vinet", "birch_murnaghan", "murnaghan"))
 
 
@@ -71,6 +75,16 @@ class CommonQhaMaker(Maker, ABC):
         will be ignored
     eos_type: str
         Equation of State type used for the fitting. Defaults to vinet.
+    min_length: float
+        min length of the supercell that will be built
+    max_length: float
+        max length of the supercell that will be built
+    prefer_90_degrees: bool
+        if set to True, supercell algorithm will first try to find a supercell
+        with 3 90 degree angles
+    get_supercell_size_kwargs: dict
+        kwargs that will be passed to get_supercell_size to determine supercell size
+
     """
 
     name: str = "QHA Maker"
@@ -85,16 +99,27 @@ class CommonQhaMaker(Maker, ABC):
     skip_analysis: bool = False
     eos_type: Literal["vinet", "birch_murnaghan", "murnaghan"] = "vinet"
     analyze_free_energy_kwargs: dict = field(default_factory=dict)
-    # TODO: implement advanced handling of
-    #  imaginary modes in phonon runs (i.e., fitting procedures)
-
-    def make(self, structure: Structure, prev_dir: str | Path = None) -> Flow:
+    min_length: float | None = 20.0
+    max_length: float | None = None
+    prefer_90_degrees: bool = True
+    allow_orthorhombic: bool = False
+    get_supercell_size_kwargs: dict = field(default_factory=dict)
+
+    def make(
+        self,
+        structure: Structure,
+        supercell_matrix: Matrix3D | None = None,
+        prev_dir: str | Path = None,
+    ) -> Flow:
         """Run an EOS flow.
 
         Parameters
         ----------
         structure : Structure
             A pymatgen structure object.
+        supercell_matrix: list
+            Instead of min_length, also a supercell_matrix can be given, e.g.
+            [[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]
         prev_dir : str or Path or None
             A previous calculation directory to copy output files from.
 
@@ -116,14 +141,32 @@ def make(self, structure: Structure, prev_dir: str | Path = None) -> Flow:
             eos_relax_maker=self.eos_relax_maker,
             static_maker=None,
             postprocessor=None,
+            linear_strain=self.linear_strain,
             number_of_frames=self.number_of_frames,
         )
 
         eos_job = self.eos.make(structure)
         qha_jobs.append(eos_job)
 
+        # implement a supercell job to get matrix for just the equillibrium structure
+        if supercell_matrix is None:
+            supercell = get_supercell_size(
+                eos_output=eos_job.output,
+                min_length=self.min_length,
+                max_length=self.max_length,
+                prefer_90_degrees=self.prefer_90_degrees,
+                allow_orthorhombic=self.allow_orthorhombic,
+                **self.get_supercell_size_kwargs,
+            )
+            qha_jobs.append(supercell)
+            supercell_matrix = supercell.output
+
+        # pass the matrix to the phonon_jobs, allow to set a consistent matrix instead
+
         phonon_jobs = get_phonon_jobs(
-            phonon_maker=self.phonon_maker, eos_output=eos_job.output
+            phonon_maker=self.phonon_maker,
+            eos_output=eos_job.output,
+            supercell_matrix=supercell_matrix,
         )
         qha_jobs.append(phonon_jobs)
         if not self.skip_analysis:

src/atomate2/common/jobs/phonons.py

@@ -14,11 +14,9 @@
 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, get_factor
+from atomate2.common.utils import get_supercell_matrix
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -82,45 +80,15 @@ def get_supercell_size(
     **kwargs:
         Additional parameters that can be set.
     """
-    kwargs.setdefault("force_diagonal", False)
-    common_kwds = dict(
-        min_length=min_length,
+    return get_supercell_matrix(
+        allow_orthorhombic=allow_orthorhombic,
         max_length=max_length,
-        min_atoms=kwargs.get("min_atoms"),
-        max_atoms=kwargs.get("max_atoms"),
-        step_size=kwargs.get("step_size", 0.1),
-        force_diagonal=kwargs["force_diagonal"],
+        min_length=min_length,
+        prefer_90_degrees=prefer_90_degrees,
+        structure=structure,
+        **kwargs,
     )
 
-    if not prefer_90_degrees:
-        transformation = CubicSupercellTransformation(
-            **common_kwds,
-            force_90_degrees=False,
-            allow_orthorhombic=allow_orthorhombic,
-        )
-        transformation.apply_transformation(structure=structure)
-    else:
-        try:
-            common_kwds.update({"max_atoms": kwargs.get("max_atoms", 1200)})
-            transformation = CubicSupercellTransformation(
-                **common_kwds,
-                force_90_degrees=True,
-                angle_tolerance=kwargs.get("angle_tolerance", 1e-2),
-                allow_orthorhombic=allow_orthorhombic,
-            )
-            transformation.apply_transformation(structure=structure)
-
-        except AttributeError:
-            transformation = CubicSupercellTransformation(
-                **common_kwds,
-                force_90_degrees=False,
-                allow_orthorhombic=allow_orthorhombic,
-            )
-            transformation.apply_transformation(structure=structure)
-
-    # matrix from pymatgen has to be transposed
-    return transformation.transformation_matrix.transpose().tolist()
-
 
 @job(data=[Structure])
 def generate_phonon_displacements(

src/atomate2/common/jobs/qha.py

@@ -9,6 +9,7 @@
 
 from atomate2.common.schemas.phonons import PhononBSDOSDoc
 from atomate2.common.schemas.qha import PhononQHADoc
+from atomate2.common.utils import get_supercell_matrix
 
 if TYPE_CHECKING:
     from pymatgen.core.structure import Structure
@@ -18,10 +19,49 @@
 logger = logging.getLogger(__name__)
 
 
+@job
+def get_supercell_size(
+    eos_output: dict,
+    min_length: float,
+    max_length: float,
+    prefer_90_degrees: bool,
+    allow_orthorhombic: bool = False,
+    **kwargs,
+) -> list[list[float]]:
+    """
+    Job to get the supercell size from an eos output.
+
+    Parameters
+    ----------
+    eos_output: dict
+        output from eos state job
+    min_length: float
+        minimum length of cell in Angstrom
+    max_length: float
+        maximum length of cell in Angstrom
+    prefer_90_degrees: bool
+        if True, the algorithm will try to find a cell with 90 degree angles first
+    allow_orthorhombic: bool
+        if True, orthorhombic supercells are allowed
+    **kwargs:
+        Additional parameters that can be set.
+    """
+    return get_supercell_matrix(
+        eos_output["relax"]["structure"][0],
+        min_length,
+        max_length,
+        prefer_90_degrees,
+        allow_orthorhombic,
+        **kwargs,
+    )
+
+
 @job(
     data=[PhononBSDOSDoc],
 )
-def get_phonon_jobs(phonon_maker: BasePhononMaker, eos_output: dict) -> Flow:
+def get_phonon_jobs(
+    phonon_maker: BasePhononMaker, eos_output: dict, supercell_matrix: list[list[float]]
+) -> Flow:
     """
     Start all relevant phonon jobs.
 
@@ -31,17 +71,20 @@ def get_phonon_jobs(phonon_maker: BasePhononMaker, eos_output: dict) -> Flow:
         Maker to start harmonic phonon runs.
     eos_output: dict
         Output from EOSMaker
-
+    supercell_matrix:
+        Supercell matrix to be passed into the phonon runs.
     """
     phonon_jobs = []
     outputs = []
     for istructure, structure in enumerate(eos_output["relax"]["structure"]):
         if eos_output["relax"]["dir_name"][istructure] is not None:
             phonon_job = phonon_maker.make(
-                structure, prev_dir=eos_output["relax"]["dir_name"][istructure]
+                structure,
+                prev_dir=eos_output["relax"]["dir_name"][istructure],
+                supercell_matrix=supercell_matrix,
             )
         else:
-            phonon_job = phonon_maker.make(structure)
+            phonon_job = phonon_maker.make(structure, supercell_matrix=supercell_matrix)
         phonon_job.append_name(f" eos deformation {istructure + 1}")
         phonon_jobs.append(phonon_job)
         outputs.append(phonon_job.output)
@@ -92,6 +135,7 @@ def analyze_free_energy(
         output.volume_per_formula_unit * output.formula_units
         for output in phonon_outputs
     ]
+    supercell_matrix: list[list[float]] = phonon_outputs[0].supercell_matrix
 
     for itemp, temp in enumerate(phonon_outputs[0].temperatures):
         temperatures.append(float(temp))
@@ -129,5 +173,6 @@ def analyze_free_energy(
         pressure=pressure,
         formula_units=next(iter(set(formula_units))),
         eos_type=eos_type,
+        supercell_matrix=supercell_matrix,
         **kwargs,
     )

src/atomate2/common/schemas/qha.py

@@ -4,6 +4,7 @@
 from typing import Optional, Union
 
 import numpy as np
+from emmet.core.math import Matrix3D
 from emmet.core.structure import StructureMetadata
 from phonopy.api_qha import PhonopyQHA
 from pydantic import Field
@@ -88,6 +89,8 @@ class PhononQHADoc(StructureMetadata, extra="allow"):  # type: ignore[call-arg]
     )
     formula_units: Optional[int] = Field(None, description="Formula units")
 
+    supercell_matrix: Optional[Matrix3D] = Field(None, description="Supercell matrix")
+
     @classmethod
     def from_phonon_runs(
         cls,
@@ -98,6 +101,7 @@ def from_phonon_runs(
         free_energies: list[list[float]],
         heat_capacities: list[list[float]],
         entropies: list[list[float]],
+        supercell_matrix: list[list[float]],
         t_max: float = None,
         pressure: float = None,
         formula_units: Union[int, None] = None,
@@ -115,6 +119,7 @@ def from_phonon_runs(
         free_energies: list of list of floats
         heat_capacities: list of list of floats
         entropies: list of list of floats
+        supercell_matrix: list of list of floats
         t_max: float
         pressure: float
         eos_type: string
@@ -232,4 +237,5 @@ def from_phonon_runs(
             heat_capacities=heat_capacities,
             entropies=entropies,
             formula_units=formula_units,
+            supercell_matrix=supercell_matrix,
         )

src/atomate2/common/utils.py

@@ -7,10 +7,79 @@
 from typing import TYPE_CHECKING, Any
 
 from monty.serialization import loadfn
+from pymatgen.transformations.advanced_transformations import (
+    CubicSupercellTransformation,
+)
 
 if TYPE_CHECKING:
     from pathlib import Path
 
+    from pymatgen.core.structure import Structure
+
+
+def get_supercell_matrix(
+    structure: Structure,
+    min_length: float,
+    max_length: float,
+    prefer_90_degrees: bool,
+    allow_orthorhombic: bool = False,
+    **kwargs,
+) -> list[list[float]]:
+    """
+    Determine supercell size with given min_length and max_length.
+
+    Parameters
+    ----------
+    structure: Structure Object
+        Input structure that will be used to determine supercell
+    min_length: float
+        minimum length of cell in Angstrom
+    max_length: float
+        maximum length of cell in Angstrom
+    prefer_90_degrees: bool
+        if True, the algorithm will try to find a cell with 90 degree angles first
+    allow_orthorhombic: bool
+        if True, orthorhombic supercells are allowed
+    **kwargs:
+        Additional parameters that can be set.
+    """
+    kwargs.setdefault("force_diagonal", False)
+    common_kwds = dict(
+        min_length=min_length,
+        max_length=max_length,
+        min_atoms=kwargs.get("min_atoms"),
+        max_atoms=kwargs.get("max_atoms"),
+        step_size=kwargs.get("step_size", 0.1),
+        force_diagonal=kwargs["force_diagonal"],
+    )
+    if not prefer_90_degrees:
+        transformation = CubicSupercellTransformation(
+            **common_kwds,
+            force_90_degrees=False,
+            allow_orthorhombic=allow_orthorhombic,
+        )
+        transformation.apply_transformation(structure=structure)
+    else:
+        try:
+            common_kwds.update({"max_atoms": kwargs.get("max_atoms", 1200)})
+            transformation = CubicSupercellTransformation(
+                **common_kwds,
+                force_90_degrees=True,
+                angle_tolerance=kwargs.get("angle_tolerance", 1e-2),
+                allow_orthorhombic=allow_orthorhombic,
+            )
+            transformation.apply_transformation(structure=structure)
+
+        except AttributeError:
+            transformation = CubicSupercellTransformation(
+                **common_kwds,
+                force_90_degrees=False,
+                allow_orthorhombic=allow_orthorhombic,
+            )
+            transformation.apply_transformation(structure=structure)
+    # matrix from pymatgen has to be transposed
+    return transformation.transformation_matrix.transpose().tolist()
+
 
 def get_transformations(
     transformations: tuple[str, ...], params: tuple[dict, ...] | None