Write out results schema to filesystem in summarize functions (#1729)

## Summary of Changes

This has two purposes:

- Make sure the results are written safely somewhere while the file
database problem is not sorted

- For calculators where nothing is written to the disk (MACE, EMT, ...)
this results file gives sense to the currently empty directories

What do you think?

EDIT: Probably not going to be as easy as I was expecting

### Checklist

- [ ] I have read the ["Guidelines"
section](https://quantum-accelerators.github.io/quacc/dev/contributing.html#guidelines)
of the contributing guide. Don't lie! 😉
- [ ] My PR is on a custom branch and is _not_ named `main`.
- [ ] I have used `black`, `isort`, and `ruff` as described in the
[style
guide](https://quantum-accelerators.github.io/quacc/dev/contributing.html#style).
- [ ] 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).

---------

Co-authored-by: Andrew S. Rosen <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/quacc/schemas/ase.py

@@ -2,11 +2,14 @@
 
 from __future__ import annotations
 
+import pickle
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import numpy as np
 from ase import units
 from ase.io import read
+from ase.vibrations import Vibrations
 from ase.vibrations.data import VibrationsData
 
 from quacc import SETTINGS, __version__
@@ -24,7 +27,6 @@
     from ase.io import Trajectory
     from ase.optimize.optimize import Optimizer
     from ase.thermochemistry import IdealGasThermo
-    from ase.vibrations import Vibrations
     from maggma.core import Store
 
     from quacc.schemas._aliases.ase import (
@@ -117,6 +119,10 @@ def summarize_run(
     )
     task_doc = clean_task_doc(unsorted_task_doc)
 
+    if SETTINGS.WRITE_PICKLE:
+        with Path(directory, "quacc_results.pkl").open("wb") as f:
+            pickle.dump(task_doc, f)
+
     if store:
         results_to_db(store, task_doc)
 
@@ -217,6 +223,10 @@ def summarize_opt_run(
     )
     task_doc = clean_task_doc(unsorted_task_doc)
 
+    if SETTINGS.WRITE_PICKLE:
+        with Path(directory, "quacc_results.pkl").open("wb") as f:
+            pickle.dump(task_doc, f)
+
     if store:
         results_to_db(store, task_doc)
 
@@ -279,6 +289,12 @@ def summarize_vib_and_thermo(
     )
     task_doc = clean_task_doc(unsorted_task_doc)
 
+    if isinstance(vib, Vibrations):
+        directory = vib.atoms.calc.directory
+        if SETTINGS.WRITE_PICKLE:
+            with Path(directory, "quacc_results.pkl").open("wb") as f:
+                pickle.dump(task_doc, f)
+
     if store:
         results_to_db(store, task_doc)
 

src/quacc/schemas/cclib.py

@@ -4,6 +4,7 @@
 
 import logging
 import os
+import pickle
 from inspect import getmembers, isclass
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -137,6 +138,10 @@ def cclib_summarize_run(
     unsorted_task_doc = run_task_doc | cclib_task_doc | additional_fields
     task_doc = clean_task_doc(unsorted_task_doc)
 
+    if SETTINGS.WRITE_PICKLE:
+        with Path(dir_path, "quacc_results.pkl").open("wb") as f:
+            pickle.dump(task_doc, f)
+
     # Store the results
     if store:
         results_to_db(store, task_doc)

src/quacc/schemas/phonons.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import pickle
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -89,6 +90,10 @@ def summarize_phonopy(
     )
     task_doc = clean_task_doc(unsorted_task_doc)
 
+    if SETTINGS.WRITE_PICKLE:
+        with Path(directory, "quacc_results.pkl").open("wb") as f:
+            pickle.dump(task_doc, f)
+
     if store:
         results_to_db(store, task_doc)
 

src/quacc/schemas/vasp.py

@@ -4,6 +4,7 @@
 
 import logging
 import os
+import pickle
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -145,6 +146,10 @@ def vasp_summarize_run(
     unsorted_task_doc = vasp_task_doc | base_task_doc | additional_fields
     task_doc = clean_task_doc(unsorted_task_doc)
 
+    if SETTINGS.WRITE_PICKLE:
+        with Path(dir_path, "quacc_results.pkl").open("wb") as f:
+            pickle.dump(task_doc, f)
+
     # Store the results
     if store:
         results_to_db(store, task_doc)

src/quacc/settings.py

@@ -117,6 +117,9 @@ class QuaccSettings(BaseSettings):
     GZIP_FILES: bool = Field(
         True, description="Whether generated files should be gzip'd."
     )
+    WRITE_PICKLE: bool = Field(
+        True, description="Whether the results schema should be written to a .pkl file."
+    )
     CHECK_CONVERGENCE: bool = Field(
         True,
         description="Whether to check for convergence, when implemented by a given recipe.",

tests/core/recipes/emt_recipes/test_emt_recipes.py

@@ -1,3 +1,6 @@
+import pickle
+from pathlib import Path
+
 import numpy as np
 import pytest
 from ase.build import bulk, molecule
@@ -27,6 +30,23 @@ def test_static_job(tmp_path, monkeypatch):
     assert output["parameters"]["asap_cutoff"] is True
     assert output["results"]["energy"] == pytest.approx(0.11074520235398744)
 
+    assert Path(output["dir_name"], "quacc_results.pkl").exists()
+
+    with open(Path(output["dir_name"], "quacc_results.pkl"), "rb") as f:
+        pickle_results = pickle.load(f)
+
+    output.pop("uuid")
+    assert pickle_results.keys() == output.keys()
+
+    assert pickle_results["nsites"] == output["nsites"]
+    assert (
+        pickle_results["parameters"]["asap_cutoff"]
+        == output["parameters"]["asap_cutoff"]
+    )
+    assert pickle_results["results"]["energy"] == output["results"]["energy"]
+    assert pickle_results["atoms"].info["test"] == output["atoms"].info["test"]
+    assert pickle_results["atoms"].info.get("_id") == output["atoms"].info.get("_id")
+
 
 def test_relax_job(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
@@ -82,6 +102,25 @@ def test_relax_job(tmp_path, monkeypatch):
     assert output["parameters"]["asap_cutoff"] is True
     assert output["results"]["energy"] == pytest.approx(0.04996032884581858)
 
+    with open(Path(output["dir_name"], "quacc_results.pkl"), "rb") as f:
+        pickle_results = pickle.load(f)
+
+    output.pop("uuid")
+
+    assert pickle_results.keys() == output.keys()
+
+    # assert things on the trajectory are the same
+    assert pickle_results["trajectory"][0] == output["trajectory"][0]
+    assert pickle_results["trajectory"][-1] == output["trajectory"][-1]
+    assert (
+        pickle_results["trajectory_results"][0]["energy"]
+        == output["trajectory_results"][0]["energy"]
+    )
+    assert (
+        pickle_results["trajectory_results"][-1]["energy"]
+        == output["trajectory_results"][-1]["energy"]
+    )
+
 
 def test_slab_dynamic_jobs(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
@@ -125,6 +164,8 @@ def test_customizer_v2():
     for result in results:
         assert result["parameters"]["asap_cutoff"] is False
 
+        assert Path(result["dir_name"], "quacc_results.pkl").exists()
+
 
 def test_all_customizers():
     atoms = bulk("Cu")

tests/core/recipes/lj_recipes/test_lj_recipes.py

@@ -1,3 +1,6 @@
+import pickle
+from pathlib import Path
+
 import numpy as np
 import pytest
 from ase.build import bulk, molecule
@@ -94,6 +97,27 @@ def test_freq_job(tmp_path, monkeypatch):
     assert len(output["parameters_thermo"]["vib_freqs"]) == 3 * len(atoms) - 6
     assert output["parameters_thermo"]["n_imag"] == 0
 
+    with open(Path(output["dir_name"], "quacc_results.pkl"), "rb") as f:
+        pickle_results = pickle.load(f)
+
+    output.pop("uuid")
+    assert pickle_results.keys() == output.keys()
+
+    # assert things on thermo and freq results
+    assert (
+        pickle_results["parameters_thermo"]["n_imag"]
+        == output["parameters_thermo"]["n_imag"]
+    )
+    assert (
+        pickle_results["parameters_thermo"]["vib_freqs"]
+        == output["parameters_thermo"]["vib_freqs"]
+    )
+    assert pickle_results["results"]["vib_freqs"] == output["results"]["vib_freqs"]
+    assert (
+        pickle_results["results"]["vib_freqs_raw"] == output["results"]["vib_freqs_raw"]
+    )
+    assert pickle_results["results"]["vib_freqs"] == output["results"]["vib_freqs"]
+
 
 def test_freq_job_threads(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)