Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add an ase_relax_job for VASP #1888

Merged
merged 8 commits into from
Mar 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/quacc/recipes/onetep/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
66 changes: 64 additions & 2 deletions src/quacc/recipes/vasp/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)

Expand All @@ -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)
56 changes: 54 additions & 2 deletions src/quacc/recipes/vasp/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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,
Expand Down
6 changes: 5 additions & 1 deletion src/quacc/schemas/_aliases/vasp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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"""
17 changes: 16 additions & 1 deletion tests/core/recipes/vasp_recipes/mocked/test_vasp_recipes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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")
Expand Down
Loading