Aims phonon tutorials

src/atomate2/utils/testing/aims.py

@@ -0,0 +1,154 @@
+"""Utilities for testing FHI-aims calculations."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Final, Literal
+
+from jobflow import CURRENT_JOB
+from monty.os.path import zpath as monty_zpath
+from pymatgen.io.aims.sets.base import AimsInputGenerator
+
+import atomate2.aims.jobs.base
+import atomate2.aims.run
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator, Sequence
+
+    from pymatgen.io.aims.sets import AimsInputSet
+    from pytest import MonkeyPatch
+
+logger = logging.getLogger("atomate2")
+
+
+_VFILES: Final = ("control.in",)
+_REF_PATHS: dict[str, str | Path] = {}
+_FAKE_RUN_AIMS_KWARGS: dict[str, dict] = {}
+
+
+def zpath(path: str | Path) -> Path:
+    """Return the path of a zip file.
+
+    Returns an existing (zipped or unzipped) file path given the unzipped
+    version. If no path exists, returns the unmodified path.
+    """
+    return Path(monty_zpath(str(path)))
+
+
+def monkeypatch_aims(
+    monkeypatch: MonkeyPatch, ref_path: Path
+) -> Generator[Callable[[Any, Any], Any], None, None]:
+    """Allow one to mock (fake) running FHI-aims.
+
+    To use the fixture successfully, the following steps must be followed:
+    1. "mock_aims" should be included as an argument to any test that would like to use
+       its functionally.
+    2. For each job in your workflow, you should prepare a reference directory
+       containing two folders "inputs" (containing the reference input files expected
+       to be produced by write_aims_input_set) and "outputs" (containing the expected
+       output files to be produced by run_aims). These files should reside in a
+       subdirectory of "tests/test_data/aims".
+    3. Create a dictionary mapping each job name to its reference directory. Note that
+       you should supply the reference directory relative to the "tests/test_data/aims"
+       folder. For example, if your calculation has one job named "static" and the
+       reference files are present in "tests/test_data/aims/Si_static", the dictionary
+       would look like: ``{"static": "Si_static"}``.
+    4. Optional (does not work yet): create a dictionary mapping each job name to
+       custom keyword arguments that will be supplied to fake_run_aims.
+       This way you can configure which control.in settings are expected for each job.
+       For example, if your calculation has one job named "static" and you wish to
+       validate that "xc" is set correctly in the control.in, your dictionary would
+       look like
+       ``{"static": {"input_settings": {"relativistic": "atomic_zora scalar"}}``.
+    5. Inside the test function, call `mock_aims(ref_paths, fake_aims_kwargs)`, where
+       ref_paths is the dictionary created in step 3 and fake_aims_kwargs is the
+       dictionary created in step 4.
+    6. Run your aims job after calling `mock_aims`.
+
+    For examples, see the tests in tests/aims/jobs/core.py.
+    """
+
+    def mock_run_aims(*args, **kwargs) -> None:  # noqa: ARG001
+        name = CURRENT_JOB.job.name
+        try:
+            ref_dir = ref_path / _REF_PATHS[name]
+        except KeyError:
+            raise ValueError(
+                f"no reference directory found for job {name!r}; "
+                f"reference paths received={_REF_PATHS}"
+            ) from None
+        fake_run_aims(ref_dir, **_FAKE_RUN_AIMS_KWARGS.get(name, {}))
+
+    get_input_set_orig = AimsInputGenerator.get_input_set
+
+    def mock_get_input_set(self: AimsInputGenerator, *args, **kwargs) -> AimsInputSet:
+        return get_input_set_orig(self, *args, **kwargs)
+
+    monkeypatch.setattr(atomate2.aims.run, "run_aims", mock_run_aims)
+    monkeypatch.setattr(atomate2.aims.jobs.base, "run_aims", mock_run_aims)
+    monkeypatch.setattr(AimsInputGenerator, "get_input_set", mock_get_input_set)
+
+    def _run(ref_paths: dict, fake_run_aims_kwargs: dict | None = None) -> None:
+        _REF_PATHS.update(ref_paths)
+        _FAKE_RUN_AIMS_KWARGS.update(fake_run_aims_kwargs or {})
+
+    yield _run
+
+    monkeypatch.undo()
+    _REF_PATHS.clear()
+    _FAKE_RUN_AIMS_KWARGS.clear()
+
+
+def fake_run_aims(
+    ref_path: str | Path,
+    input_settings: Sequence[str] | None = None,  # noqa: ARG001
+    check_inputs: Sequence[Literal["control.in"]] = _VFILES,  # noqa: ARG001
+    clear_inputs: bool = False,
+) -> None:
+    """
+    Emulate running aims and validate aims input files.
+
+    Parameters
+    ----------
+    ref_path
+        Path to reference directory with aims input files in the folder named 'inputs'
+        and output files in the folder named 'outputs'.
+    input_settings
+        A list of input settings to check.
+    check_inputs
+        A list of aims input files to check. Supported options are "aims.inp"
+    clear_inputs
+        Whether to clear input files before copying in the reference aims outputs.
+    """
+    logger.info("Running fake aims.")
+
+    ref_path = Path(ref_path)
+
+    logger.info("Verified inputs successfully")
+
+    if clear_inputs:
+        clear_aims_inputs()
+
+    copy_aims_outputs(ref_path)
+
+    # pretend to run aims by copying pre-generated outputs from reference dir
+    logger.info("Generated fake aims outputs")
+
+
+def clear_aims_inputs() -> None:
+    """Clean up FHI-aims input files."""
+    for aims_file in ("control.in", "geometry.in", "parameters.json"):
+        if Path(aims_file).exists():
+            Path(aims_file).unlink()
+    logger.info("Cleared aims inputs")
+
+
+def copy_aims_outputs(ref_path: str | Path) -> None:
+    """Copy FHI-aims output files from the reference directory."""
+    import shutil
+
+    output_path = Path(ref_path) / "outputs"
+    for output_file in output_path.iterdir():
+        if output_file.is_file():
+            shutil.copy(output_file, ".")

tests/aims/test_flows/test_anharmonic_quantification.py

@@ -1,6 +1,9 @@
+from pathlib import Path
+
 import numpy as np
 import pytest
 from jobflow import run_locally
+from pymatgen.core import SETTINGS, Structure
 from pymatgen.io.aims.sets.core import SocketIOSetGenerator, StaticSetGenerator
 
 from atomate2.aims.flows.anharmonicity import AnharmonicityMaker
@@ -11,8 +14,11 @@
     PhononDisplacementMakerSocket,
 )
 
+si_structure_file = Path(__file__).parents[2] / "test_data/structures/Si_diamond.cif"
+
 
-def test_anharmonic_quantification_oneshot(si, clean_dir, mock_aims, species_dir):
+def test_anharmonic_quantification_oneshot(clean_dir, mock_aims, species_dir):
+    si = Structure.from_file(si_structure_file)
     # mapping from job name to directory containing test files
     ref_paths = {
         "Relaxation calculation": "phonon-relax-si",
@@ -21,94 +27,63 @@ def test_anharmonic_quantification_oneshot(si, clean_dir, mock_aims, species_dir
         "phonon static aims anharmonicity quant. 1/1": "anharm-os-si",
     }
 
+    SETTINGS["AIMS_SPECIES_DIR"] = species_dir / "tight"
     # settings passed to fake_run_aims; adjust these to check for certain input settings
     fake_run_aims_kwargs = {}
 
     # automatically use fake FHI-aims
     mock_aims(ref_paths, fake_run_aims_kwargs)
 
-    parameters = {
-        "species_dir": (species_dir / "light").as_posix(),
-        "rlsy_symmetry": "all",
-        "sc_accuracy_rho": 1e-06,
-        "sc_accuracy_forces": 0.0001,
-        "relativistic": "atomic_zora scalar",
-    }
-
-    parameters_phonon_disp = dict(compute_forces=True, **parameters)
-    parameters_phonon_disp["rlsy_symmetry"] = None
-
     phonon_maker = PhononMaker(
-        bulk_relax_maker=RelaxMaker.full_relaxation(
-            user_params=parameters, user_kpoints_settings={"density": 5.0}
-        ),
-        static_energy_maker=StaticMaker(
-            input_set_generator=StaticSetGenerator(
-                user_params=parameters, user_kpoints_settings={"density": 5.0}
-            )
-        ),
+        min_length=3.0,
+        generate_frequencies_eigenvectors_kwargs={"tstep": 100},
+        create_thermal_displacements=True,
+        store_force_constants=True,
+        born_maker=None,
         use_symmetrized_structure="primitive",
-        phonon_displacement_maker=PhononDisplacementMaker(
-            input_set_generator=StaticSetGenerator(
-                user_params=parameters_phonon_disp,
-                user_kpoints_settings={"density": 5.0},
-            )
-        ),
     )
 
     maker = AnharmonicityMaker(
         phonon_maker=phonon_maker,
     )
     maker.name = "anharmonicity"
-    flow = maker.make(si, supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3))
+    flow = maker.make(
+        si,
+        supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3),
+        one_shot_approx=True,
+        seed=1234,
+    )
 
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(flow, create_folders=True, ensure_success=True)
     dct = responses[flow.job_uuids[-1]][1].output.sigma_dict
-    assert np.round(dct["one-shot"], 3) == 0.104
+    assert np.round(dct["one-shot"], 3) == 0.120
+
 
+def test_anharmonic_quantification_full(clean_dir, mock_aims, species_dir):
+    si = Structure.from_file(si_structure_file)
 
-def test_anharmonic_quantification_full(si, clean_dir, mock_aims, species_dir):
     ref_paths = {
-        "Relaxation calculation": "phonon-relax-si-full",
-        "phonon static aims 1/1": "phonon-disp-si-full",
-        "SCF Calculation": "phonon-energy-si-full",
+        "Relaxation calculation": "phonon-relax-si",
+        "phonon static aims 1/1": "phonon-disp-si",
+        "SCF Calculation": "phonon-energy-si",
         "phonon static aims anharmonicity quant. 1/1": "anharm-si-full",
     }
 
+    SETTINGS["AIMS_SPECIES_DIR"] = species_dir / "tight"
     # settings passed to fake_run_aims; adjust these to check for certain input settings
     fake_run_aims_kwargs = {}
 
     # automatically use fake FHI-aims
     mock_aims(ref_paths, fake_run_aims_kwargs)
 
-    parameters = {
-        "species_dir": (species_dir / "light").as_posix(),
-        "rlsy_symmetry": "all",
-        "sc_accuracy_rho": 1e-06,
-        "sc_accuracy_forces": 0.0001,
-        "relativistic": "atomic_zora scalar",
-    }
-
-    parameters_phonon_disp = dict(compute_forces=True, **parameters)
-    parameters_phonon_disp["rlsy_symmetry"] = None
-
     phonon_maker = PhononMaker(
-        bulk_relax_maker=RelaxMaker.full_relaxation(
-            user_params=parameters, user_kpoints_settings={"density": 5.0}
-        ),
-        static_energy_maker=StaticMaker(
-            input_set_generator=StaticSetGenerator(
-                user_params=parameters, user_kpoints_settings={"density": 5.0}
-            )
-        ),
+        min_length=3.0,
+        generate_frequencies_eigenvectors_kwargs={"tstep": 100},
+        create_thermal_displacements=True,
+        store_force_constants=True,
+        born_maker=None,
         use_symmetrized_structure="primitive",
-        phonon_displacement_maker=PhononDisplacementMaker(
-            input_set_generator=StaticSetGenerator(
-                user_params=parameters_phonon_disp,
-                user_kpoints_settings={"density": 5.0},
-            )
-        ),
     )
 
     maker = AnharmonicityMaker(
@@ -125,7 +100,8 @@ def test_anharmonic_quantification_full(si, clean_dir, mock_aims, species_dir):
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(flow, create_folders=True, ensure_success=True)
     dct = responses[flow.job_uuids[-1]][1].output.sigma_dict
-    assert pytest.approx(dct["full"], 0.001) == 0.12012
+
+    assert pytest.approx(dct["full"], 0.001) == 0.144264
 
 
 def test_mode_resolved_anharmonic_quantification(si, clean_dir, mock_aims, species_dir):

tests/aims/test_flows/test_phonon_workflow.py

@@ -1,18 +1,21 @@
 """Test various makers"""
 
 import json
+from pathlib import Path
 
 import pytest
 
+si_structure_file = Path(__file__).parents[2] / "test_data/structures/Si_diamond.cif"
 
-def test_phonon_flow(si, clean_dir, mock_aims, species_dir):
+
+def test_phonon_flow(clean_dir, mock_aims, species_dir):
     import numpy as np
     from jobflow import run_locally
-    from pymatgen.io.aims.sets.core import StaticSetGenerator
+    from pymatgen.core import Structure
 
     from atomate2.aims.flows.phonons import PhononMaker
-    from atomate2.aims.jobs.core import RelaxMaker, StaticMaker
-    from atomate2.aims.jobs.phonons import PhononDisplacementMaker
+
+    si = Structure.from_file(si_structure_file)
 
     # mapping from job name to directory containing test files
     ref_paths = {
@@ -26,26 +29,15 @@ def test_phonon_flow(si, clean_dir, mock_aims, species_dir):
 
     # automatically use fake FHI-aims
     mock_aims(ref_paths, fake_run_aims_kwargs)
-
-    parameters = {
-        "k_grid": [2, 2, 2],
-        "species_dir": (species_dir / "light").as_posix(),
-    }
     # generate job
 
-    parameters_phonon_disp = dict(compute_forces=True, **parameters)
     maker = PhononMaker(
-        bulk_relax_maker=RelaxMaker.full_relaxation(user_params=parameters),
-        static_energy_maker=StaticMaker(
-            input_set_generator=StaticSetGenerator(user_params=parameters)
-        ),
+        min_length=3.0,
+        generate_frequencies_eigenvectors_kwargs={"tstep": 100},
+        create_thermal_displacements=True,
+        store_force_constants=True,
+        born_maker=None,
         use_symmetrized_structure="primitive",
-        phonon_displacement_maker=PhononDisplacementMaker(
-            input_set_generator=StaticSetGenerator(
-                user_params=parameters_phonon_disp,
-                user_kpoints_settings={"density": 5.0, "even": True},
-            )
-        ),
     )
     maker.name = "phonons"
     flow = maker.make(si, supercell_matrix=np.ones((3, 3)) - 2 * np.eye(3))
@@ -84,11 +76,11 @@ def test_phonon_flow(si, clean_dir, mock_aims, species_dir):
     assert output.born is None
     assert not output.has_imaginary_modes
 
-    assert output.temperatures == list(range(0, 500, 10))
+    assert output.temperatures == list(range(0, 500, 100))
     assert output.heat_capacities[0] == 0.0
-    assert np.round(output.heat_capacities[-1], 2) == 23.06
+    assert np.round(output.heat_capacities[-1], 2) == 21.95
     assert output.phonopy_settings.schema_json() == json.dumps(phonopy_settings_schema)
-    assert np.round(output.phonon_bandstructure.bands[-1, 0], 2) == 14.41
+    assert np.round(output.phonon_bandstructure.bands[-1, 0], 2) == 15.1
 
 
 @pytest.mark.skip(reason="Currently not mocked and needs FHI-aims binary")

tests/test_data/aims/anharm-os-si/inputs/parameters.json

@@ -1 +1 @@
-{"xc": "pbe", "relativistic": "atomic_zora scalar", "compute_forces": true, "k_grid": [2, 2, 2], "species_dir": "/home/purcellt/git/atomate2/tests/aims/species_dir/light", "rlsy_symmetry": null, "sc_accuracy_rho": 1e-06, "sc_accuracy_forces": 0.0001}
+{"xc": "pbe", "relativistic": "atomic_zora scalar", "compute_forces": true, "k_grid": [6, 6, 6]}