Update Atoms info from calculated results (#221)

* Ensure atoms info is consistent with files written

* Add write-kwargs to MD

* Add arch to Atoms.info and default writing results

---------

Co-authored-by: Jacob Wilkins <[email protected]>

janus_core/calculations/eos.py

@@ -5,16 +5,15 @@
 
 from ase import Atoms
 from ase.eos import EquationOfState
-from ase.io import write
 from ase.units import kJ
 from codecarbon import OfflineEmissionsTracker
 from numpy import float64, linspace
 from numpy.typing import NDArray
 
 from janus_core.calculations.geom_opt import optimize
-from janus_core.helpers.janus_types import ASEWriteArgs, EoSNames, EoSResults, PathLike
+from janus_core.helpers.janus_types import EoSNames, EoSResults, OutputKwargs, PathLike
 from janus_core.helpers.log import config_logger, config_tracker
-from janus_core.helpers.utils import none_to_dict
+from janus_core.helpers.utils import none_to_dict, output_structs
 
 
 def _calc_volumes_energies(  # pylint: disable=too-many-locals
@@ -26,7 +25,7 @@ def _calc_volumes_energies(  # pylint: disable=too-many-locals
     minimize_all: bool = False,
     minimize_kwargs: Optional[dict[str, Any]] = None,
     write_structures: bool = False,
-    write_kwargs: Optional[ASEWriteArgs] = None,
+    write_kwargs: Optional[OutputKwargs] = None,
     logger: Optional[Logger] = None,
     tracker: Optional[OfflineEmissionsTracker] = None,
 ) -> tuple[NDArray[float64], list[float], list[float]]:
@@ -50,7 +49,7 @@ def _calc_volumes_energies(  # pylint: disable=too-many-locals
         chemical formula of the structure.
     write_structures : bool
         True to write out all genereated structures. Default is False.
-    write_kwargs : Optional[ASEWriteArgs],
+    write_kwargs : Optional[OutputKwargs],
         Keyword arguments to pass to ase.io.write to save generated structures.
         Default is {}.
     logger : Optional[Logger]
@@ -88,10 +87,15 @@ def _calc_volumes_energies(  # pylint: disable=too-many-locals
         volumes.append(c_struct.get_volume())
         energies.append(c_struct.get_potential_energy())
 
-        if write_structures:
-            # Always append first original structure
-            write_kwargs["append"] = True
-            write(images=c_struct, **write_kwargs)
+        # Always append first original structure
+        write_kwargs["append"] = True
+        # Write structures, but no need to set info c_struct is not used elsewhere
+        output_structs(
+            images=c_struct,
+            write_results=write_structures,
+            set_info=False,
+            write_kwargs=write_kwargs,
+        )
 
     if logger:
         tracker.stop_task()
@@ -113,7 +117,7 @@ def calc_eos(
     minimize_kwargs: Optional[dict[str, Any]] = None,
     write_results: bool = True,
     write_structures: bool = False,
-    write_kwargs: Optional[ASEWriteArgs] = None,
+    write_kwargs: Optional[OutputKwargs] = None,
     file_prefix: Optional[PathLike] = None,
     log_kwargs: Optional[dict[str, Any]] = None,
     tracker_kwargs: Optional[dict[str, Any]] = None,
@@ -145,7 +149,7 @@ def calc_eos(
         True to write out results of equation of state calculations. Default is True.
     write_structures : bool
         True to write out all genereated structures. Default is False.
-    write_kwargs : Optional[ASEWriteArgs],
+    write_kwargs : Optional[OutputKwargs],
         Keyword arguments to pass to ase.io.write to save generated structures.
         Default is {}.
     file_prefix : Optional[PathLike]
@@ -205,8 +209,10 @@ def calc_eos(
             }
         optimize(struct, **minimize_kwargs)
 
-    if write_structures:
-        write(images=struct, **write_kwargs)
+        # Optionally write structure to file
+        output_structs(
+            images=struct, write_results=write_structures, write_kwargs=write_kwargs
+        )
 
     # Set constant volume for geometry optimization of generated structures
     if "filter_kwargs" in minimize_kwargs:

janus_core/calculations/geom_opt.py

@@ -7,15 +7,15 @@
 
 from ase import Atoms, filters, units
 from ase.filters import FrechetCellFilter
-from ase.io import read, write
+from ase.io import read
 import ase.optimize
 from ase.optimize import LBFGS
 from ase.optimize.optimize import Optimizer
 from numpy import linalg
 
-from janus_core.helpers.janus_types import ASEOptArgs, ASEWriteArgs
+from janus_core.helpers.janus_types import ASEOptArgs, OutputKwargs
 from janus_core.helpers.log import config_logger, config_tracker
-from janus_core.helpers.utils import none_to_dict, spacegroup
+from janus_core.helpers.utils import none_to_dict, output_structs, spacegroup
 
 
 def _set_functions(
@@ -126,8 +126,8 @@ def optimize(  # pylint: disable=too-many-arguments,too-many-locals,too-many-bra
     optimizer: Callable = LBFGS,
     opt_kwargs: Optional[ASEOptArgs] = None,
     write_results: bool = False,
-    write_kwargs: Optional[ASEWriteArgs] = None,
-    traj_kwargs: Optional[ASEWriteArgs] = None,
+    write_kwargs: Optional[OutputKwargs] = None,
+    traj_kwargs: Optional[OutputKwargs] = None,
     log_kwargs: Optional[dict[str, Any]] = None,
     tracker_kwargs: Optional[dict[str, Any]] = None,
 ) -> Atoms:
@@ -161,10 +161,10 @@ def optimize(  # pylint: disable=too-many-arguments,too-many-locals,too-many-bra
         Keyword arguments to pass to optimizer. Default is {}.
     write_results : bool
         True to write out optimized structure. Default is False.
-    write_kwargs : Optional[ASEWriteArgs],
+    write_kwargs : Optional[OutputKwargs],
         Keyword arguments to pass to ase.io.write to save optimized structure.
         Default is {}.
-    traj_kwargs : Optional[ASEWriteArgs]
+    traj_kwargs : Optional[OutputKwargs]
         Keyword arguments to pass to ase.io.write to save optimization trajectory.
         Must include "filename" keyword. Default is {}.
     log_kwargs : Optional[dict[str, Any]]
@@ -239,13 +239,20 @@ def optimize(  # pylint: disable=too-many-arguments,too-many-locals,too-many-bra
         )
 
     # Write out optimized structure
-    if write_results:
-        write(images=struct, **write_kwargs)
+    output_structs(
+        struct,
+        write_results=write_results,
+        write_kwargs=write_kwargs,
+    )
 
     # Reformat trajectory file from binary
     if traj_kwargs:
         traj = read(opt_kwargs["trajectory"], index=":")
-        write(images=traj, **traj_kwargs)
+        output_structs(
+            traj,
+            write_results=True,
+            write_kwargs=traj_kwargs,
+        )
 
     if logger:
         tracker.stop()

janus_core/calculations/md.py

@@ -12,7 +12,7 @@
 
 from ase import Atoms, units
 from ase.geometry.analysis import Analysis
-from ase.io import read, write
+from ase.io import read
 from ase.md.langevin import Langevin
 from ase.md.npt import NPT as ASE_NPT
 from ase.md.velocitydistribution import (
@@ -29,12 +29,13 @@
 from janus_core.helpers.janus_types import (
     CorrelationKwargs,
     Ensembles,
+    OutputKwargs,
     PathLike,
     PostProcessKwargs,
 )
 from janus_core.helpers.log import config_logger, config_tracker
 from janus_core.helpers.post_process import compute_rdf, compute_vaf
-from janus_core.helpers.utils import FileNameMixin
+from janus_core.helpers.utils import FileNameMixin, output_structs
 
 DENS_FACT = (units.m / 1.0e2) ** 3 / units.mol
 
@@ -111,6 +112,9 @@ class MolecularDynamics(FileNameMixin):  # pylint: disable=too-many-instance-att
         heating.
     temp_time : Optional[float]
         Time between heating steps, in fs. Default is None, which disables heating.
+    write_kwargs : Optional[OutputKwargs],
+        Keyword arguments to pass to `output_structs` when saving trajectory and final
+        files. Default is {}.
     post_process_kwargs : Optional[PostProcessKwargs]
         Keyword arguments to control post-processing operations.
     correlation_kwargs : Optional[CorrelationKwargs]
@@ -182,6 +186,7 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
         temp_end: Optional[float] = None,
         temp_step: Optional[float] = None,
         temp_time: Optional[float] = None,
+        write_kwargs: Optional[OutputKwargs] = None,
         post_process_kwargs: Optional[PostProcessKwargs] = None,
         correlation_kwargs: Optional[list[CorrelationKwargs]] = None,
         log_kwargs: Optional[dict[str, Any]] = None,
@@ -262,6 +267,9 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
             disables heating.
         temp_time : Optional[float]
             Time between heating steps, in fs. Default is None, which disables heating.
+        write_kwargs : Optional[OutputKwargs],
+            Keyword arguments to pass to `output_structs` when saving trajectory and
+            final files. Default is {}.
         post_process_kwargs : Optional[PostProcessKwargs]
             Keyword arguments to control post-processing operations.
         correlation_kwargs : Optional[list[CorrelationKwargs]]
@@ -300,6 +308,7 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
         self.temp_end = temp_end
         self.temp_step = temp_step
         self.temp_time = temp_time * units.fs if temp_time else None
+        self.write_kwargs = write_kwargs if write_kwargs is not None else {}
         self.post_process_kwargs = (
             post_process_kwargs if post_process_kwargs is not None else {}
         )
@@ -312,6 +321,12 @@ def __init__(  # pylint: disable=too-many-arguments,too-many-locals,too-many-sta
 
         FileNameMixin.__init__(self, struct, struct_name, file_prefix, ensemble)
 
+        self.write_kwargs.setdefault(
+            "columns", ["symbols", "positions", "momenta", "masses"]
+        )
+        if "append" in self.write_kwargs:
+            raise ValueError("`append` cannot be specified when writing files")
+
         self.log_kwargs = (
             log_kwargs if log_kwargs else {}
         )  # pylint: disable=duplicate-code
@@ -595,11 +610,13 @@ def _write_traj(self) -> None:
                 self.dyn.nsteps > self.traj_start + self.traj_start % self.traj_every
             )
 
-            self.dyn.atoms.write(
-                self.traj_file,
-                write_info=True,
-                columns=["symbols", "positions", "momenta", "masses"],
-                append=append,
+            write_kwargs = self.write_kwargs
+            write_kwargs["filename"] = self.traj_file
+            write_kwargs["append"] = append
+            output_structs(
+                images=self.struct,
+                write_results=True,
+                write_kwargs=write_kwargs,
             )
 
     def _write_final_state(self) -> None:
@@ -611,12 +628,13 @@ def _write_final_state(self) -> None:
         # Append if final file has been created
         append = self.created_final_file
 
-        write(
-            self.final_file,
-            self.struct,
-            write_info=True,
-            columns=["symbols", "positions", "momenta", "masses"],
-            append=append,
+        write_kwargs = self.write_kwargs
+        write_kwargs["filename"] = self.final_file
+        write_kwargs["append"] = append
+        output_structs(
+            images=self.struct,
+            write_results=True,
+            write_kwargs=write_kwargs,
         )
 
     def _post_process(self) -> None:
@@ -708,11 +726,12 @@ def _write_restart(self) -> None:
         """Write restart file and (optionally) rotate files saved."""
         step = self.offset + self.dyn.nsteps
         if step > 0:
-            write(
-                self._restart_file,
-                self.struct,
-                write_info=True,
-                columns=["symbols", "positions", "momenta", "masses"],
+            write_kwargs = self.write_kwargs
+            write_kwargs["filename"] = self._restart_file
+            output_structs(
+                images=self.struct,
+                write_results=True,
+                write_kwargs=write_kwargs,
             )
             if self.rotate_restart:
                 self.restart_files.append(self._restart_file)