Repository: Quantum-Accelerators/quacc
Pull request #2235 (2024-06-11 23:25:28)

Get rid of `.traj_atoms` hack


Patch:
diff --git a/src/quacc/runners/ase.py b/src/quacc/runners/ase.py
index 347cc2c9..b3b67d49 100644
--- a/src/quacc/runners/ase.py
+++ b/src/quacc/runners/ase.py
@@ -218,6 +218,7 @@ class Runner:
             optimizer_kwargs,
         )
         run_kwargs = run_kwargs or {}
+        traj_filename = "opt.traj"
 
         # Check if trajectory kwarg is specified
         if "trajectory" in optimizer_kwargs:
@@ -233,7 +234,7 @@ class Runner:
             optimizer_kwargs.pop("restart", None)
 
         # Define the Trajectory object
-        traj_file = self.tmpdir / "opt.traj"
+        traj_file = self.tmpdir / traj_filename
         traj = Trajectory(traj_file, "w", atoms=self.atoms)
         optimizer_kwargs["trajectory"] = traj
 
@@ -265,11 +266,10 @@ class Runner:
         except Exception as exception:
             terminate(self.tmpdir, exception)
 
-        # Store the trajectory atoms
-        dyn.traj_atoms = read(traj_file, index=":")
-
         # Perform cleanup operations
         self._cleanup()
+        traj.filename = zpath(self.job_results_dir / traj_filename)
+        dyn.trajectory = traj
 
         return dyn
 
diff --git a/src/quacc/schemas/ase.py b/src/quacc/schemas/ase.py
index c046e0c7..f8caa30b 100644
--- a/src/quacc/schemas/ase.py
+++ b/src/quacc/schemas/ase.py
@@ -137,7 +137,7 @@ def summarize_opt_run(
         ASE Optimizer object.
     trajectory
         ASE Trajectory object or list[Atoms] from reading a trajectory file. If
-        None, the trajectory must be found in dyn.traj_atoms.
+        None, the trajectory must be found in `dyn.trajectory.filename`.
     check_convergence
         Whether to check the convergence of the calculation. Defaults to True in
         settings.
@@ -167,11 +167,7 @@ def summarize_opt_run(
 
     # Get trajectory
     if not trajectory:
-        trajectory = (
-            dyn.traj_atoms
-            if hasattr(dyn, "traj_atoms")
-            else read(dyn.trajectory.filename, index=":")
-        )
+        trajectory = read(dyn.trajectory.filename, index=":")
     trajectory_results = [atoms.calc.results for atoms in trajectory]
     for traj_atoms in trajectory:
         traj_atoms.calc = None
diff --git a/src/quacc/schemas/cclib.py b/src/quacc/schemas/cclib.py
index 2051b260..1da15a21 100644
--- a/src/quacc/schemas/cclib.py
+++ b/src/quacc/schemas/cclib.py
@@ -157,7 +157,7 @@ def cclib_summarize_run(
 
 
 def summarize_cclib_opt_run(
-    optimizer: Optimizer,
+    dyn: Optimizer,
     logfile_extensions: str | list[str],
     trajectory: Trajectory | list[Atoms] | None = None,
     directory: Path | str | None = None,
@@ -186,7 +186,7 @@ def summarize_cclib_opt_run(
 
     Parameters
     ----------
-    optimizer
+    dyn
         The ASE optimizer object
     logfile_extensions
         Possible extensions of the log file (e.g. ".log", ".out", ".txt",
@@ -196,7 +196,7 @@ def summarize_cclib_opt_run(
         used. For an exact match only, put in the full file name.
     trajectory
         ASE Trajectory object or list[Atoms] from reading a trajectory file. If
-        None, the trajectory must be found in dyn.traj_atoms.
+        None, the trajectory must be found in `dyn.trajectory.filename`.
     directory
         The path to the folder containing the calculation outputs. A value of
         None specifies the calculator directory.
@@ -220,7 +220,7 @@ def summarize_cclib_opt_run(
     """
     store = SETTINGS.STORE if store == _DEFAULT_SETTING else store
 
-    final_atoms = get_final_atoms_from_dynamics(optimizer)
+    final_atoms = get_final_atoms_from_dynamics(dyn)
     directory = Path(directory or final_atoms.calc.directory)
     cclib_summary = cclib_summarize_run(
         final_atoms,
@@ -232,7 +232,7 @@ def summarize_cclib_opt_run(
         store=None,
     )
     opt_run_summary = summarize_opt_run(
-        optimizer,
+        dyn,
         trajectory=trajectory,
         check_convergence=check_convergence,
         charge_and_multiplicity=(
diff --git a/src/quacc/schemas/vasp.py b/src/quacc/schemas/vasp.py
index b9a0004f..4a258c8b 100644
--- a/src/quacc/schemas/vasp.py
+++ b/src/quacc/schemas/vasp.py
@@ -193,7 +193,7 @@ def summarize_vasp_opt_run(
         The ASE optimizer object
     trajectory
         ASE Trajectory object or list[Atoms] from reading a trajectory file. If
-        None, the trajectory must be found in dyn.traj_atoms.
+        None, the trajectory must be found in dyn.trajectory.filename.
     directory
         Path to VASP outputs. A value of None specifies the calculator directory.
     move_magmoms


Test patch:
diff --git a/tests/core/runners/test_ase.py b/tests/core/runners/test_ase.py
index 9e501e26..c3f62442 100644
--- a/tests/core/runners/test_ase.py
+++ b/tests/core/runners/test_ase.py
@@ -11,6 +11,7 @@ import pytest
 from ase.build import bulk, molecule
 from ase.calculators.emt import EMT
 from ase.calculators.lj import LennardJones
+from ase.io import read
 from ase.optimize import BFGS, BFGSLineSearch
 from ase.optimize.sciopt import SciPyFminBFGS
 
@@ -97,7 +98,7 @@ def test_run_opt1(tmp_path, monkeypatch):
         atoms[0].position += 0.1
 
         dyn = Runner(atoms, EMT(), copy_files={Path(): "test_file.txt"}).run_opt()
-        traj = dyn.traj_atoms
+        traj = read(dyn.trajectory.filename, index=":")
         results_dir = _find_results_dir()
 
         assert atoms.calc is None
@@ -117,13 +118,13 @@ def test_run_opt2(tmp_path, monkeypatch):
     dyn = Runner(atoms, EMT(), copy_files={Path(): "test_file.txt"}).run_opt(
         optimizer=BFGS, optimizer_kwargs={"restart": None}
     )
-    traj = dyn.traj_atoms
+    traj = read(dyn.trajectory.filename, index=":")
     assert traj[-1].calc.results is not None
 
     dyn = Runner(traj[-1], EMT(), copy_files={Path(): "test_file.txt"}).run_opt(
         optimizer=BFGSLineSearch, optimizer_kwargs={"restart": None}
     )
-    traj = dyn.traj_atoms
+    traj = read(dyn.trajectory.filename, index=":")
     assert traj[-1].calc.results is not None
 
 
@@ -133,7 +134,7 @@ def test_run_scipy_opt(tmp_path, monkeypatch):
     atoms[0].position += 0.1
 
     dyn = Runner(atoms, EMT()).run_opt(optimizer=SciPyFminBFGS)
-    traj = dyn.traj_atoms
+    traj = read(dyn.trajectory.filename, index=":")
     assert traj[-1].calc.results is not None
     assert dyn.todict().get("restart") is None
 
diff --git a/tests/core/runners/test_sella.py b/tests/core/runners/test_sella.py
index 8f0dd853..cf6e388c 100644
--- a/tests/core/runners/test_sella.py
+++ b/tests/core/runners/test_sella.py
@@ -6,6 +6,7 @@ sella = pytest.importorskip("sella")
 from ase.build import bulk, molecule
 from ase.calculators.emt import EMT
 from ase.calculators.lj import LennardJones
+from ase.io import read
 from sella import Sella
 
 from quacc.runners.ase import Runner
@@ -19,7 +20,7 @@ def test_sella(tmp_path, monkeypatch):
     dyn = Runner(atoms, EMT()).run_opt(
         optimizer=Sella, optimizer_kwargs={"restart": None}
     )
-    traj = dyn.traj_atoms
+    traj = read(dyn.trajectory.filename, index=":")
     assert traj[-1].calc.results is not None
     assert dyn.user_internal is False
 
@@ -27,7 +28,7 @@ def test_sella(tmp_path, monkeypatch):
     dyn = Runner(atoms, LennardJones()).run_opt(
         optimizer=Sella, optimizer_kwargs={"restart": None}
     )
-    traj = dyn.traj_atoms
+    traj = read(dyn.trajectory.filename, index=":")
     assert traj[-1].calc.results is not None
     assert dyn.user_internal is True
 
@@ -37,7 +38,7 @@ def test_dof(tmp_path, monkeypatch):
 
     atoms = molecule("C2H6")
     dyn = Runner(atoms, LennardJones()).run_opt(optimizer=Sella, fmax=1.0)
-    traj = dyn.traj_atoms
+    traj = read(dyn.trajectory.filename, index=":")
     assert traj[-1].calc.results is not None
     assert dyn.internal.ndof == 24
     assert dyn.internal.nbonds == 0
@@ -46,6 +47,6 @@ def test_dof(tmp_path, monkeypatch):
 
     atoms = bulk("Cu")
     dyn = Runner(atoms, EMT()).run_opt(optimizer=Sella, fmax=1.0)
-    traj = dyn.traj_atoms
+    traj = read(dyn.trajectory.filename, index=":")
     assert traj[-1].calc.results is not None
     assert dyn.internal is None