Add an ase_relax_job for VASP (#1888)

Closes #1887.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/quacc/recipes/onetep/core.py

@@ -69,6 +69,7 @@ def static_job(
 @job
 def ase_relax_job(
     atoms: Atoms,
+    relax_cell: bool = False,
     copy_files: SourceDirectory | dict[SourceDirectory, Filenames] | None = None,
     opt_params: dict[str, Any] | None = None,
     **calc_kwargs,
@@ -102,7 +103,7 @@ def ase_relax_job(
         {"keywords": {"write_forces": True, "forces_output_detail": "verbose"}},
     )
 
-    opt_defaults = {"optimizer": LBFGS}
+    opt_defaults = {"optimizer": LBFGS, "relax_cell": relax_cell}
 
     return base_opt_fn(
         atoms,

src/quacc/recipes/vasp/_base.py

@@ -4,8 +4,10 @@
 
 from typing import TYPE_CHECKING
 
+from quacc.atoms.core import get_final_atoms_from_dyn
 from quacc.calculators.vasp import Vasp
-from quacc.runners.ase import run_calc
+from quacc.runners.ase import run_calc, run_opt
+from quacc.schemas.ase import summarize_opt_run
 from quacc.schemas.vasp import vasp_summarize_run
 from quacc.utils.dicts import recursive_dict_merge
 
@@ -52,7 +54,7 @@ def base_fn(
     Returns
     -------
     VaspSchema
-        Dictionary of results from [quacc.schemas.vasp.vasp_summarize_run][]
+        Dictionary of results
     """
     calc_flags = recursive_dict_merge(calc_defaults, calc_swaps)
 
@@ -64,3 +66,63 @@ def base_fn(
         report_mp_corrections=report_mp_corrections,
         additional_fields=additional_fields,
     )
+
+
+def base_opt_fn(
+    atoms: Atoms,
+    preset: str | None = None,
+    calc_defaults: dict[str, Any] | None = None,
+    calc_swaps: dict[str, Any] | None = None,
+    opt_defaults: dict[str, Any] | None = None,
+    opt_params: dict[str, Any] | None = None,
+    report_mp_corrections: bool = False,
+    additional_fields: dict[str, Any] | None = None,
+    copy_files: SourceDirectory | dict[SourceDirectory, Filenames] | None = None,
+) -> VaspSchema:
+    """
+    Base job function for VASP recipes.
+
+    Parameters
+    ----------
+    atoms
+        Atoms object
+    preset
+        Preset to use from `quacc.calculators.vasp.presets`.
+    calc_defaults
+        Default parameters for the recipe.
+    calc_swaps
+        Dictionary of custom kwargs for the Vasp calculator. Set a value to
+        `None` to remove a pre-existing key entirely. For a list of available
+        keys, refer to [quacc.calculators.vasp.vasp.Vasp][].
+    opt_defaults
+        Default arguments for the ASE optimizer.
+    opt_params
+        Dictionary of custom kwargs for [quacc.runners.ase.run_opt][]
+    report_mp_corrections
+        Whether to report the Materials Project corrections in the results.
+    additional_fields
+        Additional fields to supply to the summarizer.
+    copy_files
+        Files to copy (and decompress) from source to the runtime directory.
+
+    Returns
+    -------
+    VaspASESchema
+        Dictionary of results
+    """
+    calc_flags = recursive_dict_merge(calc_defaults, calc_swaps)
+    opt_flags = recursive_dict_merge(opt_defaults, opt_params)
+
+    atoms.calc = Vasp(atoms, preset=preset, **calc_flags)
+    dyn = run_opt(atoms, copy_files=copy_files, **opt_flags)
+
+    opt_run_summary = summarize_opt_run(dyn, additional_fields=additional_fields)
+
+    final_atoms = get_final_atoms_from_dyn(dyn)
+
+    vasp_summary = vasp_summarize_run(
+        final_atoms,
+        report_mp_corrections=report_mp_corrections,
+        additional_fields=additional_fields,
+    )
+    return recursive_dict_merge(vasp_summary, opt_run_summary)

src/quacc/recipes/vasp/core.py

@@ -10,14 +10,14 @@
 from pymatgen.io.vasp import Vasprun
 
 from quacc import flow, job
-from quacc.recipes.vasp._base import base_fn
+from quacc.recipes.vasp._base import base_fn, base_opt_fn
 
 if TYPE_CHECKING:
     from typing import Any
 
     from ase.atoms import Atoms
 
-    from quacc.schemas._aliases.vasp import DoubleRelaxSchema, VaspSchema
+    from quacc.schemas._aliases.vasp import DoubleRelaxSchema, VaspASESchema, VaspSchema
     from quacc.utils.files import Filenames, SourceDirectory
 
 
@@ -180,6 +180,58 @@ def double_relax_flow(
     return {"relax1": summary1, "relax2": summary2}
 
 
+@job
+def ase_relax_job(
+    atoms: Atoms,
+    preset: str | None = "BulkSet",
+    relax_cell: bool = True,
+    opt_params: dict[str, Any] | None = None,
+    copy_files: SourceDirectory | dict[SourceDirectory, Filenames] | None = None,
+    **calc_kwargs,
+) -> VaspASESchema:
+    """
+    Relax a structure.
+
+    Parameters
+    ----------
+    atoms
+        Atoms object
+    preset
+        Preset to use from `quacc.calculators.vasp.presets`.
+    relax_cell
+        True if a volume relaxation should be performed. False if only the positions
+        should be updated.
+    copy_files
+        Files to copy (and decompress) from source to the runtime directory.
+    **calc_kwargs
+        Custom kwargs for the Vasp calculator. Set a value to
+        `None` to remove a pre-existing key entirely. For a list of available
+        keys, refer to the [quacc.calculators.vasp.vasp.Vasp][] calculator.
+
+    Returns
+    -------
+    VaspASESchema
+        Dictionary of results. See the type-hint for the data structure.
+    """
+
+    calc_defaults = {
+        "lcharg": False,
+        "lwave": False,
+        "nsw": 0,
+    }
+    opt_defaults = {"relax_cell": relax_cell}
+    return base_opt_fn(
+        atoms,
+        preset=preset,
+        calc_defaults=calc_defaults,
+        calc_swaps=calc_kwargs,
+        opt_defaults=opt_defaults,
+        opt_params=opt_params,
+        additional_fields={"name": "VASP ASE Relax"},
+        copy_files=copy_files,
+    )
+
+
 @job
 def non_scf_job(
     atoms: Atoms,

src/quacc/schemas/_aliases/vasp.py

@@ -4,7 +4,7 @@
 
 from typing import TypedDict
 
-from quacc.schemas._aliases.ase import RunSchema
+from quacc.schemas._aliases.ase import OptSchema, RunSchema
 from quacc.schemas._aliases.emmet import TaskDoc
 
 
@@ -72,3 +72,7 @@ class QMOFRelaxSchema(VaspSchema):
     position_relax_lowacc: VaspSchema
     volume_relax_lowacc: VaspSchema | None
     double_relax: VaspSchema
+
+
+class VaspASESchema(VaspSchema, OptSchema):
+    """Type hint associated with VASP relaxations run via ASE"""

tests/core/recipes/vasp_recipes/mocked/test_vasp_recipes.py

@@ -8,7 +8,7 @@
 from quacc import SETTINGS
 from quacc.recipes.vasp.core import (
     double_relax_flow,
-    non_scf_job,
+    non_scf_job,ase_relax_job,
     relax_job,
     static_job,
 )
@@ -161,6 +161,21 @@ def test_doublerelax_flow(tmp_path, monkeypatch):
     assert double_relax_flow(atoms, relax1_kwargs={"kpts": [1, 1, 1]})
 
 
+
+def test_ase_relax_job(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+
+    atoms = bulk("Al")
+
+    output = ase_relax_job(atoms)
+    assert output["nsites"] == len(atoms)
+    assert output["parameters"]["nsw"] == 0
+    assert output["parameters"]["lwave"] is False
+    assert output["parameters"]["lcharg"] is False
+    assert output["parameters"]["encut"] == 520
+    assert output["fmax"] == 0.01
+    assert len(output["trajectory_results"]) > 1
+
 def test_non_scf_job1(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
     copy(MOCKED_DIR / "vasprun.xml.gz", tmp_path / "vasprun.xml.gz")