From e3036f82acf05c7cb89c338915550fe500f72930 Mon Sep 17 00:00:00 2001 From: Troy Smith Date: Tue, 21 Apr 2026 10:09:57 -0400 Subject: [PATCH 1/3] Scan CalcType + ScanData + Tests --- src/qcdata/__init__.py | 1 + src/qcdata/models/base_models.py | 5 +-- src/qcdata/models/data.py | 75 ++++++++++++++++++++++++++++++-- tests/conftest.py | 15 +++++++ tests/test_scan_data.py | 60 +++++++++++++++++++++++++ 5 files changed, 149 insertions(+), 7 deletions(-) create mode 100644 tests/test_scan_data.py diff --git a/src/qcdata/__init__.py b/src/qcdata/__init__.py index e1ef28c..a98ba40 100644 --- a/src/qcdata/__init__.py +++ b/src/qcdata/__init__.py @@ -27,6 +27,7 @@ "SinglePointData", "OptimizationData", "ConformerSearchData", + "ScanData", "SinglePointResults", "OptimizationResults", "ConformerSearchResults", diff --git a/src/qcdata/models/base_models.py b/src/qcdata/models/base_models.py index ede9608..e2df3e0 100644 --- a/src/qcdata/models/base_models.py +++ b/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""" diff --git a/src/qcdata/models/data.py b/src/qcdata/models/data.py index 3b0b949..79e3629 100644 --- a/src/qcdata/models/data.py +++ b/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,6 +361,75 @@ def conformers_filtered( ) +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] StructuredDataType = TypeVar("StructuredDataType", bound=StructuredData) Data = Union[Files, StructuredData] diff --git a/tests/conftest.py b/tests/conftest.py index a3c68d9..1ada2c5 100644 --- a/tests/conftest.py +++ b/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}, + ) diff --git a/tests/test_scan_data.py b/tests/test_scan_data.py new file mode 100644 index 0000000..3aadd95 --- /dev/null +++ b/tests/test_scan_data.py @@ -0,0 +1,60 @@ +from qcdata import ProgramOutput, ScanData + +ScanData.model_rebuild() + + +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 From e943ba41cf40e3fa73296ca85678187686ea050f Mon Sep 17 00:00:00 2001 From: Colton Hicks Date: Tue, 21 Apr 2026 11:50:30 -0700 Subject: [PATCH 2/3] Removed unused import --- tests/test_scan_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_scan_data.py b/tests/test_scan_data.py index 3aadd95..df34f39 100644 --- a/tests/test_scan_data.py +++ b/tests/test_scan_data.py @@ -1,4 +1,4 @@ -from qcdata import ProgramOutput, ScanData +from qcdata import ScanData ScanData.model_rebuild() From 11cc1794a8092a416e7b649c6590c747acf1aece Mon Sep 17 00:00:00 2001 From: Colton Hicks Date: Tue, 21 Apr 2026 11:57:35 -0700 Subject: [PATCH 3/3] Moved ScanData.model_rebuild() to outputs.py. Added ScanData to StructuredData union. Removed ProgramInput from test_scan_data.py --- src/qcdata/models/data.py | 3 +-- src/qcdata/models/outputs.py | 10 +++++++++- tests/test_scan_data.py | 2 -- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/qcdata/models/data.py b/src/qcdata/models/data.py index 79e3629..f756c52 100644 --- a/src/qcdata/models/data.py +++ b/src/qcdata/models/data.py @@ -430,12 +430,11 @@ def save( super().save(filepath, exclude_none, exclude_unset, indent, **kwargs) -StructuredData = Union[SinglePointData, OptimizationData, ConformerSearchData] +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 diff --git a/src/qcdata/models/outputs.py b/src/qcdata/models/outputs.py index 81feea7..230ef21 100644 --- a/src/qcdata/models/outputs.py +++ b/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.""" diff --git a/tests/test_scan_data.py b/tests/test_scan_data.py index df34f39..5c745b9 100644 --- a/tests/test_scan_data.py +++ b/tests/test_scan_data.py @@ -1,7 +1,5 @@ from qcdata import ScanData -ScanData.model_rebuild() - def test_scan_data_properties(opt_output): """Test that the number of energies matches the number of conformers"""