Easier handling of trajectory in opt schema (#1582)

Co-authored-by: deepsource-autofix[bot] <62050782+deepsource-autofix[bot]@users.noreply.github.com>

src/quacc/recipes/newtonnet/core.py

@@ -219,17 +219,18 @@ def _add_stdev_and_hess(summary: dict[str, Any]) -> dict[str, Any]:
  Hessian values.
  """
 
- for conf in summary["trajectory"]:
+ for i, atoms in enumerate(summary["trajectory"]):
  ml_calculator = NewtonNet(
  model_path=SETTINGS.NEWTONNET_MODEL_PATH,
  settings_path=SETTINGS.NEWTONNET_CONFIG_PATH,
  )
- atoms = conf["atoms"]
  atoms.calc = ml_calculator
  results = run_calc(atoms).calc.results
- conf["hessian"] = results["hessian"]
- conf["energy_std"] = results["energy_disagreement"]
- conf["forces_std"] = results["forces_disagreement"]
- conf["hessian_std"] = results["hessian_disagreement"]
+ summary["trajectory_results"][i]["hessian"] = results["hessian"]
+ summary["trajectory_results"][i]["energy_std"] = results["energy_disagreement"]
+ summary["trajectory_results"][i]["forces_std"] = results["forces_disagreement"]
+ summary["trajectory_results"][i]["hessian_std"] = results[
+ "hessian_disagreement"
+ ]
 
  return summary

src/quacc/schemas/_aliases/ase.py

@@ -3,6 +3,7 @@
 
 from typing import Any, TypedDict
 
+from ase.atoms import Atoms
 from numpy.typing import NDArray
 
 from quacc.schemas._aliases.atoms import AtomsSchema
@@ -29,7 +30,7 @@ class OptSchema(RunSchema):
  parameters_opt: dict[str, Any] # from Optimizer.todict()
  converged: bool
  nsteps: int
- trajectory: list[AtomsSchema]
+ trajectory: list[Atoms]
  trajectory_results: list[results]
 
 

src/quacc/schemas/_aliases/cclib.py

@@ -3,10 +3,10 @@
 
 from typing import Any, TypedDict
 
+from ase.atoms import Atoms
 from numpy.typing import NDArray
 
 from quacc.schemas._aliases.ase import OptSchema, RunSchema
-from quacc.schemas._aliases.atoms import AtomsSchema
 
 
 class AdditionalAttributes(TypedDict, total=False):
@@ -129,7 +129,7 @@ class cclibBaseSchema(TypedDict):
  logfile: str
  attributes: AllAttributes
  pop_analysis: PopAnalysisAttributes | None
- trajectory: list[AtomsSchema]
+ trajectory: list[Atoms]
 
 
 class cclibSchema(cclibBaseSchema, RunSchema):

src/quacc/schemas/ase.py

@@ -199,10 +199,7 @@ def summarize_opt_run(
  "parameters_opt": dyn.todict(),
  "converged": is_converged,
  "nsteps": dyn.get_number_of_steps(),
- "trajectory": [
- atoms_to_metadata(atoms, charge_and_multiplicity=charge_and_multiplicity)
- for atoms in trajectory
- ],
+ "trajectory": trajectory,
  "trajectory_results": [atoms.calc.results for atoms in trajectory],
  }
 

src/quacc/schemas/cclib.py

@@ -14,7 +14,6 @@
 
 from quacc import SETTINGS
 from quacc.schemas.ase import summarize_run
-from quacc.schemas.atoms import atoms_to_metadata
 from quacc.utils.dicts import clean_task_doc
 from quacc.utils.files import find_recent_logfile
 from quacc.wflow_tools.db import results_to_db
@@ -121,7 +120,7 @@ def cclib_summarize_run(
  positions = [row[1:] for row in coords_obj]
  input_atoms = Atoms(symbols=symbols, positions=positions)
  else:
- input_atoms = cclib_task_doc["trajectory"][0]["atoms"]
+ input_atoms = cclib_task_doc["trajectory"][0]
 
  # Get the base task document for the ASE run
  run_task_doc = summarize_run(
@@ -206,19 +205,11 @@ def _make_cclib_schema(
  if cpu_time := attributes["metadata"].get("cpu_time"):
  attributes["metadata"]["cpu_time"] = [*map(str, cpu_time)]
 
- # Store charge and multiplicity since we use it frequently
- charge = cclib_obj.charge
- mult = cclib_obj.mult
-
  # Construct the trajectory
  coords = cclib_obj.atomcoords
  trajectory = [
  Atoms(numbers=list(cclib_obj.atomnos), positions=coord) for coord in coords
  ]
- traj_metadata = [
- atoms_to_metadata(traj, charge_and_multiplicity=(charge, mult))
- for traj in trajectory
- ]
 
  # Get the final energy to store as its own key/value pair
  final_scf_energy = (
@@ -265,7 +256,7 @@ def _make_cclib_schema(
  "logfile": str(logfile).split(":")[-1],
  "attributes": attributes | additional_attributes,
  "pop_analysis": popanalysis_attributes or None,
- "trajectory": traj_metadata,
+ "trajectory": trajectory,
  }
 
 

tests/core/recipes/emt_recipes/test_emt_recipes.py

@@ -40,8 +40,8 @@ def test_relax_job(tmp_path, monkeypatch):
  assert np.max(np.linalg.norm(output["results"]["forces"], axis=1)) < 0.01
  assert len(output["trajectory"]) == 30
  assert output["atoms"] != output["input_atoms"]["atoms"]
- assert output["trajectory"][0]["atoms"] == output["input_atoms"]["atoms"]
- assert output["trajectory"][-1]["atoms"] == output["atoms"]
+ assert output["trajectory"][0] == output["input_atoms"]["atoms"]
+ assert output["trajectory"][-1] == output["atoms"]
  assert (
  output["trajectory_results"][0]["energy"]
  > output["trajectory_results"][-1]["energy"]

tests/core/recipes/newtonnet_recipes/test_newtonnet_recipes.py

@@ -134,11 +134,6 @@ def test_ts_job_with_custom_hessian(tmp_path, monkeypatch):
 
  # Perform assertions on the result
  assert isinstance(output, dict)
-
- # assert output["results"]["energy"] == pytest.approx(-8.855604432470276)
- # assert output["freq_job"]["vib"]["results"]["vib_energies"][0] == pytest.approx(
- # 0.2256022513686731
- # )
  assert "thermo" in output["freq_job"]
 
 

tests/core/schemas/test_cclib_schema.py

@@ -100,8 +100,8 @@ def test_cclib_summarize_run(tmp_path, monkeypatch):
  7.6760596087545006
  )
  assert len(results["trajectory"]) == 7
- assert results["trajectory"][0]["atoms"] == results["input_atoms"]["atoms"]
- assert results["trajectory"][-1]["atoms"] == results["atoms"]
+ assert results["trajectory"][0] == results["input_atoms"]["atoms"]
+ assert results["trajectory"][-1] == results["atoms"]
  assert results["test"] == "hi"
 
  # test document can be jsanitized and decoded