Improve phonon handling (#1138)

CHANGELOG.md

@@ -6,6 +6,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [0.3.12]
 
+### Added
+
+- Added a basic phonon runner in `quacc.runners.calc.phonons`
+- Added a TBLite phonon recipe `phonon_job`
+- Added more type hints for input keyword arguments
+- Added the `quacc info` command-line argument for error reporting
+
 ### Changed
 
 - Consolidated VASP INCAR co-pilot settings

src/quacc/calculators/qchem/custodian.py

@@ -50,7 +50,7 @@ def run_custodian(
  -------
  None
  """
- # Adapted from atomate.qchem.firetasks.run_calc
+ # Adapted from atomate.qchem.firetasks.run_ase_calc
  from custodian import Custodian
  from custodian.qchem.handlers import QChemErrorHandler
  from custodian.qchem.jobs import QCJob

src/quacc/recipes/dftb/core.py

@@ -6,7 +6,7 @@
 from ase.calculators.dftb import Dftb
 
 from quacc import SETTINGS, job
-from quacc.runners.calc import run_calc
+from quacc.runners.calc import run_ase_calc
 from quacc.schemas.ase import summarize_run
 from quacc.utils.dicts import merge_dicts
 from quacc.utils.files import check_logfile
@@ -188,7 +188,7 @@ def _base_job(
  flags = merge_dicts(defaults, calc_swaps)
 
  atoms.calc = Dftb(**flags)
- final_atoms = run_calc(atoms, geom_file=GEOM_FILE, copy_files=copy_files)
+ final_atoms = run_ase_calc(atoms, geom_file=GEOM_FILE, copy_files=copy_files)
 
  if SETTINGS.CHECK_CONVERGENCE:
  if check_logfile(LOG_FILE, "SCC is NOT converged"):

src/quacc/recipes/emt/core.py

@@ -11,7 +11,7 @@
 from ase.optimize import FIRE
 
 from quacc import job
-from quacc.runners.calc import run_ase_opt, run_calc
+from quacc.runners.calc import run_ase_calc, run_ase_opt
 from quacc.schemas.ase import summarize_opt_run, summarize_run
 from quacc.utils.dicts import merge_dicts
 
@@ -57,7 +57,7 @@ def static_job(
  calc_swaps = calc_swaps or {}
 
  atoms.calc = EMT(**calc_swaps)
- final_atoms = run_calc(atoms, copy_files=copy_files)
+ final_atoms = run_ase_calc(atoms, copy_files=copy_files)
 
  return summarize_run(
  final_atoms,

src/quacc/recipes/gaussian/core.py

@@ -7,7 +7,7 @@
 from ase.calculators.gaussian import Gaussian
 
 from quacc import SETTINGS, job
-from quacc.runners.calc import run_calc
+from quacc.runners.calc import run_ase_calc
 from quacc.schemas.cclib import cclib_summarize_run
 from quacc.utils.dicts import merge_dicts
 
@@ -230,6 +230,6 @@ def _base_job(
  flags = merge_dicts(defaults, calc_swaps)
 
  atoms.calc = Gaussian(command=GAUSSIAN_CMD, **flags)
- atoms = run_calc(atoms, geom_file=LOG_FILE, copy_files=copy_files)
+ atoms = run_ase_calc(atoms, geom_file=LOG_FILE, copy_files=copy_files)
 
  return cclib_summarize_run(atoms, LOG_FILE, additional_fields=additional_fields)

src/quacc/recipes/gulp/core.py

@@ -8,7 +8,7 @@
 from ase.calculators.gulp import GULP
 
 from quacc import SETTINGS, job
-from quacc.runners.calc import run_calc
+from quacc.runners.calc import run_ase_calc
 from quacc.schemas.ase import summarize_run
 from quacc.utils.dicts import merge_dicts
 
@@ -246,7 +246,7 @@ def _base_job(
  options=gulp_options,
  library=library,
  )
- final_atoms = run_calc(
+ final_atoms = run_ase_calc(
  atoms,
  geom_file=GEOM_FILE_PBC if atoms.pbc.any() else GEOM_FILE_NOPBC,
  copy_files=copy_files,

src/quacc/recipes/lj/core.py

@@ -12,7 +12,7 @@
 
 from quacc import job
 from quacc.builders.thermo import build_ideal_gas
-from quacc.runners.calc import run_ase_opt, run_ase_vib, run_calc
+from quacc.runners.calc import run_ase_calc, run_ase_opt, run_ase_vib
 from quacc.schemas.ase import summarize_opt_run, summarize_run, summarize_vib_and_thermo
 from quacc.utils.dicts import merge_dicts
 
@@ -59,7 +59,7 @@ def static_job(
  calc_swaps = calc_swaps or {}
 
  atoms.calc = LennardJones(**calc_swaps)
- final_atoms = run_calc(atoms, copy_files=copy_files)
+ final_atoms = run_ase_calc(atoms, copy_files=copy_files)
 
  return summarize_run(
  final_atoms, input_atoms=atoms, additional_fields={"name": "LJ Static"}

src/quacc/recipes/newtonnet/core.py

@@ -9,7 +9,7 @@
 
 from quacc import SETTINGS, job
 from quacc.builders.thermo import build_ideal_gas
-from quacc.runners.calc import run_ase_opt, run_calc
+from quacc.runners.calc import run_ase_calc, run_ase_opt
 from quacc.schemas.ase import (
  summarize_ideal_gas_thermo,
  summarize_opt_run,
@@ -83,7 +83,7 @@ def static_job(
  flags = merge_dicts(defaults, calc_swaps)
 
  atoms.calc = NewtonNet(**flags)
- final_atoms = run_calc(atoms, copy_files=copy_files)
+ final_atoms = run_ase_calc(atoms, copy_files=copy_files)
 
  return summarize_run(
  final_atoms,
@@ -206,7 +206,7 @@ def freq_job(
 
  ml_calculator = NewtonNet(**flags)
  atoms.calc = ml_calculator
- final_atoms = run_calc(atoms, copy_files=copy_files)
+ final_atoms = run_ase_calc(atoms, copy_files=copy_files)
 
  summary = summarize_run(
  final_atoms,
@@ -261,7 +261,7 @@ def _add_stdev_and_hess(summary: dict[str, Any]) -> dict[str, Any]:
  )
  atoms = conf["atoms"]
  atoms.calc = ml_calculator
- results = run_calc(atoms).calc.results
+ results = run_ase_calc(atoms).calc.results
  conf["hessian"] = results["hessian"]
  conf["energy_std"] = results["energy_disagreement"]
  conf["forces_std"] = results["forces_disagreement"]

src/quacc/recipes/orca/core.py

@@ -8,7 +8,7 @@
 from ase.calculators.orca import ORCA, OrcaProfile
 
 from quacc import SETTINGS, job
-from quacc.runners.calc import run_calc
+from quacc.runners.calc import run_ase_calc
 from quacc.schemas.cclib import cclib_summarize_run
 from quacc.utils.dicts import merge_dicts
 
@@ -275,7 +275,7 @@ def _base_job(
  orcasimpleinput=orcasimpleinput,
  orcablocks=orcablocks,
  )
- atoms = run_calc(atoms, geom_file=GEOM_FILE, copy_files=copy_files)
+ atoms = run_ase_calc(atoms, geom_file=GEOM_FILE, copy_files=copy_files)
 
  return cclib_summarize_run(
  atoms,

src/quacc/recipes/psi4/core.py

@@ -7,7 +7,7 @@
 from monty.dev import requires
 
 from quacc import job
-from quacc.runners.calc import run_calc
+from quacc.runners.calc import run_ase_calc
 from quacc.schemas.ase import summarize_run
 from quacc.utils.dicts import merge_dicts
 
@@ -136,7 +136,7 @@ def _base_job(
  flags = merge_dicts(defaults, calc_swaps)
 
  atoms.calc = Psi4(**flags)
- final_atoms = run_calc(atoms, copy_files=copy_files)
+ final_atoms = run_ase_calc(atoms, copy_files=copy_files)
 
  return summarize_run(
  final_atoms,

src/quacc/recipes/qchem/core.py

@@ -8,7 +8,7 @@
 
 from quacc import SETTINGS, job
 from quacc.calculators.qchem import QChem
-from quacc.runners.calc import run_ase_opt, run_calc
+from quacc.runners.calc import run_ase_calc, run_ase_opt
 from quacc.schemas.ase import summarize_opt_run, summarize_run
 from quacc.utils.dicts import merge_dicts, remove_dict_nones
 
@@ -499,7 +499,7 @@ def _base_job(
  qchem_flags = remove_dict_nones(defaults)
 
  atoms.calc = QChem(atoms, **qchem_flags)
- final_atoms = run_calc(atoms, copy_files=copy_files)
+ final_atoms = run_ase_calc(atoms, copy_files=copy_files)
 
  return summarize_run(
  final_atoms,

src/quacc/recipes/tblite/core.py

@@ -8,7 +8,7 @@
 
 from quacc import job
 from quacc.builders.thermo import build_ideal_gas
-from quacc.runners.calc import run_ase_opt, run_ase_phonons, run_ase_vib, run_calc
+from quacc.runners.calc import run_ase_calc, run_ase_opt, run_ase_phonons, run_ase_vib
 from quacc.schemas.ase import (
  summarize_opt_run,
  summarize_phonon_run,
@@ -76,7 +76,7 @@ def static_job(
  flags = merge_dicts(defaults, calc_swaps)
  atoms.calc = TBLite(**flags)
 
- final_atoms = run_calc(atoms, copy_files=copy_files)
+ final_atoms = run_ase_calc(atoms, copy_files=copy_files)
  return summarize_run(
  final_atoms,
  input_atoms=atoms,

src/quacc/recipes/vasp/core.py

@@ -5,7 +5,7 @@
 
 from quacc import job
 from quacc.calculators.vasp import Vasp
-from quacc.runners.calc import run_calc
+from quacc.runners.calc import run_ase_calc
 from quacc.schemas.vasp import vasp_summarize_run
 from quacc.utils.dicts import merge_dicts
 
@@ -250,6 +250,6 @@ def _base_job(
  flags = merge_dicts(defaults, calc_swaps, remove_nones=False)
 
  atoms.calc = Vasp(atoms, preset=preset, **flags)
- atoms = run_calc(atoms, copy_files=copy_files)
+ atoms = run_ase_calc(atoms, copy_files=copy_files)
 
  return vasp_summarize_run(atoms, additional_fields=additional_fields)

src/quacc/runners/calc.py

@@ -56,7 +56,7 @@ class PhononReadKwargs(TypedDict, total=False):
  born: bool # default = False
 
 
-def run_calc(
+def run_ase_calc(
  atoms: Atoms, geom_file: str | None = None, copy_files: list[str] | None = None
 ) -> Atoms:
  """

src/quacc/schemas/_aliases/ase.py

@@ -1,6 +1,8 @@
 """Aliases for type hinting `quacc.schemas.ase`"""
 from typing import Any, TypedDict
 
+from numpy.typing import NDArray
+
 from quacc.schemas._aliases.atoms import AtomsSchema
 
 results = dict[str, Any] # from atoms.calc.results
@@ -55,9 +57,11 @@ class VibSchema(AtomsSchema):
  results: VibResults
 
 
-class PhononSchema(AtomsSchema):
+class PhononSchema(RunSchema):
  """Schema for `quacc.schemas.ase.summarize_phonon_run`"""
 
+ force_constant: NDArray
+
 
 class ParametersThermo(TypedDict):
  temperature: float

src/quacc/schemas/ase.py

@@ -13,7 +13,7 @@
 from quacc import SETTINGS
 from quacc.runners.prep import prep_next_run as prep_next_run_
 from quacc.schemas.atoms import atoms_to_metadata
-from quacc.utils.dicts import merge_dicts, remove_dict_nones, sort_dict
+from quacc.utils.dicts import recursive_merge_dicts, sort_dict
 from quacc.utils.files import get_uri
 from quacc.wflow.db import results_to_db
 
@@ -105,8 +105,10 @@ def summarize_run(
  atoms_to_store, charge_and_multiplicity=charge_and_multiplicity
  )
 
- unsorted_task_doc = final_atoms_metadata | inputs | results | additional_fields
- task_doc = sort_dict(remove_dict_nones(unsorted_task_doc))
+ unsorted_task_doc = recursive_merge_dicts(
+ final_atoms_metadata, inputs, results, additional_fields
+ )
+ task_doc = sort_dict(unsorted_task_doc)
 
  if store:
  results_to_db(store, task_doc)
@@ -201,8 +203,10 @@ def summarize_opt_run(
  }
 
  # Create a dictionary of the inputs/outputs
- unsorted_task_doc = base_task_doc | opt_fields | additional_fields
- task_doc = sort_dict(remove_dict_nones(unsorted_task_doc))
+ unsorted_task_doc = recursive_merge_dicts(
+ base_task_doc, opt_fields, additional_fields
+ )
+ task_doc = sort_dict(unsorted_task_doc)
 
  if store:
  results_to_db(store, task_doc)
@@ -312,8 +316,10 @@ def summarize_vib_run(
  }
  }
 
- unsorted_task_doc = atoms_metadata | inputs | results | additional_fields
- task_doc = sort_dict(remove_dict_nones(unsorted_task_doc))
+ unsorted_task_doc = recursive_merge_dicts(
+ atoms_metadata, inputs, results, additional_fields
+ )
+ task_doc = sort_dict(unsorted_task_doc)
 
  if store:
  results_to_db(store, task_doc)
@@ -352,15 +358,22 @@ def summarize_phonon_run(
  additional_fields = additional_fields or {}
  store = SETTINGS.PRIMARY_STORE if store is None else store
 
- # do stuff with phonons
+ # Create an atoms object with calculator results
  atoms = phonons.atoms
- atoms_metadata = atoms_to_metadata(
- atoms, charge_and_multiplicity=charge_and_multiplicity
+ atoms.calc = phonons.calc
+
+ # Base task doc
+ base_task_doc = summarize_run(
+ atoms,
+ charge_and_multiplicity=charge_and_multiplicity,
+ store=False,
  )
 
- # store more stuff
- unsorted_task_doc = atoms_metadata | additional_fields
- task_doc = sort_dict(remove_dict_nones(unsorted_task_doc))
+ # TODO: Store more specific properties from the phonons object
+ results = {"results": {"force_constant": phonons.get_force_constant()}}
+
+ unsorted_task_doc = recursive_merge_dicts(base_task_doc, results, additional_fields)
+ task_doc = sort_dict(unsorted_task_doc)
 
  if store:
  results_to_db(store, task_doc)
@@ -443,8 +456,10 @@ def summarize_ideal_gas_thermo(
  igt.atoms, charge_and_multiplicity=charge_and_multiplicity
  )
 
- unsorted_task_doc = atoms_metadata | inputs | results | additional_fields
- task_doc = sort_dict(remove_dict_nones(unsorted_task_doc))
+ unsorted_task_doc = recursive_merge_dicts(
+ atoms_metadata, inputs, results, additional_fields
+ )
+ task_doc = sort_dict(unsorted_task_doc)
 
  if store:
  results_to_db(store, task_doc)
@@ -505,8 +520,10 @@ def summarize_vib_and_thermo(
  store=False,
  )
 
- unsorted_task_doc = merge_dicts(vib_task_doc, thermo_task_doc) | additional_fields
- task_doc = sort_dict(remove_dict_nones(unsorted_task_doc))
+ unsorted_task_doc = recursive_merge_dicts(
+ vib_task_doc, thermo_task_doc, additional_fields
+ )
+ task_doc = sort_dict(unsorted_task_doc)
 
  if store:
  results_to_db(store, task_doc)