Repository: Quantum-Accelerators/quacc
Pull request #1831 (2024-03-04 05:25:20)

Refactor schemas


Patch:
diff --git a/docs/gen_ref_pages.py b/docs/gen_ref_pages.py
index f243a8d8..8dcd3b09 100644
--- a/docs/gen_ref_pages.py
+++ b/docs/gen_ref_pages.py
@@ -26,7 +26,7 @@ for path in sorted(Path("src").rglob("*.py")):
         ident = ".".join(parts)
         fd.write(f"::: {ident}")
 
-    mkdocs_gen_files.set_edit_path(full_doc_path, Path("../") / path)
+    mkdocs_gen_files.set_edit_path(full_doc_path, Path("..") / path)
 
 with mkdocs_gen_files.open(Path("reference", "SUMMARY.md"), "w") as nav_file:
     nav_file.writelines(nav.build_literate_nav())
diff --git a/src/quacc/schemas/_aliases/ase.py b/src/quacc/schemas/_aliases/ase.py
index 4f5e0dd5..728f017e 100644
--- a/src/quacc/schemas/_aliases/ase.py
+++ b/src/quacc/schemas/_aliases/ase.py
@@ -53,8 +53,6 @@ class VibResults(TypedDict):
 
 
 class VibSchema(AtomsSchema):
-    """Schema for [quacc.schemas.ase.summarize_vib_run][]"""
-
     parameters: Parameters | None
     parameters_vib: ParametersVib | None
     results: VibResults
diff --git a/src/quacc/schemas/ase.py b/src/quacc/schemas/ase.py
index c6ab70b8..dc828211 100644
--- a/src/quacc/schemas/ase.py
+++ b/src/quacc/schemas/ase.py
@@ -223,7 +223,69 @@ def summarize_opt_run(
     return task_doc
 
 
-def summarize_vib_run(
+def summarize_vib_and_thermo(
+    vib: Vibrations,
+    igt: IdealGasThermo,
+    temperature: float = 298.15,
+    pressure: float = 1.0,
+    charge_and_multiplicity: tuple[int, int] | None = None,
+    additional_fields: dict[str, Any] | None = None,
+    store: Store | bool | None = None,
+) -> VibThermoSchema:
+    """
+    Get tabulated results from an ASE Vibrations run and ASE IdealGasThermo object and
+    store them in a database-friendly format.
+
+    Parameters
+    ----------
+    vib
+        ASE Vibrations object.
+    igt
+        ASE IdealGasThermo object.
+    temperature
+        Temperature in Kelvins.
+    pressure
+        Pressure in bar.
+    charge_and_multiplicity
+        Charge and spin multiplicity of the Atoms object, only used for Molecule
+        metadata.
+    additional_fields
+        Additional fields to add to the task document.
+    store
+        Maggma Store object to store the results in. If None,
+        `SETTINGS.STORE` will be used.
+
+    Returns
+    -------
+    VibThermoSchema
+        A dictionary that merges the `VibSchema` and `ThermoSchema`.
+    """
+    additional_fields = additional_fields or {}
+    store = SETTINGS.STORE if store is None else store
+
+    vib_task_doc = _summarize_vib_run(
+        vib, charge_and_multiplicity=charge_and_multiplicity, store=False
+    )
+    thermo_task_doc = _summarize_ideal_gas_thermo(
+        igt,
+        temperature=temperature,
+        pressure=pressure,
+        charge_and_multiplicity=charge_and_multiplicity,
+        store=False,
+    )
+
+    unsorted_task_doc = recursive_dict_merge(
+        vib_task_doc, thermo_task_doc, additional_fields
+    )
+    task_doc = clean_task_doc(unsorted_task_doc)
+
+    if store:
+        results_to_db(store, task_doc)
+
+    return task_doc
+
+
+def _summarize_vib_run(
     vib: Vibrations | VibrationsData,
     charge_and_multiplicity: tuple[int, int] | None = None,
     additional_fields: dict[str, Any] | None = None,
@@ -337,7 +399,7 @@ def summarize_vib_run(
     return task_doc
 
 
-def summarize_ideal_gas_thermo(
+def _summarize_ideal_gas_thermo(
     igt: IdealGasThermo,
     temperature: float = 298.15,
     pressure: float = 1.0,
@@ -422,65 +484,3 @@ def summarize_ideal_gas_thermo(
         results_to_db(store, task_doc)
 
     return task_doc
-
-
-def summarize_vib_and_thermo(
-    vib: Vibrations,
-    igt: IdealGasThermo,
-    temperature: float = 298.15,
-    pressure: float = 1.0,
-    charge_and_multiplicity: tuple[int, int] | None = None,
-    additional_fields: dict[str, Any] | None = None,
-    store: Store | bool | None = None,
-) -> VibThermoSchema:
-    """
-    Get tabulated results from an ASE Vibrations run and ASE IdealGasThermo object and
-    store them in a database-friendly format.
-
-    Parameters
-    ----------
-    vib
-        ASE Vibrations object.
-    igt
-        ASE IdealGasThermo object.
-    temperature
-        Temperature in Kelvins.
-    pressure
-        Pressure in bar.
-    charge_and_multiplicity
-        Charge and spin multiplicity of the Atoms object, only used for Molecule
-        metadata.
-    additional_fields
-        Additional fields to add to the task document.
-    store
-        Maggma Store object to store the results in. If None,
-        `SETTINGS.STORE` will be used.
-
-    Returns
-    -------
-    VibThermoSchema
-        A dictionary that merges the `VibSchema` and `ThermoSchema`.
-    """
-    additional_fields = additional_fields or {}
-    store = SETTINGS.STORE if store is None else store
-
-    vib_task_doc = summarize_vib_run(
-        vib, charge_and_multiplicity=charge_and_multiplicity, store=False
-    )
-    thermo_task_doc = summarize_ideal_gas_thermo(
-        igt,
-        temperature=temperature,
-        pressure=pressure,
-        charge_and_multiplicity=charge_and_multiplicity,
-        store=False,
-    )
-
-    unsorted_task_doc = recursive_dict_merge(
-        vib_task_doc, thermo_task_doc, additional_fields
-    )
-    task_doc = clean_task_doc(unsorted_task_doc)
-
-    if store:
-        results_to_db(store, task_doc)
-
-    return task_doc


Test patch:
diff --git a/tests/core/schemas/test_ase.py b/tests/core/schemas/test_ase.py
index 7fe05dae..10e3de0a 100644
--- a/tests/core/schemas/test_ase.py
+++ b/tests/core/schemas/test_ase.py
@@ -14,10 +14,10 @@ from maggma.stores import MemoryStore
 from monty.json import MontyDecoder, jsanitize
 
 from quacc.schemas.ase import (
-    summarize_ideal_gas_thermo,
+    _summarize_ideal_gas_thermo,
+    _summarize_vib_run,
     summarize_opt_run,
     summarize_run,
-    summarize_vib_run,
 )
 
 FILE_DIR = Path(__file__).parent
@@ -175,7 +175,7 @@ def test_summarize_vib_run(tmp_path, monkeypatch):
     vib = Vibrations(atoms)
     vib.run()
 
-    results = summarize_vib_run(vib)
+    results = _summarize_vib_run(vib)
     assert results["atoms"] == input_atoms
     assert results["natoms"] == len(atoms)
     assert results["parameters_vib"]["delta"] == vib.delta
@@ -209,7 +209,7 @@ def test_summarize_vib_run(tmp_path, monkeypatch):
     vib.run()
 
     store = MemoryStore()
-    summarize_vib_run(vib, store=store)
+    _summarize_vib_run(vib, store=store)
     assert store.count() == 1
 
     # Make sure info tags are handled appropriately
@@ -219,7 +219,7 @@ def test_summarize_vib_run(tmp_path, monkeypatch):
     vib = Vibrations(atoms)
     vib.run()
 
-    results = summarize_vib_run(vib)
+    results = _summarize_vib_run(vib)
     assert results.get("atoms_info", {}) != {}
     assert results["atoms_info"].get("test_dict", None) == {"hi": "there", "foo": "bar"}
     assert results["atoms"].info.get("test_dict", None) == {"hi": "there", "foo": "bar"}
@@ -235,7 +235,7 @@ def test_summarize_vib_run(tmp_path, monkeypatch):
     vib = Vibrations(atoms)
     vib.run()
 
-    results = summarize_vib_run(vib)
+    results = _summarize_vib_run(vib)
     assert results["atoms"] == input_atoms
     assert results["nsites"] == len(atoms)
     assert results["parameters_vib"]["delta"] == vib.delta
@@ -251,7 +251,7 @@ def test_summarize_ideal_gas_thermo(tmp_path, monkeypatch):
     # Make sure metadata is made
     atoms = molecule("N2")
     igt = IdealGasThermo([0.34], "linear", atoms=atoms, spin=0, symmetrynumber=2)
-    results = summarize_ideal_gas_thermo(igt)
+    results = _summarize_ideal_gas_thermo(igt)
     assert results["natoms"] == len(atoms)
     assert results["atoms"] == atoms
     assert results["parameters_thermo"]["vib_energies"] == [0.34]
@@ -262,9 +262,9 @@ def test_summarize_ideal_gas_thermo(tmp_path, monkeypatch):
     # Test DB
     atoms = molecule("N2")
     igt = IdealGasThermo([0.34], "linear", atoms=atoms, spin=0, symmetrynumber=2)
-    summarize_ideal_gas_thermo(igt)
+    _summarize_ideal_gas_thermo(igt)
     store = MemoryStore()
-    summarize_ideal_gas_thermo(igt, store=store)
+    _summarize_ideal_gas_thermo(igt, store=store)
     assert store.count() == 1
 
     # Make sure right number of vib energies are reported
@@ -272,7 +272,7 @@ def test_summarize_ideal_gas_thermo(tmp_path, monkeypatch):
     igt = IdealGasThermo(
         [0.0, 0.34], "linear", atoms=atoms, potentialenergy=-1, spin=0, symmetrynumber=2
     )
-    results = summarize_ideal_gas_thermo(igt)
+    results = _summarize_ideal_gas_thermo(igt)
     assert results["natoms"] == len(atoms)
     assert results["atoms"] == atoms
     assert results["parameters_thermo"]["vib_energies"] == [0.34]
@@ -286,7 +286,7 @@ def test_summarize_ideal_gas_thermo(tmp_path, monkeypatch):
     igt = IdealGasThermo(
         [0.0, 0.34], "linear", atoms=atoms, potentialenergy=-1, spin=0, symmetrynumber=2
     )
-    results = summarize_ideal_gas_thermo(igt)
+    results = _summarize_ideal_gas_thermo(igt)
     assert results.get("atoms_info", {}) != {}
     assert results["atoms_info"].get("test_dict", None) == {"hi": "there", "foo": "bar"}
     assert results["atoms"].info.get("test_dict", None) == {"hi": "there", "foo": "bar"}
@@ -315,7 +315,7 @@ def test_summarize_ideal_gas_thermo(tmp_path, monkeypatch):
         spin=0.5,
         symmetrynumber=6,
     )
-    results = summarize_ideal_gas_thermo(igt, temperature=1000.0, pressure=20.0)
+    results = _summarize_ideal_gas_thermo(igt, temperature=1000.0, pressure=20.0)
     assert results["natoms"] == len(atoms)
     assert results["atoms"] == atoms
     assert len(results["parameters_thermo"]["vib_energies"]) == 6
@@ -335,7 +335,7 @@ def test_summarize_ideal_gas_thermo(tmp_path, monkeypatch):
     MontyDecoder().process_decoded(d)
 
     with pytest.raises(ValueError):
-        summarize_ideal_gas_thermo(igt, charge_and_multiplicity=[0, 1])
+        _summarize_ideal_gas_thermo(igt, charge_and_multiplicity=[0, 1])
 
 
 def test_errors(tmp_path, monkeypatch):