Repository: Quantum-Accelerators/quacc
Pull request #1552 (2024-01-15 22:26:26)

Espresso: ASE optimizers
## Requirements

### Checklist

- [x] I have read the [Contributing Guide](https://quantum-accelerators.github.io/quacc/dev/contributing.html). Don't lie! 😉
- [x] My PR is on a custom branch and is _not_ named `main`.
- [x] I have used `black`, `isort`, and `ruff` as described in the [style guide](https://quantum-accelerators.github.io/quacc/dev/contributing.html#style).
- [x] I have added relevant, comprehensive [unit tests](https://quantum-accelerators.github.io/quacc/dev/contributing.html#unit-tests).

### Notes

- Your PR will likely not be merged without proper and thorough tests.
- If you are an external contributor, you will see a comment from [@buildbot-princeton](https://github.com/buildbot-princeton). This is solely for the maintainers.
- When your code is ready for review, ping one of the [active maintainers](https://quantum-accelerators.github.io/quacc/about/contributors.html#active-maintainers).

## Summary of Changes

External optimizers for espresso!

Default does not have to be LBFGS. GPAW did test a few of them: https://wiki.fysik.dtu.dk/gpaw/devel/ase_optimize/ase_optimize.html

buildbot-princeton: Can one of the admins verify this patch?
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/1552?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`85bb036`)](https://app.codecov.io/gh/Quantum-Accelerators/quacc/commit/85bb036066ff217e013ae04f9948c0c3d37b3e4c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) 99.07% compared to head [(`1a9932f`)](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/1552?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) 99.07%.
> Report is 18 commits behind head on main.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1552      +/-   ##
==========================================
- Coverage   99.07%   99.07%   -0.01%     
==========================================
  Files          79       79              
  Lines        3029     3013      -16     
==========================================
- Hits         3001     2985      -16     
  Misses         28       28              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/1552?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).

Andrew-S-Rosen: Thanks, @tomdemeyere. Looks good to me. I just fixed some outdated docstrings. Setting this to auto-merge.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6afe0532..ea324f22 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Added
 
 - Support for ONETEP recipes (@nekkrad)
+- New Espresso recipe for ASE external relaxations (@tomdemeyere)
 
 ### Fixed
 
diff --git a/docs/user/recipes/recipes_list.md b/docs/user/recipes/recipes_list.md
index d4aebd6a..a0790f58 100644
--- a/docs/user/recipes/recipes_list.md
+++ b/docs/user/recipes/recipes_list.md
@@ -190,6 +190,7 @@ The list of available quacc recipes is shown below. The "Req'd Extras" column sp
 | ------------------------ | ---------------- | --------------------------------------------------- | ------------ |
 | Espresso Static          | `#!Python @job`  | [quacc.recipes.espresso.core.static_job][]          |              |
 | Espresso Relax           | `#!Python @job`  | [quacc.recipes.espresso.core.relax_job][]           |              |
+| Espresso ASE Relax       | `#!Python @job`  | [quacc.recipes.espresso.core.relax_job][]           |              |
 | Espresso Post Processing | `#!Python @job`  | [quacc.recipes.espresso.core.post_processing_job][] |              |
 | Espresso Phonon          | `#!Python @job`  | [quacc.recipes.espresso.phonons.phonon_job][]       |              |
 | Espresso Grid Phonon     | `#!Python @flow` | [quacc.recipes.espresso.phonons.grid_phonon_flow][] |              |
diff --git a/src/quacc/calculators/espresso/espresso.py b/src/quacc/calculators/espresso/espresso.py
index 602e7212..ea0a5683 100644
--- a/src/quacc/calculators/espresso/espresso.py
+++ b/src/quacc/calculators/espresso/espresso.py
@@ -229,8 +229,6 @@ class EspressoTemplate(EspressoTemplate_):
                     path.mkdir(parents=True, exist_ok=True)
                     input_data[section][d_key] = path
 
-        self.outdirs = [path for path in self.outdirs.values() if path is not None]
-
         parameters["input_data"] = input_data
 
         return parameters
diff --git a/src/quacc/recipes/espresso/_base.py b/src/quacc/recipes/espresso/_base.py
index e8b6c1c8..bf14c6ab 100644
--- a/src/quacc/recipes/espresso/_base.py
+++ b/src/quacc/recipes/espresso/_base.py
@@ -11,8 +11,8 @@ from quacc.calculators.espresso.espresso import (
     EspressoProfile,
     EspressoTemplate,
 )
-from quacc.runners.ase import run_calc
-from quacc.schemas.ase import summarize_run
+from quacc.runners.ase import run_calc, run_opt
+from quacc.schemas.ase import summarize_opt_run, summarize_run
 from quacc.utils.dicts import recursive_dict_merge
 
 if TYPE_CHECKING:
@@ -64,6 +64,134 @@ def base_fn(
         Dictionary of results from [quacc.schemas.ase.summarize_run][]
     """
 
+    atoms = _prepare_atoms(
+        atoms=atoms,
+        preset=preset,
+        template=template,
+        profile=profile,
+        calc_defaults=calc_defaults,
+        calc_swaps=calc_swaps,
+        parallel_info=parallel_info,
+    )
+
+    geom_file = template.outputname if template.binary == "pw" else None
+
+    final_atoms = run_calc(atoms, geom_file=geom_file, copy_files=copy_files)
+
+    return summarize_run(
+        final_atoms, atoms, additional_fields=additional_fields
+    )
+
+
+def base_opt_fn(
+    atoms: Atoms | None = None,
+    preset: str | None = None,
+    relax_cell: bool = False,
+    template: EspressoTemplate | None = None,
+    profile: EspressoProfile | 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,
+    parallel_info: dict[str, Any] | None = None,
+    additional_fields: dict[str, Any] | None = None,
+    copy_files: list[str] | None = None,
+) -> RunSchema:
+    """
+    Base function to carry out espresso recipes.
+
+    Parameters
+    ----------
+    atoms
+        Atoms object
+    preset
+        Name of the preset to use
+    relax_cell
+        Whether to relax the cell or not.
+    template
+        EspressoTemplate to use
+    profile
+        EspressoProfile to use
+    calc_defaults
+        The default calculator parameters.
+    calc_swaps
+        Custom kwargs for the espresso calculator. Set a value to
+        `quacc.Remove` to remove a pre-existing key entirely. For a list of available
+        keys, refer to the `ase.calculators.espresso.Espresso` calculator.
+    opt_defaults
+        The default optimization parameters.
+    opt_params
+        Dictionary of parameters to pass to the optimizer. pass "optimizer"
+        to change the optimizer being used. "fmax" and "max_steps" are commonly
+        used keywords. See the ASE documentation for more information.
+    parallel_info
+        Dictionary of parallelization information.
+    additional_fields
+        Any additional fields to supply to the summarizer.
+    copy_files
+        Files to copy to the runtime directory.
+
+    Returns
+    -------
+    RunSchema
+        Dictionary of results from [quacc.schemas.ase.summarize_run][]
+    """
+
+    atoms = _prepare_atoms(
+        atoms=atoms,
+        preset=preset,
+        template=template,
+        profile=profile,
+        calc_defaults=calc_defaults,
+        calc_swaps=calc_swaps,
+        parallel_info=parallel_info,
+    )
+
+    opt_flags = recursive_dict_merge(opt_defaults, opt_params)
+
+    dyn = run_opt(atoms, relax_cell=relax_cell, copy_files=copy_files, **opt_flags)
+
+    return summarize_opt_run(dyn, additional_fields=additional_fields)
+
+
+def _prepare_atoms(
+    atoms: Atoms | None = None,
+    preset: str | None = None,
+    template: EspressoTemplate | None = None,
+    profile: EspressoProfile | None = None,
+    calc_defaults: dict[str, Any] | None = None,
+    calc_swaps: dict[str, Any] | None = None,
+    parallel_info: dict[str, Any] | None = None,
+) -> Atoms:
+    """
+    Commonly used preparation function to merge parameters
+    and attach an Esprseso calculator accordingly.
+
+    Parameters
+    ----------
+    atoms
+        Atoms object
+    preset
+        Name of the preset to use
+    template
+        EspressoTemplate to use
+    profile
+        EspressoProfile to use
+    calc_defaults
+        The default calculator parameters.
+    calc_swaps
+        Custom kwargs for the espresso calculator. Set a value to
+        `quacc.Remove` to remove a pre-existing key entirely. For a list of available
+        keys, refer to the `ase.calculators.espresso.Espresso` calculator.
+    parallel_info
+        Dictionary of parallelization information.
+
+    Returns
+    -------
+    Atoms
+        Atoms object with attached Espresso calculator.
+    """
+
     atoms = Atoms() if atoms is None else atoms
 
     calc_defaults["input_data"] = Namelist(calc_defaults.get("input_data"))
@@ -85,8 +213,4 @@ def base_fn(
         **calc_flags,
     )
 
-    geom_file = template.outputname if template.binary == "pw" else None
-
-    final_atoms = run_calc(atoms, geom_file=geom_file, copy_files=copy_files)
-
-    return summarize_run(final_atoms, atoms, additional_fields=additional_fields)
+    return atoms
diff --git a/src/quacc/recipes/espresso/core.py b/src/quacc/recipes/espresso/core.py
index 6ec6e35a..32660127 100644
--- a/src/quacc/recipes/espresso/core.py
+++ b/src/quacc/recipes/espresso/core.py
@@ -3,12 +3,15 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+from ase.optimize import LBFGS
+
 from quacc import job
 from quacc.calculators.espresso.espresso import EspressoTemplate
-from quacc.recipes.espresso._base import base_fn
+from quacc.recipes.espresso._base import base_fn, base_opt_fn
 
 if TYPE_CHECKING:
     from pathlib import Path
+    from typing import Any
 
     from ase.atoms import Atoms
 
@@ -148,6 +151,83 @@ def relax_job(
     )
 
 
+@job
+def ase_relax_job(
+    atoms: Atoms,
+    preset: str | None = "sssp_1.3.0_pbe_efficiency",
+    relax_cell: bool = False,
+    parallel_info: dict[str] | None = None,
+    opt_params: dict[str, Any] | None = None,
+    copy_files: str | Path | list[str | Path] | None = None,
+    **calc_kwargs,
+) -> RunSchema:
+    """
+    Function to carry out a structure relaxation with pw.x using ASE
+    external optimizers.
+
+    Parameters
+    ----------
+    atoms
+        The Atoms object.
+    preset
+        The name of a YAML file containing a list of parameters to use as
+        a "preset" for the calculator. quacc will automatically look in the
+        `ESPRESSO_PRESET_DIR` (default: quacc/calculators/espresso/presets).
+    relax_cell
+        Whether to relax the cell or not.
+    parallel_info
+        Dictionary containing information about the parallelization of the
+        calculation. See the ASE documentation for more information.
+    opt_params
+        Dictionary of parameters to pass to the optimizer. pass "optimizer"
+        to change the optimizer being used. "fmax" and "max_steps" are commonly
+        used keywords. See the ASE documentation for more information.
+    copy_files
+        List of files to copy to the calculation directory. Useful for copying
+        files from a previous calculation. This parameter can either be a string
+        or a list of strings.
+
+        If a string is provided, it is assumed to be a path to a directory,
+        all of the child tree structure of that directory is going to be copied to the
+        scratch of this calculation. For phonon_job this is what most users will want to do.
+
+        If a list of strings is provided, each string point to a specific file. In this case
+        it is important to note that no directory structure is going to be copied, everything
+        is copied at the root of the temporary directory.
+    **calc_kwargs
+        Additional keyword arguments to pass to the Espresso calculator. See the
+        docstring of [quacc.calculators.espresso.espresso.Espresso][] for more information.
+
+    Returns
+    -------
+    RunSchema
+        Dictionary of results from [quacc.schemas.ase.summarize_run][].
+        See the type-hint for the data structure.
+    """
+
+    calc_defaults = {
+        "input_data": {
+            "control": {"calculation": "scf", "tstress": relax_cell, "tprnfor": True}
+        }
+    }
+
+    opt_defaults = {"fmax": 0.01, "max_steps": 1000, "optimizer": LBFGS}
+
+    return base_opt_fn(
+        atoms,
+        preset=preset,
+        relax_cell=relax_cell,
+        template=EspressoTemplate("pw"),
+        calc_defaults=calc_defaults,
+        calc_swaps=calc_kwargs,
+        opt_defaults=opt_defaults,
+        opt_params=opt_params,
+        parallel_info=parallel_info,
+        additional_fields={"name": "pw.x ExternalRelax"},
+        copy_files=copy_files,
+    )
+
+
 @job
 def post_processing_job(
     prev_dir: str | Path,


Test patch:
diff --git a/tests/core/calculators/espresso/test_espresso.py b/tests/core/calculators/espresso/test_espresso.py
index 3b40facf..e9b812be 100644
--- a/tests/core/calculators/espresso/test_espresso.py
+++ b/tests/core/calculators/espresso/test_espresso.py
@@ -152,7 +152,6 @@ def test_outdir_handler(tmp_path, monkeypatch):
     test_path = "../test3/test2/test1"
     fake_template = EspressoTemplate()
     new_parameters = fake_template._outdir_handler(parameters, Path())
-    assert len(fake_template.outdirs) == 1
     assert new_parameters["input_data"]["system"]["outdir"] == test_path
     assert not Path(test_path).exists()
 
@@ -160,7 +159,6 @@ def test_outdir_handler(tmp_path, monkeypatch):
     parameters["input_data"]["system"]["outdir"] = test_path
     fake_template = EspressoTemplate()
     new_parameters = fake_template._outdir_handler(parameters, Path())
-    assert len(fake_template.outdirs) == 1
     assert new_parameters["input_data"]["system"]["outdir"] == test_path
     assert not Path(test_path).exists()
 
@@ -168,7 +166,6 @@ def test_outdir_handler(tmp_path, monkeypatch):
     parameters["input_data"]["system"]["outdir"] = test_path
     fake_template = EspressoTemplate()
     new_parameters = fake_template._outdir_handler(parameters, Path())
-    assert len(fake_template.outdirs) == 0
     assert (
         new_parameters["input_data"]["system"]["outdir"]
         == test_path.expanduser().resolve()
diff --git a/tests/core/recipes/espresso_recipes/test_core.py b/tests/core/recipes/espresso_recipes/test_core.py
index 6ddcca68..9d402de2 100644
--- a/tests/core/recipes/espresso_recipes/test_core.py
+++ b/tests/core/recipes/espresso_recipes/test_core.py
@@ -1,14 +1,22 @@
+from __future__ import annotations
+
 from pathlib import Path
 from shutil import which
 from subprocess import CalledProcessError
 
 import pytest
 from ase.build import bulk
+from ase.optimize import BFGS
 from numpy.testing import assert_allclose, assert_array_equal
 
 from quacc import SETTINGS
 from quacc.calculators.espresso.espresso import EspressoTemplate
-from quacc.recipes.espresso.core import post_processing_job, relax_job, static_job
+from quacc.recipes.espresso.core import (
+    ase_relax_job,
+    post_processing_job,
+    relax_job,
+    static_job,
+)
 from quacc.utils.files import copy_decompress_files
 
 pytestmark = pytest.mark.skipif(which("pw.x") is None, reason="QE not installed")
@@ -177,9 +185,7 @@ def test_static_job_test_run(tmp_path, monkeypatch):
 
     pseudopotentials = {"Si": "Si.upf"}
 
-    EspressoTemplate._test_run(
-        {"input_data": {"control": {"prefix": "test"}}}, Path(".")
-    )
+    EspressoTemplate._test_run({"input_data": {"control": {"prefix": "test"}}}, Path())
 
     assert Path("test.EXIT").exists()
 
@@ -217,6 +223,66 @@ def test_relax_job(tmp_path, monkeypatch):
     assert new_input_data["control"]["calculation"] == "relax"
 
 
+def test_ase_relax_job(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+
+    copy_decompress_files([DATA_DIR / "Si.upf.gz"], tmp_path)
+
+    atoms = bulk("Si")
+    atoms[0].position += 0.05
+
+    pseudopotentials = {"Si": "Si.upf"}
+    input_data = {"control": {"pseudo_dir": tmp_path}}
+
+    results = ase_relax_job(
+        atoms,
+        input_data=input_data,
+        pseudopotentials=pseudopotentials,
+        kpts=None,
+        opt_params={"max_steps": 10, "fmax": 1.0e-1, "optimizer": BFGS},
+    )
+
+    with pytest.raises(AssertionError):
+        assert_allclose(
+            results["atoms"].get_positions(), atoms.get_positions(), atol=1.0e-4
+        )
+    assert_allclose(results["atoms"].get_cell(), atoms.get_cell(), atol=1.0e-3)
+
+    assert len(results["trajectory"]) == 3
+    assert_allclose(
+        results["trajectory_results"][-1]["free_energy"], -293.71198070497894
+    )
+    new_input_data = results["parameters"]["input_data"]
+    assert new_input_data["control"]["calculation"] == "scf"
+
+
+def test_ase_relax_cell_job(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+
+    copy_decompress_files([DATA_DIR / "Si.upf.gz"], tmp_path)
+
+    atoms = bulk("Si")
+    atoms[0].position += 0.05
+
+    pseudopotentials = {"Si": "Si.upf"}
+    input_data = {
+        "control": {"pseudo_dir": tmp_path},
+        "occupations": "smearing",
+        "smearing": "gaussian",
+        "degauss": 0.005,
+    }
+
+    with pytest.raises(RuntimeError):
+        ase_relax_job(
+            atoms,
+            relax_cell=True,
+            input_data=input_data,
+            pseudopotentials=pseudopotentials,
+            kpts=None,
+            opt_params={"max_steps": 2, "fmax": 1.0e-1, "optimizer": BFGS},
+        )
+
+
 def test_relax_job_cell(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)