diff --git a/janus_core/single_point.py b/janus_core/single_point.py index 5de81564..231ab59f 100644 --- a/janus_core/single_point.py +++ b/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,18 +183,25 @@ 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: @@ -200,4 +209,7 @@ def run_single_point( if "stress" in properties or len(properties) == 0: results["stress"] = self._get_stress() + if write_kwargs: + write(images=self.struct, **write_kwargs) + return results