Skip to content

Commit

Permalink
Clean up schemas (#2022)
Browse files Browse the repository at this point in the history
## Summary of Changes

Some minor internal refactor.

### 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 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: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
Andrew-S-Rosen and pre-commit-ci[bot] authored Apr 21, 2024
1 parent 352d3cc commit 568f93f
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 40 deletions.
40 changes: 7 additions & 33 deletions src/quacc/schemas/ase.py
Original file line number Diff line number Diff line change
Expand Up @@ -279,19 +279,17 @@ def summarize_vib_and_thermo(
vib_task_doc = _summarize_vib_run(
vib,
charge_and_multiplicity=charge_and_multiplicity,
store=None,
additional_fields=additional_fields,
)
thermo_task_doc = _summarize_ideal_gas_thermo(
igt,
temperature=temperature,
pressure=pressure,
charge_and_multiplicity=charge_and_multiplicity,
store=None,
additional_fields=additional_fields,
)

unsorted_task_doc = recursive_dict_merge(vib_task_doc, thermo_task_doc)
unsorted_task_doc = recursive_dict_merge(
vib_task_doc, thermo_task_doc, additional_fields
)
task_doc = clean_task_doc(unsorted_task_doc)

if isinstance(vib, Vibrations):
Expand All @@ -312,8 +310,6 @@ def summarize_vib_and_thermo(
def _summarize_vib_run(
vib: Vibrations | VibrationsData,
charge_and_multiplicity: tuple[int, int] | None = None,
additional_fields: dict[str, Any] | None = None,
store: Store | None = None,
) -> VibSchema:
"""
Get tabulated results from an ASE Vibrations object and store them in a database-
Expand All @@ -326,19 +322,12 @@ def _summarize_vib_run(
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. Defaults to `SETTINGS.STORE`.
Returns
-------
VibSchema
Dictionary representation of the task document
"""
additional_fields = additional_fields or {}
store = SETTINGS.STORE if store == _DEFAULT_SETTING else store

vib_freqs_raw = vib.get_frequencies().tolist()
vib_energies_raw = vib.get_energies().tolist()

Expand Down Expand Up @@ -411,22 +400,16 @@ def _summarize_vib_run(
}
}

unsorted_task_doc = atoms_metadata | inputs | results | additional_fields
task_doc = clean_task_doc(unsorted_task_doc)

if store:
results_to_db(store, task_doc)
unsorted_task_doc = atoms_metadata | inputs | results

return task_doc
return clean_task_doc(unsorted_task_doc)


def _summarize_ideal_gas_thermo(
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 | None = _DEFAULT_SETTING,
) -> ThermoSchema:
"""
Get tabulated results from an ASE IdealGasThermo object and store them in a
Expand All @@ -445,17 +428,12 @@ def _summarize_ideal_gas_thermo(
metadata.
additional_fields
Additional fields to add to the task document.
store
Maggma Store object to store the results in. Defaults to `SETTINGS.STORE`
Returns
-------
ThermoSchema
Dictionary representation of the task document
"""
additional_fields = additional_fields or {}
store = SETTINGS.STORE if store == _DEFAULT_SETTING else store

spin_multiplicity = round(2 * igt.spin + 1)

inputs = {
Expand Down Expand Up @@ -494,10 +472,6 @@ def _summarize_ideal_gas_thermo(
igt.atoms, charge_and_multiplicity=charge_and_multiplicity
)

unsorted_task_doc = atoms_metadata | inputs | results | additional_fields
task_doc = clean_task_doc(unsorted_task_doc)
unsorted_task_doc = atoms_metadata | inputs | results

if store:
results_to_db(store, task_doc)

return task_doc
return clean_task_doc(unsorted_task_doc)
7 changes: 0 additions & 7 deletions tests/core/schemas/test_ase.py
Original file line number Diff line number Diff line change
Expand Up @@ -234,10 +234,6 @@ def test_summarize_vib_run(tmp_path, monkeypatch):
vib = Vibrations(atoms)
vib.run()

store = MemoryStore()
_summarize_vib_run(vib, store=store)
assert store.count() == 1

# Make sure info tags are handled appropriately
atoms = molecule("N2")
atoms.info["test_dict"] = {"hi": "there", "foo": "bar"}
Expand Down Expand Up @@ -287,9 +283,6 @@ def test_summarize_ideal_gas_thermo(tmp_path, monkeypatch):
atoms = molecule("N2")
igt = IdealGasThermo([0.34], "linear", atoms=atoms, spin=0, symmetrynumber=2)
_summarize_ideal_gas_thermo(igt)
store = MemoryStore()
_summarize_ideal_gas_thermo(igt, store=store)
assert store.count() == 1

# Make sure right number of vib energies are reported
atoms = molecule("N2")
Expand Down

0 comments on commit 568f93f

Please sign in to comment.