diff --git a/docs/api/molgraph.md b/docs/api/molgraph.md index a66d3af..ff4f1f4 100644 --- a/docs/api/molgraph.md +++ b/docs/api/molgraph.md @@ -96,4 +96,7 @@ List of atomic radii for each atom in the molecule ### `x, y, z` Lists of coordinates for each atom +### `comment` +Optional comment or title from the second line of XYZ file + These properties are automatically updated when reading an XYZ file or modifying the molecular structure. diff --git a/tests/test_molgraph.py b/tests/test_molgraph.py index b4f8259..8946a44 100644 --- a/tests/test_molgraph.py +++ b/tests/test_molgraph.py @@ -44,6 +44,7 @@ def test_read_xyz(self, mol_graph: MolGraph) -> None: assert mol_graph.elements == ["O", "H", "H"] assert len(mol_graph.atomic_radii) == 3 assert mol_graph.atomic_radii[0] == DEFAULT_RADII["O"] + assert mol_graph.comment == "Water molecule" def test_to_plotly(self, mol_graph: MolGraph) -> None: """Test Plotly figure generation.""" @@ -170,6 +171,48 @@ def test_validate_true_invalid(self, tmp_path: Path) -> None: ): mol.read_xyz(xyz_file, validate=True) + def test_xyz_comment_handling(self, tmp_path: Path) -> None: + """Test handling of XYZ file comments.""" + + # Test with normal comment + xyz_with_comment = """3 + Test Comment + O 0.0 0.0 0.0 + H 0.757 0.586 0.0 + H -0.757 0.586 0.0 + """ + file_path = tmp_path / "with_comment.xyz" + file_path.write_text(xyz_with_comment) + + mol = MolGraph() + mol.read_xyz(str(file_path)) + assert mol.comment == "Test Comment" + + # Test with empty comment + xyz_empty_comment = """3 + + O 0.0 0.0 0.0 + H 0.757 0.586 0.0 + H -0.757 0.586 0.0 + """ + file_path = tmp_path / "empty_comment.xyz" + file_path.write_text(xyz_empty_comment) + + mol = MolGraph() + mol.read_xyz(str(file_path)) + assert mol.comment == "" + + # Test with xyz_start=0 (should not set comment) + xyz_no_header = """O 0.0 0.0 0.0 + H 0.757 0.586 0.0 + H -0.757 0.586 0.0""" + file_path = tmp_path / "no_header.xyz" + file_path.write_text(xyz_no_header) + + mol = MolGraph() + mol.read_xyz(str(file_path), xyz_start=0) + assert mol.comment == "" + def test_error_handling(self) -> None: """Test error conditions.""" mol = MolGraph() diff --git a/xyz2graph/xyz2graph.py b/xyz2graph/xyz2graph.py index 235f84f..13ba917 100755 --- a/xyz2graph/xyz2graph.py +++ b/xyz2graph/xyz2graph.py @@ -150,6 +150,7 @@ class MolGraph: x: List of x-coordinates for each atom y: List of y-coordinates for each atom z: List of z-coordinates for each atom + comment: Optional comment or title from the XYZ file adj_list: Dictionary mapping atom indices to their connected neighbors atomic_radii: List of atomic radii for each atom bond_lengths: Dictionary mapping pairs of connected atoms to their bond lengths @@ -163,6 +164,7 @@ class MolGraph: x: List[float] = field(default_factory=list) y: List[float] = field(default_factory=list) z: List[float] = field(default_factory=list) + comment: str = field(default="") adj_list: Dict[int, Set[int]] = field(default_factory=dict) atomic_radii: List[float] = field(default_factory=list) bond_lengths: Dict[FrozenSet[int], float] = field(default_factory=dict) @@ -326,6 +328,11 @@ def read_xyz( "Could not read atom count from first line" ) from err + # Store comment if available (line 1 in standard XYZ format) + if xyz_start > 1: + self.comment = lines[1].strip() + logger.debug(f"Read comment: {self.comment}") + coordinates = self._parse_coordinates(lines[xyz_start:]) self.elements, self.x, self.y, self.z = coordinates