Write out results schema to filesystem in summarize functions

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)