Add option to save singlepoint results

janus_core/single_point.py

@@ -1,9 +1,9 @@
 """Prepare and perform single point calculations."""
 
-import pathlib
+from pathlib import Path
 from typing import Any, Literal, Optional, Union
 
-from ase.io import read
+from ase.io import read, write
 from numpy import ndarray
 
 from janus_core.mlip_calculators import architectures, choose_calculator, devices
@@ -98,7 +98,7 @@ def read_structure(self, **kwargs) -> None:
             Keyword arguments passed to ase.io.read.
         """
         self.struct = read(self.structure, **kwargs)
-        self.structname = pathlib.Path(self.structure).stem
+        self.structname = Path(self.structure).stem
 
     def set_calculator(
         self, read_kwargs: Optional[dict[str, Any]] = None, **kwargs
@@ -171,7 +171,9 @@ def _get_stress(self) -> Union[ndarray, list[ndarray]]:
         return self.struct.get_stress()
 
     def run_single_point(
-        self, properties: Optional[Union[str, list[str]]] = None
+        self,
+        properties: Optional[Union[str, list[str]]] = None,
+        write_kwargs: Optional[dict[str, Any]] = None,
     ) -> dict[str, Any]:
         """
         Run single point calculations.
@@ -181,23 +183,33 @@ def run_single_point(
         properties : Optional[Union[str, list[str]]]
             Physical properties to calculate. If not specified, "energy",
             "forces", and "stress" will be returned.
+        write_kwargs : Optional[dict[str, Any]] = None,
+            Keyword arguments to pass to ase.io.write to save structure with
+            results of calculations. Must include "filename" keyword. Default is {}.
 
         Returns
         -------
         dict[str, Any]
             Dictionary of calculated results.
         """
-        results = {}
         if properties is None:
             properties = []
         if isinstance(properties, str):
             properties = [properties]
 
+        write_kwargs = write_kwargs if write_kwargs else {}
+        if write_kwargs and "filename" not in write_kwargs:
+            raise ValueError("'filename' must be included in write_kwargs")
+
+        results = {}
         if "energy" in properties or len(properties) == 0:
             results["energy"] = self._get_potential_energy()
         if "forces" in properties or len(properties) == 0:
             results["forces"] = self._get_forces()
         if "stress" in properties or len(properties) == 0:
             results["stress"] = self._get_stress()
 
+        if write_kwargs:
+            write(images=self.struct, **write_kwargs)
+
         return results