Code cleanup enabled by Parsl lifting of [] into App (#1062)

- [X] I have read the [Developer 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 added relevant unit tests. Note: Your PR will likely not be merged without proper tests. Waiting on: - Parsl/parsl#2904 Closes #1059 (and #737).
Quantum-Accelerators · Oct 14, 2023 · 839c07b · 839c07b
1 parent f09b3e2
commit 839c07b
Show file tree

Hide file tree

Showing 27 changed files with 146 additions and 229 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,17 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.3.7]
+
+### Changed
+
+- Recipes now take `Atoms` instead of `Atoms | dict` as input
+- Recipes no longer require the use of `fetch_atoms`
+
+### Removed
+
+- Removed the `fetch_atoms` function
+
 ## [0.3.6]
 
 ### Added

diff --git a/pyproject.toml b/pyproject.toml
@@ -45,7 +45,7 @@ defects = ["pymatgen-analysis-defects>=2023.8.22", "shakenbreak>=3.2.0"]
 jobflow = ["jobflow>=0.1.14", "fireworks>=2.0.3"]
 newtonnet = ["torch>=2.0.1", "scikit-learn>=1.3.0"]
 sella = ["sella>=2.3.2"]
-parsl = ["parsl>=2023.8.14"]
+parsl = ["parsl@git+https://github.com/Parsl/parsl.git@benc-lifted-dict"]
 prefect = ["prefect>=2.13.1", "prefect-dask>=0.2.4", "dask-jobqueue>=0.8.2"]
 redun = ["redun>=0.16.2"]
 tblite = ["tblite[ase]>=0.3.0; platform_system=='Linux'"]

diff --git a/src/quacc/__init__.py b/src/quacc/__init__.py
@@ -7,11 +7,10 @@
 from ase import __version__ as ase_version
 from ase.io.jsonio import decode, encode
 
-from quacc.schemas.atoms import fetch_atoms
 from quacc.settings import QuaccSettings
 from quacc.wflow.decorators import flow, job, subflow
 
-__all__ = ["flow", "job", "subflow", "fetch_atoms"]
+__all__ = ["flow", "job", "subflow"]
 
 
 def atoms_as_dict(s: Atoms) -> dict:

diff --git a/src/quacc/recipes/dftb/core.py b/src/quacc/recipes/dftb/core.py
@@ -5,7 +5,7 @@
 
 from ase.calculators.dftb import Dftb
 
-from quacc import SETTINGS, fetch_atoms, job
+from quacc import SETTINGS, job
 from quacc.runners.calc import run_calc
 from quacc.schemas.ase import summarize_run
 from quacc.utils.dicts import merge_dicts
@@ -24,7 +24,7 @@
 
 @job
 def static_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     method: Literal["GFN1-xTB", "GFN2-xTB", "DFTB"] = "GFN2-xTB",
     kpts: tuple | list[tuple] | dict | None = None,
     calc_swaps: dict | None = None,
@@ -48,8 +48,7 @@ def static_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     method
         Method to use.
     kpts
@@ -89,7 +88,7 @@ def static_job(
 
 @job
 def relax_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     method: Literal["GFN1-xTB", "GFN2-xTB", "DFTB"] = "GFN2-xTB",
     kpts: tuple | list[tuple] | dict | None = None,
     relax_cell: bool = False,
@@ -118,8 +117,7 @@ def relax_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     method
         Method to use.
     kpts
@@ -165,7 +163,7 @@ def relax_job(
 
 
 def _base_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     defaults: dict | None = None,
     calc_swaps: dict | None = None,
     additional_fields: dict | None = None,
@@ -177,8 +175,7 @@ def _base_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     defaults
         The default calculator parameters to use.
     calc_swaps
@@ -195,7 +192,6 @@ def _base_job(
         Dictionary of results, specified in [quacc.schemas.ase.summarize_run][]
     """
 
-    atoms = fetch_atoms(atoms)
     flags = merge_dicts(defaults, calc_swaps)
 
     atoms.calc = Dftb(**flags)

diff --git a/src/quacc/recipes/emt/core.py b/src/quacc/recipes/emt/core.py
@@ -10,7 +10,7 @@
 from ase.calculators.emt import EMT
 from ase.optimize import FIRE
 
-from quacc import fetch_atoms, job
+from quacc import job
 from quacc.runners.calc import run_ase_opt, run_calc
 from quacc.schemas.ase import summarize_opt_run, summarize_run
 from quacc.utils.dicts import merge_dicts
@@ -23,7 +23,7 @@
 
 @job
 def static_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     calc_swaps: dict | None = None,
     copy_files: list[str] | None = None,
 ) -> RunSchema:
@@ -41,8 +41,7 @@ def static_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     calc_swaps
         Dictionary of custom kwargs for the EMT calculator.
     copy_files
@@ -53,7 +52,6 @@ def static_job(
     RunSchema
         Dictionary of results, specified in [quacc.schemas.ase.summarize_run][]
     """
-    atoms = fetch_atoms(atoms)
     calc_swaps = calc_swaps or {}
 
     atoms.calc = EMT(**calc_swaps)
@@ -68,7 +66,7 @@ def static_job(
 
 @job
 def relax_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     relax_cell: bool = False,
     calc_swaps: dict | None = None,
     opt_swaps: dict | None = None,
@@ -94,8 +92,7 @@ def relax_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     relax_cell
         Whether to relax the cell
     calc_swaps
@@ -111,7 +108,6 @@ def relax_job(
     OptSchema
         Dictionary of results, specified in [quacc.schemas.ase.summarize_opt_run][]
     """
-    atoms = fetch_atoms(atoms)
     calc_swaps = calc_swaps or {}
 
     opt_defaults = {"fmax": 0.01, "max_steps": 1000, "optimizer": FIRE}

diff --git a/src/quacc/recipes/emt/defects.py b/src/quacc/recipes/emt/defects.py
@@ -5,7 +5,7 @@
 
 from pymatgen.analysis.defects.generators import VacancyGenerator
 
-from quacc import fetch_atoms, flow, job, subflow
+from quacc import flow, job, subflow
 from quacc.atoms.defects import make_defects_from_bulk
 from quacc.recipes.emt.core import relax_job, static_job
 
@@ -24,7 +24,7 @@
 
 @flow
 def bulk_to_defects_flow(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     defect_gen: (
         AntiSiteGenerator
         | ChargeInterstitialGenerator
@@ -80,7 +80,6 @@ def bulk_to_defects_flow(
 
     @job
     def _make_defects(atoms):
-        atoms = fetch_atoms(atoms)
         return make_defects_from_bulk(
             atoms,
             defect_gen=defect_gen,
@@ -96,7 +95,7 @@ def _relax_distributed(defects):
     def _relax_and_static_distributed(defects):
         return [
             static_job(
-                relax_job(defect, **defect_relax_kwargs),
+                relax_job(defect, **defect_relax_kwargs)["atoms"],
                 **defect_static_kwargs,
             )
             for defect in defects

diff --git a/src/quacc/recipes/emt/slabs.py b/src/quacc/recipes/emt/slabs.py
@@ -3,7 +3,7 @@
 
 from typing import TYPE_CHECKING
 
-from quacc import fetch_atoms, flow, job, subflow
+from quacc import flow, job, subflow
 from quacc.atoms.slabs import make_slabs_from_bulk
 from quacc.recipes.emt.core import relax_job, static_job
 
@@ -15,7 +15,7 @@
 
 @flow
 def bulk_to_slabs_flow(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     make_slabs_kwargs: dict | None = None,
     run_static: bool = True,
     slab_relax_kwargs: dict | None = None,
@@ -33,8 +33,7 @@ def bulk_to_slabs_flow(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     make_slabs_kwargs
         Additional keyword arguments to pass to
         [quacc.atoms.slabs.make_slabs_from_bulk][]
@@ -60,7 +59,6 @@ def bulk_to_slabs_flow(
 
     @job
     def _make_slabs(atoms):
-        atoms = fetch_atoms(atoms)
         return make_slabs_from_bulk(atoms, **make_slabs_kwargs)
 
     @subflow
@@ -71,7 +69,7 @@ def _relax_distributed(slabs):
     def _relax_and_static_distributed(slabs):
         return [
             static_job(
-                relax_job(slab, **slab_relax_kwargs),
+                relax_job(slab, **slab_relax_kwargs)["atoms"],
                 **slab_static_kwargs,
             )
             for slab in slabs

diff --git a/src/quacc/recipes/gaussian/core.py b/src/quacc/recipes/gaussian/core.py
@@ -6,7 +6,7 @@
 
 from ase.calculators.gaussian import Gaussian
 
-from quacc import fetch_atoms, job
+from quacc import job
 from quacc.runners.calc import run_calc
 from quacc.schemas.cclib import cclib_summarize_run
 from quacc.utils.dicts import merge_dicts
@@ -22,7 +22,7 @@
 
 @job
 def static_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     charge: int,
     spin_multiplicity: int,
     xc: str = "wb97x-d",
@@ -59,8 +59,7 @@ def static_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     charge
         Charge of the system.
     spin_multiplicity
@@ -148,8 +147,7 @@ def relax_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     charge
         Charge of the system.
     spin_multiplicity
@@ -200,7 +198,7 @@ def relax_job(
 
 
 def _base_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     charge: int,
     spin_multiplicity: int,
     defaults: dict | None = None,
@@ -214,8 +212,7 @@ def _base_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     charge
         Charge of the system.
     spin_multiplicity
@@ -235,7 +232,6 @@ def _base_job(
     cclibSchema
         Dictionary of results, as specified in [quacc.schemas.cclib.cclib_summarize_run][]
     """
-    atoms = fetch_atoms(atoms)
     flags = merge_dicts(defaults, calc_swaps)
 
     atoms.calc = Gaussian(**flags)

diff --git a/src/quacc/recipes/gulp/core.py b/src/quacc/recipes/gulp/core.py
@@ -6,7 +6,7 @@
 
 from ase.calculators.gulp import GULP
 
-from quacc import SETTINGS, fetch_atoms, job
+from quacc import SETTINGS, job
 from quacc.runners.calc import run_calc
 from quacc.schemas.ase import summarize_run
 from quacc.utils.dicts import merge_dicts
@@ -24,7 +24,7 @@
 
 @job
 def static_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     use_gfnff: bool = True,
     library: str | None = None,
     keyword_swaps: dict | None = None,
@@ -58,8 +58,7 @@ def static_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     use_gfnff
         True if (p)GFN-FF should be used; False if not.
     library
@@ -97,7 +96,7 @@ def static_job(
 
 @job
 def relax_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     use_gfnff: bool = True,
     library: str | None = None,
     relax_cell: bool = False,
@@ -135,8 +134,7 @@ def relax_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     use_gfnff
         True if (p)GFN-FF should be used; False if not.
     library
@@ -178,7 +176,7 @@ def relax_job(
 
 
 def _base_job(
-    atoms: Atoms | dict,
+    atoms: Atoms,
     library: str | None = None,
     keyword_defaults: dict | None = None,
     option_defaults: dict | None = None,
@@ -193,8 +191,7 @@ def _base_job(
     Parameters
     ----------
     atoms
-        Atoms object or a dictionary with the key "atoms" and an Atoms object as
-        the value
+        Atoms object
     library
         Filename of the potential library file, if required.
     keyword_defaults
@@ -215,7 +212,6 @@ def _base_job(
     RunSchema
         Dictionary of results from [quacc.schemas.ase.summarize_run][]
     """
-    atoms = fetch_atoms(atoms)
     keyword_defaults = keyword_defaults or {}
 
     if not atoms.pbc.any():