Repository: Quantum-Accelerators/quacc
Pull request #2137 (2024-05-13 04:27:32)

Clean up `customize_funcs`
## Summary of Changes

This makes it easier to handle updating default `job_params`.

### Checklist

- [X] I have read the ["Guidelines" section](https://quantum-accelerators.github.io/quacc/dev/contributing.html#guidelines) of the contributing guide. Don't lie! 😉
- [X] My PR is on a custom branch and is _not_ named `main`.
- [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).



Patch:
diff --git a/src/quacc/calculators/espresso/utils.py b/src/quacc/calculators/espresso/utils.py
index 4c9732ec..3c1d8450 100644
--- a/src/quacc/calculators/espresso/utils.py
+++ b/src/quacc/calculators/espresso/utils.py
@@ -39,7 +39,7 @@ def get_pseudopotential_info(
     float
         The max(ecutrho) value
     dict[str, str]
-        The pseudopotentials dictinoary, e.g. {"O": "O.pbe-n-kjpaw_psl.0.1.UPF"}
+        The pseudopotentials dictionary, e.g. {"O": "O.pbe-n-kjpaw_psl.0.1.UPF"}
     """
     unique_elements = list(set(atoms.get_chemical_symbols()))
     ecutwfc, ecutrho = 0, 0
diff --git a/src/quacc/recipes/emt/defects.py b/src/quacc/recipes/emt/defects.py
index 587b1a56..98d9be64 100644
--- a/src/quacc/recipes/emt/defects.py
+++ b/src/quacc/recipes/emt/defects.py
@@ -98,15 +98,15 @@ def bulk_to_defects_flow(
         or [quacc.schemas.ase.summarize_opt_run][].
         See the return type-hint for the data structure.
     """
-    make_defects_kwargs = recursive_dict_merge(
-        make_defects_kwargs, {"defect_gen": defect_gen, "defect_charge": defect_charge}
-    )
     relax_job_, static_job_ = customize_funcs(
         ["relax_job", "static_job"],
         [relax_job, static_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
+    make_defects_kwargs = recursive_dict_merge(
+        make_defects_kwargs, {"defect_gen": defect_gen, "defect_charge": defect_charge}
+    )
 
     return bulk_to_defects_subflow(
         atoms,
diff --git a/src/quacc/recipes/emt/elastic.py b/src/quacc/recipes/emt/elastic.py
index ccfe6203..85b0ebec 100644
--- a/src/quacc/recipes/emt/elastic.py
+++ b/src/quacc/recipes/emt/elastic.py
@@ -63,7 +63,7 @@ def bulk_to_deformations_flow(
     relax_job_, static_job_ = customize_funcs(
         ["relax_job", "static_job"],
         [relax_job, static_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
diff --git a/src/quacc/recipes/emt/phonons.py b/src/quacc/recipes/emt/phonons.py
index d030a83a..523b3a00 100644
--- a/src/quacc/recipes/emt/phonons.py
+++ b/src/quacc/recipes/emt/phonons.py
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING
 from quacc import flow
 from quacc.recipes.common.phonons import phonon_subflow
 from quacc.recipes.emt.core import relax_job, static_job
-from quacc.utils.dicts import recursive_dict_merge
 from quacc.wflow_tools.customizers import customize_funcs
 
 if TYPE_CHECKING:
@@ -71,7 +70,7 @@ def phonon_flow(
     run_relax
         Whether to run a relaxation beforehand.
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -83,13 +82,12 @@ def phonon_flow(
         Dictionary of results from [quacc.schemas.phonons.summarize_phonopy][].
         See the return type-hint for the data structure.
     """
-    calc_defaults = {"relax_job": {"opt_params": {"fmax": 1e-3}}}
-    job_params = recursive_dict_merge(calc_defaults, job_params)
-
+    job_param_defaults = {"relax_job": {"opt_params": {"fmax": 1e-3}}}
     relax_job_, static_job_ = customize_funcs(
         ["relax_job", "static_job"],
         [relax_job, static_job],
-        parameters=job_params,
+        param_defaults=job_param_defaults,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
     if run_relax:
diff --git a/src/quacc/recipes/emt/slabs.py b/src/quacc/recipes/emt/slabs.py
index 17c10b56..e74301a6 100644
--- a/src/quacc/recipes/emt/slabs.py
+++ b/src/quacc/recipes/emt/slabs.py
@@ -47,7 +47,7 @@ def bulk_to_slabs_flow(
     make_slabs_kwargs
         Additional keyword arguments to pass to [quacc.atoms.slabs.make_slabs_from_bulk][]
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -63,7 +63,7 @@ def bulk_to_slabs_flow(
     relax_job_, static_job_ = customize_funcs(
         ["relax_job", "static_job"],
         [relax_job, static_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
diff --git a/src/quacc/recipes/espresso/bands.py b/src/quacc/recipes/espresso/bands.py
index 1358d532..c0a67bbb 100644
--- a/src/quacc/recipes/espresso/bands.py
+++ b/src/quacc/recipes/espresso/bands.py
@@ -295,7 +295,7 @@ def bands_flow(
         Dictionary containing information about the parallelization of the
         calculation. See the ASE documentation for more information.
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -310,7 +310,7 @@ def bands_flow(
     (bands_pw_job_, bands_pp_job_, fermi_surface_job_) = customize_funcs(
         ["bands_pw_job", "bands_pp_job", "fermi_surface_job"],
         [bands_pw_job, bands_pp_job, fermi_surface_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
diff --git a/src/quacc/recipes/espresso/dos.py b/src/quacc/recipes/espresso/dos.py
index 61faca2a..36df2d75 100644
--- a/src/quacc/recipes/espresso/dos.py
+++ b/src/quacc/recipes/espresso/dos.py
@@ -167,7 +167,7 @@ def dos_flow(
     atoms
         Atoms object
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -204,14 +204,12 @@ def dos_flow(
     static_job_, non_scf_job_, dos_job_ = customize_funcs(
         ["static_job", "non_scf_job", "dos_job"],
         [static_job, non_scf_job, dos_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
     static_results = static_job_(atoms)
-
     non_scf_results = non_scf_job_(atoms, prev_outdir=static_results["dir_name"])
-
     dos_results = dos_job_(prev_outdir=static_results["dir_name"])
 
     return {
@@ -249,7 +247,7 @@ def projwfc_flow(
     atoms
         Atoms object
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -286,14 +284,12 @@ def projwfc_flow(
     static_job_, non_scf_job_, projwfc_job_ = customize_funcs(
         ["static_job", "non_scf_job", "projwfc_job"],
         [static_job, non_scf_job, projwfc_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
     static_results = static_job_(atoms)
-
     non_scf_results = non_scf_job_(atoms, prev_outdir=static_results["dir_name"])
-
     projwfc_results = projwfc_job_(prev_outdir=static_results["dir_name"])
 
     return {
diff --git a/src/quacc/recipes/espresso/phonons.py b/src/quacc/recipes/espresso/phonons.py
index 0ec8e955..a147ddcf 100644
--- a/src/quacc/recipes/espresso/phonons.py
+++ b/src/quacc/recipes/espresso/phonons.py
@@ -289,7 +289,7 @@ def phonon_dos_flow(
     pw_job, ph_job, fc_job, dos_job = customize_funcs(
         ["relax_job", "phonon_job", "q2r_job", "matdyn_job"],
         [relax_job, phonon_job, q2r_job, matdyn_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
@@ -365,7 +365,7 @@ def grid_phonon_flow(
         If nblocks = 0, each job will contain all the representations for a
         single q-point.
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -476,7 +476,7 @@ def grid_phonon_flow(
     pw_job, ph_init_job, ph_job, ph_recover_job = customize_funcs(
         ["relax_job", "ph_init_job", "ph_job", "ph_recover_job"],
         [relax_job, phonon_job, phonon_job, phonon_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
diff --git a/src/quacc/recipes/mlp/phonons.py b/src/quacc/recipes/mlp/phonons.py
index 035fd063..e7e42141 100644
--- a/src/quacc/recipes/mlp/phonons.py
+++ b/src/quacc/recipes/mlp/phonons.py
@@ -10,7 +10,6 @@ from monty.dev import requires
 from quacc import flow
 from quacc.recipes.common.phonons import phonon_subflow
 from quacc.recipes.mlp.core import relax_job, static_job
-from quacc.utils.dicts import recursive_dict_merge
 from quacc.wflow_tools.customizers import customize_funcs
 
 has_deps = find_spec("phonopy") is not None and find_spec("seekpath") is not None
@@ -81,7 +80,7 @@ def phonon_flow(
     t_max
         Max temperature (K).
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -93,16 +92,15 @@ def phonon_flow(
         Dictionary of results from [quacc.schemas.phonons.summarize_phonopy][].
         See the type-hint for the data structure.
     """
-    calc_defaults = {
-        "relax_job": {"method": method, "opt_params": {"fmax": 1e-3}},
-        "static_job": {"method": method},
+    job_param_defaults = {
+        "all": {"method": method},
+        "relax_job": {"opt_params": {"fmax": 1e-3}},
     }
-    job_params = recursive_dict_merge(calc_defaults, job_params)
-
     relax_job_, static_job_ = customize_funcs(
         ["relax_job", "static_job"],
         [relax_job, static_job],
-        parameters=job_params,
+        param_defaults=job_param_defaults,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
     if run_relax:
diff --git a/src/quacc/recipes/tblite/phonons.py b/src/quacc/recipes/tblite/phonons.py
index 40b5b9af..d616c63b 100644
--- a/src/quacc/recipes/tblite/phonons.py
+++ b/src/quacc/recipes/tblite/phonons.py
@@ -10,7 +10,6 @@ from monty.dev import requires
 from quacc import flow
 from quacc.recipes.common.phonons import phonon_subflow
 from quacc.recipes.tblite.core import relax_job, static_job
-from quacc.utils.dicts import recursive_dict_merge
 from quacc.wflow_tools.customizers import customize_funcs
 
 has_deps_tblite = find_spec("tblite") is not None
@@ -84,7 +83,7 @@ def phonon_flow(
     t_max
         Max temperature (K).
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -96,13 +95,13 @@ def phonon_flow(
         Dictionary of results from [quacc.schemas.phonons.summarize_phonopy][].
         See the type-hint for the data structure.
     """
-    calc_defaults = {"relax_job": {"opt_params": {"fmax": 1e-3}}}
-    job_params = recursive_dict_merge(calc_defaults, job_params)
+    job_param_defaults = {"relax_job": {"opt_params": {"fmax": 1e-3}}}
 
     relax_job_, static_job_ = customize_funcs(
         ["relax_job", "static_job"],
         [relax_job, static_job],
-        parameters=job_params,
+        param_defaults=job_param_defaults,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
     if run_relax:
diff --git a/src/quacc/recipes/vasp/mp.py b/src/quacc/recipes/vasp/mp.py
index c1fbc06e..5406da52 100644
--- a/src/quacc/recipes/vasp/mp.py
+++ b/src/quacc/recipes/vasp/mp.py
@@ -285,7 +285,7 @@ def mp_gga_relax_flow(
     (mp_gga_relax_job_, mp_gga_static_job_) = customize_funcs(
         ["mp_gga_relax_job", "mp_gga_static_job"],
         [mp_gga_relax_job, mp_gga_static_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
@@ -361,7 +361,7 @@ def mp_metagga_relax_flow(
                 "mp_metagga_static_job",
             ],
             [mp_metagga_prerelax_job, mp_metagga_relax_job, mp_metagga_static_job],
-            parameters=job_params,
+            param_swaps=job_params,
             decorators=job_decorators,
         )
     )
diff --git a/src/quacc/recipes/vasp/slabs.py b/src/quacc/recipes/vasp/slabs.py
index 475a8dc2..2d0c84c5 100644
--- a/src/quacc/recipes/vasp/slabs.py
+++ b/src/quacc/recipes/vasp/slabs.py
@@ -148,7 +148,7 @@ def bulk_to_slabs_flow(
     run_static
         Whether to run static calculations.
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -163,7 +163,7 @@ def bulk_to_slabs_flow(
     relax_job_, static_job_ = customize_funcs(
         ["relax_job", "static_job"],
         [relax_job, static_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
@@ -208,7 +208,7 @@ def slab_to_ads_flow(
     make_ads_kwargs
         Additional keyword arguments to pass to [quacc.atoms.slabs.make_adsorbate_structures][]
     job_params
-        Custom parameters to pass to each Job in the Flow. This is a dictinoary where
+        Custom parameters to pass to each Job in the Flow. This is a dictionary where
         the keys are the names of the jobs and the values are dictionaries of parameters.
     job_decorators
         Custom decorators to apply to each Job in the Flow. This is a dictionary where
@@ -223,7 +223,7 @@ def slab_to_ads_flow(
     relax_job_, static_job_ = customize_funcs(
         ["relax_job", "static_job"],
         [relax_job, static_job],
-        parameters=job_params,
+        param_swaps=job_params,
         decorators=job_decorators,
     )
 
diff --git a/src/quacc/wflow_tools/customizers.py b/src/quacc/wflow_tools/customizers.py
index eae9a938..03b14d4b 100644
--- a/src/quacc/wflow_tools/customizers.py
+++ b/src/quacc/wflow_tools/customizers.py
@@ -6,6 +6,8 @@ from copy import deepcopy
 from functools import partial
 from typing import TYPE_CHECKING, Literal
 
+from quacc.utils.dicts import recursive_dict_merge
+
 if TYPE_CHECKING:
     from typing import Any, Callable
 
@@ -138,7 +140,8 @@ def update_parameters(
 def customize_funcs(
     names: list[str] | str,
     funcs: list[Callable] | Callable,
-    parameters: dict[str, dict[str, Any]] | None = None,
+    param_defaults: dict[str, dict[str, Any]] | None = None,
+    param_swaps: dict[str, dict[str, Any]] | None = None,
     decorators: dict[str, Callable | None] | None = None,
 ) -> tuple[Callable, ...] | Callable:
     """
@@ -150,8 +153,12 @@ def customize_funcs(
         The names of the functions to customize, in the order they should be returned.
     funcs
         The functions to customize, in the order they are described in `names`.
-    parameters
-        Custom parameters to apply to each function. The keys of this dictionary correspond
+    param_defaults
+        Default parameters to apply to each function. The keys of this dictionary correspond
+        to the strings in `names`. If the key `"all"` is present, it will be applied to all
+        functions. If the value is `None`, no custom parameters will be applied to that function.
+    param_swaps
+        User-overrides of parameters to apply to each function. The keys of this dictionary correspond
         to the strings in `names`. If the key `"all"` is present, it will be applied to all
         functions. If the value is `None`, no custom parameters will be applied to that function.
     decorators
@@ -164,7 +171,7 @@ def customize_funcs(
     tuple[Callable, ...] | Callable
         The customized functions, returned in the same order as provided in `funcs`.
     """
-    parameters = parameters or {}
+    parameters = recursive_dict_merge(param_defaults, param_swaps)
     decorators = decorators or {}
     updated_funcs = []
 


Test patch:
diff --git a/tests/core/wflow/test_customizers.py b/tests/core/wflow/test_customizers.py
index d681235d..425af957 100644
--- a/tests/core/wflow/test_customizers.py
+++ b/tests/core/wflow/test_customizers.py
@@ -14,13 +14,13 @@ def test_basic_customizers():
         return a * b * c * d
 
     add_, mult_ = customize_funcs(
-        ["add", "mult"], [add, mult], parameters={"add": {"b": 2}}
+        ["add", "mult"], [add, mult], param_swaps={"add": {"b": 2}}
     )
     assert add_(1) == 5
     assert mult_(1) == 4
 
     add_, mult_ = customize_funcs(
-        ["add", "mult"], [add, mult], parameters={"add": {"b": 2}, "mult": {"b": 2}}
+        ["add", "mult"], [add, mult], param_swaps={"add": {"b": 2}, "mult": {"b": 2}}
     )
     assert add_(1) == 5
     assert mult_(1) == 8
@@ -28,20 +28,20 @@ def test_basic_customizers():
     add_, mult_ = customize_funcs(
         ["add", "mult"],
         [add, mult],
-        parameters={"add": {"b": 2}, "mult": {"b": 2}},
+        param_swaps={"add": {"b": 2}, "mult": {"b": 2}},
         decorators={"add": job(), "mult": job()},
     )
     assert add_(1) == 5
     assert mult_(1) == 8
 
     add_, mult_ = customize_funcs(
-        ["add", "mult"], [add, mult], parameters={"all": {"b": 2}}
+        ["add", "mult"], [add, mult], param_swaps={"all": {"b": 2}}
     )
     assert add_(1) == 5
     assert mult_(1) == 8
 
     with pytest.raises(ValueError):
-        customize_funcs(["add", "mult"], [add, mult], parameters={"bad": {"b": 2}})
+        customize_funcs(["add", "mult"], [add, mult], param_swaps={"bad": {"b": 2}})
 
     with pytest.raises(ValueError):
-        customize_funcs("all", [add], parameters={"all": {"b": 2}})
+        customize_funcs("all", [add], param_swaps={"all": {"b": 2}})