Fix for writing non-unique site labels in CifWriter (#3767)

* Add check for non-unique labels in CifWriter

* Append _$NUM if label already exists

* Test that labels work as intended

* Remove redundant check

* Update tests/io/test_cif.py

Co-authored-by: Janosh Riebesell <[email protected]>
Signed-off-by: Stef Smeets <[email protected]>

* Add warning for duplicate sites and method for relabeling

* Add test for relabeling

* Work around groupby caveat

* return self from SiteColl.relabel_sites(), advise to use SiteColl.copy() to avoid in-place mod in doc string

* Add ignore_uniq parameter to relabel_sites method in SiteCollection class

* improve coverage in TestStructure.test_relabel_sites

* assert unique labels do not warn in test_cif_writer_non_unique_labels

---------

Signed-off-by: Stef Smeets <[email protected]>
Co-authored-by: Janosh Riebesell <[email protected]>

Author: stefsmeets

pymatgen/core/structure.py

@@ -335,6 +335,32 @@ def labels(self) -> list[str]:
         """Site labels as a list."""
         return [site.label for site in self]
 
+    def relabel_sites(self, ignore_uniq: bool = False) -> Self:
+        """Relabel sites to ensure they are unique.
+
+        Site labels are updated in-place, and relabeled by suffixing _1, _2, ..., _n for duplicates.
+        Call Structure.copy().relabel_sites() to avoid modifying the original structure.
+
+        Args:
+            ignore_uniq (bool): If True, do not relabel sites that already have unique labels.
+                Defaults to False.
+
+        Returns:
+            SiteCollection: self with relabeled sites.
+        """
+        grouped = collections.defaultdict(list)
+        for site in self:
+            grouped[site.label].append(site)
+
+        for label, sites in grouped.items():
+            if len(sites) == 0 or (len(sites) == 1 and ignore_uniq):
+                continue
+
+            for idx, site in enumerate(sites):
+                site.label = f"{label}_{idx + 1}"
+
+        return self
+
     def __contains__(self, site: object) -> bool:
         return site in self.sites
 

pymatgen/io/cif.py

@@ -1568,6 +1568,14 @@ def __init__(
                     atom_site_occupancy.append(str(occu))
                     count += 1
 
+        if len(set(atom_site_label)) != len(atom_site_label):
+            warnings.warn(
+                "Site labels are not unique, which is not compliant with the CIF spec "
+                "(https://www.iucr.org/__data/iucr/cifdic_html/1/cif_core.dic/Iatom_site_label.html):"
+                f"`{atom_site_label}`.",
+                UserWarning,
+            )
+
         block["_atom_site_type_symbol"] = atom_site_type_symbol
         block["_atom_site_label"] = atom_site_label
         block["_atom_site_symmetry_multiplicity"] = atom_site_symmetry_multiplicity

tests/core/test_structure.py

@@ -1005,6 +1005,25 @@ def test_replace_species_labels(self):
         new2 = struct.replace_species({"Si": replacement}, in_place=False)
         assert new2.labels == [label] * len(struct)
 
+    def test_relabel_sites(self):
+        struct = self.struct.copy()
+        assert self.struct.labels == ["Si", "Si"]
+        relabeled = struct.relabel_sites()
+        assert relabeled is struct
+        assert relabeled.labels == struct.labels == ["Si_1", "Si_2"]
+
+        struct.replace_species({"Si": "Li"})
+        assert struct.labels == ["Li", "Li"]
+        struct.relabel_sites()
+        assert struct.labels == ["Li_1", "Li_2"]
+
+        li_si = self.struct.copy().replace(0, "Li")
+        assert li_si.labels == ["Li", "Si"]
+        li_si.relabel_sites(ignore_uniq=True)  # check no-op for unique labels
+        assert li_si.labels == ["Li", "Si"]
+        li_si.relabel_sites()  # check no-op for unique labels
+        assert li_si.labels == ["Li_1", "Si_1"]
+
     def test_append_insert_remove_replace_substitute(self):
         struct = self.struct
         struct.insert(1, "O", [0.5, 0.5, 0.5])

tests/io/test_cif.py

@@ -1084,3 +1084,26 @@ def test_bibtex(self):
 }
 """
         assert self.mcif_ncl.get_bibtex_string() == ref_bibtex_string
+
+
+def test_cif_writer_non_unique_labels(capsys):
+    # https://github.com/materialsproject/pymatgen/issues/3761
+    parser = CifParser(f"{TEST_FILES_DIR}/cif/garnet.cif")
+    struct = parser.parse_structures()[0]
+
+    assert struct.labels[0:3] == ["Ca1", "Ca1", "Ca1"]
+    assert len(set(struct.labels)) != len(struct.labels)
+
+    # This should raise a warning
+    with pytest.warns(UserWarning, match="Site labels are not unique, which is not compliant with the CIF spec"):
+        CifWriter(struct)
+
+    struct.relabel_sites()
+    assert struct.labels[0:3] == ["Ca1_1", "Ca1_2", "Ca1_3"]
+
+    _ = capsys.readouterr()
+    # This should not raise a warning
+    CifWriter(struct)
+    stdout, stderr = capsys.readouterr()
+    assert stdout == ""
+    assert stderr == ""