Equation of State workflow for FHI-aims (#889)

* Added EOS FHI-aims workflow

* Added tmp_path to phonon tests

* Added test data for EOS workflow

* Tests changed for older Pymatgen version

* simplify from_parameters

* Change the test reference value

---------

Co-authored-by: Janosh Riebesell <[email protected]>

src/atomate2/aims/flows/eos.py

@@ -0,0 +1,82 @@
+"""Equation of state workflow for FHI-aims. Based on the common EOS workflow."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from atomate2.aims.flows.core import DoubleRelaxMaker
+from atomate2.aims.jobs.core import RelaxMaker
+from atomate2.common.flows.eos import CommonEosMaker
+
+if TYPE_CHECKING:
+    from jobflow import Maker
+
+
+@dataclass
+class AimsEosMaker(CommonEosMaker):
+    """
+    Generate equation of state data (based on common EOS maker).
+
+    First relaxes a structure using initial_relax_maker, then perform a series of
+    deformations on the relaxed structure, and evaluate single-point energies with
+    static_maker.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    initial_relax_maker : .Maker | None
+        Maker to relax the input structure, defaults to double relaxation.
+    eos_relax_maker : .Maker
+        Maker to relax deformed structures for the EOS fit.
+    static_maker : .Maker | None
+        Maker to generate statics after each relaxation, defaults to None.
+    strain : tuple[float]
+        Percentage linear strain to apply as a deformation, default = -5% to 5%.
+    number_of_frames : int
+        Number of strain calculations to do for EOS fit, default = 6.
+    postprocessor : .atomate2.common.jobs.EOSPostProcessor
+        Optional postprocessing step, defaults to
+        `atomate2.common.jobs.PostProcessEosEnergy`.
+    _store_transformation_information : .bool = False
+        Whether to store the information about transformations. Unfortunately
+        needed at present to handle issues with emmet and pydantic validation
+    """
+
+    name: str = "aims eos"
+    initial_relax_maker: Maker | None = field(
+        default_factory=lambda: DoubleRelaxMaker.from_parameters({})
+    )
+    eos_relax_maker: Maker | None = field(
+        default_factory=lambda: RelaxMaker.fixed_cell_relaxation(
+            user_params={"species_dir": "tight"}
+        )
+    )
+
+    @classmethod
+    def from_parameters(cls, parameters: dict[str, Any], **kwargs) -> AimsEosMaker:
+        """Creation of AimsEosMaker from parameters.
+
+        Parameters
+        ----------
+        parameters : dict
+            Dictionary of common parameters for both makers. The one exception is
+            `species_dir` which can be either a string or a dict with keys [`initial`,
+            `eos`]. If a string is given, it will be interpreted as the `species_dir`
+            for the `eos` Maker; the initial double relaxation will be done then with
+            the default `light` and `tight` species' defaults.
+        kwargs
+            Keyword arguments passed to `CommonEosMaker`.
+        """
+        species_dir = parameters.setdefault("species_dir", "tight")
+        initial_params = parameters.copy()
+        eos_params = parameters.copy()
+        if isinstance(species_dir, dict):
+            initial_params["species_dir"] = species_dir.get("initial")
+            eos_params["species_dir"] = species_dir.get("eos", "tight")
+        return cls(
+            initial_relax_maker=DoubleRelaxMaker.from_parameters(initial_params),
+            eos_relax_maker=RelaxMaker.fixed_cell_relaxation(user_params=eos_params),
+            **kwargs,
+        )

tests/aims/test_flows/test_eos.py

@@ -0,0 +1,104 @@
+"""Test FHI-aims Equation of State workflow"""
+
+import os
+
+import pytest
+from jobflow import run_locally
+from pymatgen.core import Structure
+
+from atomate2.aims.flows.eos import AimsEosMaker
+from atomate2.aims.jobs.core import RelaxMaker
+
+cwd = os.getcwd()
+
+# mapping from job name to directory containing test files
+ref_paths = {
+    "Relaxation calculation 1": "double-relax-si/relax-1",
+    "Relaxation calculation 2": "double-relax-si/relax-2",
+    "Relaxation calculation (fixed cell) deformation 0": "eos-si/0",
+    "Relaxation calculation (fixed cell) deformation 1": "eos-si/1",
+    "Relaxation calculation (fixed cell) deformation 2": "eos-si/2",
+    "Relaxation calculation (fixed cell) deformation 3": "eos-si/3",
+}
+
+
+def test_eos(mock_aims, tmp_path, species_dir):
+    """A test for the equation of state flow"""
+
+    # a relaxed structure for the test
+    a = 2.80791457
+    si = Structure(
+        lattice=[[0.0, a, a], [a, 0.0, a], [a, a, 0.0]],
+        species=["Si", "Si"],
+        coords=[[0, 0, 0], [0.25, 0.25, 0.25]],
+    )
+
+    # settings passed to fake_run_aims
+    fake_run_kwargs = {}
+
+    # automatically use fake AIMS
+    mock_aims(ref_paths, fake_run_kwargs)
+
+    # generate flow
+    eos_relax_maker = RelaxMaker.fixed_cell_relaxation(
+        user_params={
+            "species_dir": (species_dir / "light").as_posix(),
+            # "species_dir": "light",
+            "k_grid": [2, 2, 2],
+        }
+    )
+
+    flow = AimsEosMaker(
+        initial_relax_maker=None, eos_relax_maker=eos_relax_maker, number_of_frames=4
+    ).make(si)
+
+    # Run the flow or job and ensure that it finished running successfully
+    os.chdir(tmp_path)
+    responses = run_locally(flow, create_folders=True, ensure_success=True)
+    os.chdir(cwd)
+
+    output = responses[flow.jobs[-1].uuid][1].output
+    assert "EOS" in output["relax"]
+    # there is no initial calculation; fit using 4 points
+    assert len(output["relax"]["energy"]) == 4
+    assert output["relax"]["EOS"]["birch_murnaghan"]["b0"] == pytest.approx(
+        0.4897486348366812
+    )
+
+
+def test_eos_from_parameters(mock_aims, tmp_path, si, species_dir):
+    """A test for the equation of state flow, created from the common parameters"""
+
+    # settings passed to fake_run_aims
+    fake_run_kwargs = {}
+
+    # automatically use fake AIMS
+    mock_aims(ref_paths, fake_run_kwargs)
+
+    # generate flow
+    flow = AimsEosMaker.from_parameters(
+        parameters={
+            # TODO: to be changed after pymatgen PR is merged
+            "species_dir": {
+                "initial": species_dir,
+                "eos": (species_dir / "light").as_posix(),
+            },
+            # "species_dir": "light",
+            "k_grid": [2, 2, 2],
+        },
+        number_of_frames=4,
+    ).make(si)
+
+    # Run the flow or job and ensure that it finished running successfully
+    os.chdir(tmp_path)
+    responses = run_locally(flow, create_folders=True, ensure_success=True)
+    os.chdir(cwd)
+
+    output = responses[flow.jobs[-1].uuid][1].output
+    assert "EOS" in output["relax"]
+    # there is an initial calculation; fit using 5 points
+    assert len(output["relax"]["energy"]) == 5
+    # the initial calculation also participates in the fit here
+    assert output["relax"]["EOS"]["birch_murnaghan"]["b0"] == pytest.approx(
+        0.5189578108402951
+    )

tests/aims/test_flows/test_phonon_workflow.py

@@ -153,9 +153,9 @@ def test_phonon_socket_flow(si, tmp_path, mock_aims, species_dir):
     )
 
     # run the flow or job and ensure that it finished running successfully
-    # os.chdir(tmp_path)
+    os.chdir(tmp_path)
     responses = run_locally(flow, create_folders=True, ensure_success=True)
-    # os.chdir(cwd)
+    os.chdir(cwd)
 
     # validation the outputs of the job
     output = responses[flow.job_uuids[-1]][1].output