Repository: Quantum-Accelerators/quacc
Pull request #2058 (2024-04-25 20:47:45)

Add support for function hooks in the optimization routines
## Summary of Changes

@tomdemeyere: I wonder if something like this would enable you to do optimizer "patching" (https://github.com/Quantum-Accelerators/quacc/discussions/1556) without introducing lots of complexity to quacc. I didn't really give this much thought so maybe I'm completely wrong. 

tomdemeyere: @Andrew-S-Rosen
This is exciting! I think we're on the same page, but I have a few questions to clarify:

- If we want to use this to, for example, change the temperature linearly during an MD run, do we need to provide all the temperatures directly inside the function that modifies opt?
``` python
def my_temperature_changer(opt):
opt.set_temperatures(...)
```

Essentially, I'm wondering if we can modify the attributes of the dyn object during the dynamics. I believe that's the case, but I just want to confirm.

- If my understanding above is correct, then the main difference compared to my original idea is that "proper patching" would provide more flexibility regarding when things run. We could execute code at the beginning or end of a step, at the very start of the simulation, or only upon convergence. Additionally, we could add custom actions whenever a specific method is called, offering even greater flexibility.
- 
> without introducing lots of complexity to quacc

If you're referring to complexity on the code side, then absolutely, this option is preferable.

In all fairness, users can achieve all of this by superclassing the desired dynamics class themselves. The real question is: should quacc help this process (like providing ready-made `fn_hook` that performs specific task)? If the answer is yes, then the method showcased here takes a small step in that direction, while full patching goes all-in by providing complete control,  with the tradeoff of increased code complexity.
Andrew-S-Rosen: > Essentially, I'm wondering if we can modify the attributes of the dyn object during the dynamics. I believe that's the case, but I just want to confirm.

That was indeed my thinking --- that by passing the `dyn` object, the user has the ability to do whatever they want to it (including accessing/modifying attributes mid-loop). I have not personally tested this, but I can't see why that wouldn't be the case.

> If my understanding above is correct, then the main difference compared to my original idea is that "proper patching" would provide more flexibility regarding when things run. We could execute code at the beginning or end of a step, at the very start of the simulation, or only upon convergence. Additionally, we could add custom actions whenever a specific method is called, offering even greater flexibility.

Yes, here we are basically providing the user with the ability to do "whatever they want" at the end of the iteration. They could, for instance, kill the optimization entirely by toggling whether it is converged or not manually. It would not be difficult to split this into a `pre_fn_hook` and `post_fn_hook` so that users can have two separate function calls at start and end of the loop. The logic is largely the same.

> In all fairness, users can achieve all of this by superclassing the desired dynamics class themselves.

This is true. One could make their own class with its own custom logic and pass that in. That said, if someone wanted to do something like "upload to the cloud after every step" it seems like it might be a bit unnecessary to have the user make a custom ASE optimizer class just for this when in reality what they want is to call a simple function they have available.

>  The real question is: should quacc help this process (like providing ready-made fn_hook that performs specific task)? If the answer is yes, then the method showcased here takes a small step in that direction, while full patching goes all-in by providing complete control, with the tradeoff of increased code complexity.

At minimum, we can give users the ability to have that flexibility. If there are a set of function hooks that are likely to be general-purpose or necessary to carry out specific recipes, I think it would be potentially worth supporting. We could have a `quacc/runners/fn_hooks.py` module to store pre-made callables, for instance.

There are likely pros/cons compared to https://github.com/Quantum-Accelerators/quacc/discussions/1556, but one reason why I have proposed this approach is that it is a very simple entry point with a lot of flexibility. One does not necessarily need to know anything about an ASE dynamics class to inject logic.
tomdemeyere: > They could, for instance, kill the optimization entirely by toggling whether it is converged or not manually. It would not be difficult to split this into a pre_fn_hook and post_fn_hook

That's pretty much it, for example changing temperature would need to be *before* every step.

I would be happy to provide some pre-made functions once you are done working on it.
Andrew-S-Rosen: Cool! I'll modify this PR to have a separate pre and post hook then. I think that should cover the vast majority of use cases. And I'll fix the remaining test failure because apparently SciPy optimizers do not play nicely with `irun`.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/2058?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) Report
Attention: Patch coverage is `85.71429%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 99.09%. Comparing base [(`b980813`)](https://app.codecov.io/gh/Quantum-Accelerators/quacc/commit/b980813e6bc4c26368ba23be2e0c18e66dcf777e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) to head [(`d11540c`)](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/2058?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators).

| [Files](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/2058?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) | Patch % | Lines |
|---|---|---|
| [src/quacc/runners/ase.py](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/2058?src=pr&el=tree&filepath=src%2Fquacc%2Frunners%2Fase.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators#diff-c3JjL3F1YWNjL3J1bm5lcnMvYXNlLnB5) | 85.71% | [1 Missing :warning: ](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/2058?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Quantum-Accelerators) |

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


```diff
@@            Coverage Diff             @@
##             main    #2058      +/-   ##
==========================================
- Coverage   99.12%   99.09%   -0.03%     
==========================================
  Files          81       81              
  Lines        3300     3303       +3     
==========================================
+ Hits         3271     3273       +2     
- Misses         29       30       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/Quantum-Accelerators/quacc/pull/2058?dropdown=coverage&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: Hmmm... it is not immediately obvious to me how to get it to activate a function _before_ each step. 🤔 I'll merge this as-is for now. Happy to modify later --- it's experimental for now.

Patch:
diff --git a/src/quacc/runners/ase.py b/src/quacc/runners/ase.py
index ccfa2dc1..69e17c17 100644
--- a/src/quacc/runners/ase.py
+++ b/src/quacc/runners/ase.py
@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import sys
 from shutil import copy, copytree
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 import numpy as np
 from ase.filters import FrechetCellFilter
@@ -147,6 +147,7 @@ def run_opt(
     optimizer: Optimizer = BFGS,
     optimizer_kwargs: OptimizerKwargs | None = None,
     store_intermediate_results: bool = False,
+    fn_hook: Callable | None = None,
     run_kwargs: dict[str, Any] | None = None,
     copy_files: SourceDirectory | dict[SourceDirectory, Filenames] | None = None,
 ) -> Optimizer:
@@ -178,6 +179,10 @@ def run_opt(
         Whether to store the files generated at each intermediate step in the
         optimization. If enabled, they will be stored in a directory named
         `stepN` where `N` is the step number, starting at 0.
+    fn_hook
+        A custom function to call after each step of the optimization.
+        The function must take the instantiated dynamics class as
+        its only argument.
     run_kwargs
         Dictionary of kwargs for the run() method of the optimizer.
     copy_files
@@ -228,19 +233,23 @@ def run_opt(
 
     # Run optimization
     with traj, optimizer(atoms, **optimizer_kwargs) as dyn:
-        if store_intermediate_results:
-            for i, _ in enumerate(dyn.irun(fmax=fmax, steps=max_steps, **run_kwargs)):
-                _copy_intermediate_files(
-                    tmpdir,
-                    i,
-                    files_to_ignore=[
-                        traj_file,
-                        optimizer_kwargs["restart"],
-                        optimizer_kwargs["logfile"],
-                    ],
-                )
-        else:
+        if optimizer.__name__.startswith("SciPy"):
+            # https://gitlab.com/ase/ase/-/issues/1475
             dyn.run(fmax=fmax, steps=max_steps, **run_kwargs)
+        else:
+            for i, _ in enumerate(dyn.irun(fmax=fmax, steps=max_steps, **run_kwargs)):
+                if store_intermediate_results:
+                    _copy_intermediate_files(
+                        tmpdir,
+                        i,
+                        files_to_ignore=[
+                            traj_file,
+                            optimizer_kwargs["restart"],
+                            optimizer_kwargs["logfile"],
+                        ],
+                    )
+                if fn_hook:
+                    fn_hook(dyn)
 
     # Store the trajectory atoms
     dyn.traj_atoms = read(traj_file, index=":")


Test patch:
diff --git a/tests/core/runners/test_ase.py b/tests/core/runners/test_ase.py
index ec308f59..f1a9df60 100644
--- a/tests/core/runners/test_ase.py
+++ b/tests/core/runners/test_ase.py
@@ -223,3 +223,17 @@ def test_unique_workdir(tmp_path, monkeypatch):
 
     SETTINGS.CREATE_UNIQUE_DIR = DEFAULT_SETTINGS.CREATE_UNIQUE_DIR
     SETTINGS.RESULTS_DIR = DEFAULT_SETTINGS.RESULTS_DIR
+
+
+def test_fn_hook(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+
+    def fn_hook(dyn):
+        if dyn.atoms:
+            raise ValueError("Test error")
+
+    atoms = bulk("Cu")
+    atoms.calc = EMT()
+
+    with pytest.raises(ValueError, match="Test error"):
+        run_opt(atoms, fn_hook=fn_hook)