Scan CalcType + ScanData + Tests

src/qcdata/__init__.py

@@ -27,6 +27,7 @@
     "SinglePointData",
     "OptimizationData",
     "ConformerSearchData",
+    "ScanData",
     "SinglePointResults",
     "OptimizationResults",
     "ConformerSearchResults",

src/qcdata/models/base_models.py

@@ -249,9 +249,7 @@ def _serialize_files(self, files, _info) -> dict[str, str]:
             for filename, data in files.items()
         }
 
-    def add_file(
-        self, filepath: Path | str, relative_dir: Path | None = None
-    ) -> None:
+    def add_file(self, filepath: Path | str, relative_dir: Path | None = None) -> None:
         """Add a file to the object. The file will be added at to the `files` attribute
             with the filename as the key and the file data as the value.
 
@@ -381,6 +379,7 @@ class CalcType(str, Enum):
     optimization = "optimization"
     transition_state = "transition_state"
     conformer_search = "conformer_search"
+    scan = "scan"
 
     def __repr__(self) -> str:
         """Custom repr for CalcType"""

src/qcdata/models/data.py

@@ -25,6 +25,7 @@
     "SinglePointData",
     "OptimizationData",
     "ConformerSearchData",
+    "ScanData",
     "StructuredData",
     "StructuredDataType",
     "Data",
@@ -150,10 +151,7 @@ def return_result(self, calctype: CalcType) -> float | SerializableNDArray:
     @model_validator(mode="after")
     def _ensure_results(self) -> Self:
         """Ensure that at least one result is present."""
-        if all(
-            result is None
-            for result in [self.energy, self.gradient, self.hessian]
-        ):
+        if all(result is None for result in [self.energy, self.gradient, self.hessian]):
             raise ValueError(
                 "SinglePointResults requires either an energy, gradient, or hessian "
                 "value."
@@ -363,12 +361,80 @@ def conformers_filtered(
         )
 
 
-StructuredData = Union[SinglePointData, OptimizationData, ConformerSearchData]
+class ScanData(Files, CalcInfoData):
+    """Computed data for a scan (may be for a relaxed or frozen).
+
+    Attributes
+    ----------
+        energies: The energies for each step of the scan.
+        structures: The Structure objects for each step of the scan.
+        trajectory: The ProgramOutput objects for each step of the scan.
+    """
+
+    trajectory: list[ProgramOutput[ProgramInput, OptimizationData]]
+
+    @property
+    def energies(self) -> np.ndarray:
+        """The energies for each step of the scan."""
+        return np.array(
+            [
+                output.data.final_energy if output.success else np.nan
+                for output in self.trajectory
+            ],
+            dtype=float,
+        )
+
+    @property
+    def structures(self) -> list[Structure]:
+        """The Structure objects for each step of the optimization."""
+        return [output.data.final_structure for output in self.trajectory]
+
+    def __repr_args__(self) -> list[tuple[str, str]]:
+        """Avoid printing the entire collection of objects in representation."""
+        return [
+            ("trajectory", "[...]"),
+            ("energies", "[...]"),
+            ("structures", "[...]"),
+        ]
+
+    def to_xyz(self) -> str:
+        """Return the trajectory as an `xyz` string."""
+        return to_multi_xyz(self.structures)
+
+    def save(
+        self,
+        filepath: Path | str,
+        exclude_none: bool = True,
+        exclude_unset: bool = True,
+        indent: int = 4,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        """Save a ScanData to a file.
+
+        Args:
+            filepath: The path to save the molecule to.
+            exclude_none: If True, attributes with a value of None will not be written
+                to the file.
+            exclude_unset: If True, attributes that have not been set will not be
+                written to the file.
+            **kwargs: Additional keyword arguments to pass to the json serializer.
+
+        Note:
+            If the filepath has a `.xyz` extension, the trajectory will be saved to a
+            multi-structure `xyz` file.
+        """
+        filepath = Path(filepath)
+        if filepath.suffix == ".xyz":
+            filepath.write_text(self.to_xyz())
+            return
+        super().save(filepath, exclude_none, exclude_unset, indent, **kwargs)
+
+
+StructuredData = Union[SinglePointData, OptimizationData, ConformerSearchData, ScanData]
 StructuredDataType = TypeVar("StructuredDataType", bound=StructuredData)
 Data = Union[Files, StructuredData]
 DataType = TypeVar("DataType", bound=Data)
 
-
 @deprecated_class("SinglePointData")
 class SinglePointResults(SinglePointData):
     """This class is deprecated and will be removed in a future release. Please use

src/qcdata/models/outputs.py

@@ -13,7 +13,14 @@
 from qcdata.helper_types import SerializableNDArray
 
 from .base_models import Files, Provenance, QCDataBaseModel
-from .data import ConformerSearchData, Data, DataType, OptimizationData, SinglePointData
+from .data import (
+    ConformerSearchData,
+    Data,
+    DataType,
+    OptimizationData,
+    ScanData,
+    SinglePointData,
+)
 from .inputs import FileInput, Inputs, ProgramInput
 from .inputs import InputType as ProgramInputType
 from .structure import Structure
@@ -205,6 +212,7 @@ class Results(ProgramOutput[ProgramInputType, DataType]):
 Results.model_rebuild()
 OptimizationData.model_rebuild()
 ConformerSearchData.model_rebuild()
+ScanData.model_rebuild()
 
 def _register_program_output_classes():
     """Required so that pickle can find the concrete classes for serialization."""

tests/conftest.py

@@ -144,3 +144,18 @@ def results_failure(prog_input_factory, sp_data):
 @pytest.fixture
 def opt_data(prog_output):
     return OptimizationData(trajectory=[prog_output])
+
+
+@pytest.fixture
+def opt_output(prog_input_factory, opt_data):
+    """Successful ProgramOutput object"""
+    input_data = prog_input_factory("optimization")
+
+    return ProgramOutput[ProgramInput, OptimizationData](
+        input_data=input_data,
+        success=True,
+        logs="program standard out...",
+        data=opt_data,
+        provenance={"program": "qcdata-test-suite", "scratch_dir": "/tmp/qcdata"},
+        extras={"some_extra": 1},
+    )

tests/test_scan_data.py

@@ -0,0 +1,58 @@
+from qcdata import ScanData
+
+
+def test_scan_data_properties(opt_output):
+    """Test that the number of energies matches the number of conformers"""
+    # No energies is fine
+    scan_res = ScanData(
+        trajectory=[opt_output],
+    )
+
+    # Test properties
+    assert scan_res.energies == [opt_output.data.final_energy]
+    assert scan_res.structures == [opt_output.data.final_structure]
+    # Test custom __repr_args__
+    repr_args = scan_res.__repr_args__()
+    assert isinstance(repr_args, list)
+    for arg in repr_args:
+        assert isinstance(arg, tuple)
+        assert len(arg) == 2
+        assert isinstance(arg[0], str)
+        assert isinstance(arg[1], str)
+
+
+def test_scan_save_to_xyz(opt_output, tmp_path):
+    scan_res = ScanData(
+        trajectory=[opt_output] * 3,
+    )
+    scan_res.save(tmp_path / "scan_res.xyz")
+
+    text = (tmp_path / "scan_res.xyz").read_text()
+
+    # Text must be de-dented exactly as below
+    correct_text = """3
+qcdata_charge=0 qcdata_multiplicity=1 qcdata__identifiers_name=water
+O  0.01340919176202180 0.01026321207824930 -0.00368477733600419
+H  0.12112430307330672 0.97600619725464122 0.08599884278042236
+H  0.75016279902412597 -0.33132205318865016 -0.54481406902570462
+3
+qcdata_charge=0 qcdata_multiplicity=1 qcdata__identifiers_name=water
+O  0.01340919176202180 0.01026321207824930 -0.00368477733600419
+H  0.12112430307330672 0.97600619725464122 0.08599884278042236
+H  0.75016279902412597 -0.33132205318865016 -0.54481406902570462
+3
+qcdata_charge=0 qcdata_multiplicity=1 qcdata__identifiers_name=water
+O  0.01340919176202180 0.01026321207824930 -0.00368477733600419
+H  0.12112430307330672 0.97600619725464122 0.08599884278042236
+H  0.75016279902412597 -0.33132205318865016 -0.54481406902570462
+"""
+    assert text == correct_text
+
+
+def test_scan_save_non_xyz(opt_output, tmp_path):
+    scan_res = ScanData(
+        trajectory=[opt_output] * 3,
+    )
+    scan_res.save(tmp_path / "scan_res.json")
+    scan_res_copy = ScanData.open(tmp_path / "scan_res.json")
+    assert scan_res == scan_res_copy