Fix surface.center_slab handling of Slab regions separated by vacuum

src/pymatgen/core/surface.py

@@ -738,57 +738,58 @@ def get_equi_sites(slab: Slab, sites: list[int]) -> list[int]:
 
 
 def center_slab(slab: Structure) -> Structure:
-    """Relocate the slab to the center such that its center
-    (the slab region) is close to z=0.5.
+    """Relocate the slab such that its center is at z=0.5.
+
+    # TODO: docstring
+    Also handles the case where the slab is split into two separate
+    regions by vacuum due to PBC, with the two regions joined.
 
     This makes it easier to find surface sites and apply
     operations like doping.
 
-    There are two possible cases:
-        1. When the slab region is completely positioned between
-        two vacuum layers in the cell but is not centered, we simply
-        shift the slab to the center along z-axis.
-        2. If the slab completely resides outside the cell either
-        from the bottom or the top, we iterate through all sites that
-        spill over and shift all sites such that it is now
-        on the other side. An edge case being, either the top
-        of the slab is at z = 0 or the bottom is at z = 1.
-
     Args:
         slab (Structure): The slab to center.
 
     Returns:
         Structure: The centered slab.
     """
-    # Get all site indices
-    all_indices = list(range(len(slab)))
-
-    # Get a reasonable cutoff radius to sample neighbors
-    bond_dists = sorted(nn[1] for nn in slab.get_neighbors(slab[0], 10) if nn[1] > 0)
-    # TODO (@DanielYang59): magic number for cutoff radius (would 3 be too large?)
-    cutoff_radius = bond_dists[0] * 3
-
-    # TODO (@DanielYang59): do we need the following complex method?
-    # Why don't we just calculate the center of the Slab and move it to z=0.5?
-    # Before moving we need to ensure there is only one Slab layer though
-
-    # If structure is case 2, shift all the sites
-    # to the other side until it is case 1
-    for site in slab:  # DEBUG (@DanielYang59): Slab position changes during loop?
-        # DEBUG (@DanielYang59): sites below z=0 is not considered (only check coord > c)
-        if any(nn[1] >= slab.lattice.c for nn in slab.get_neighbors(site, cutoff_radius)):
-            # TODO (@DanielYang59): the magic offset "0.05" seems unnecessary,
-            # as the Slab would be centered later anyway
-            shift = 1 - site.frac_coords[2] + 0.05
-            slab.translate_sites(all_indices, [0, 0, shift])
-
-    # Now the slab is case 1, move it to the center
-    weights = [site.species.weight for site in slab]
-    center_of_mass = np.average(slab.frac_coords, weights=weights, axis=0)
-    shift = 0.5 - center_of_mass[2]
-
-    slab.translate_sites(all_indices, [0, 0, shift])
+    # Cutoff radius guess (3 times the smallest bond length) to sample neighbors
+    cutoff_radius: float = sorted(nn[1] for nn in slab.get_neighbors(slab[0], 10) if nn[1] > 0)[0] * 3
+
+    slab_regions: list[tuple[float, float]] = get_slab_regions(slab, blength=cutoff_radius)
+
+    if len(slab_regions) == 1:
+        pass
+
+    # If the slab is separated by vacuum, combine them by putting the top region
+    # under the bottom region
+    elif len(slab_regions) == 2:
+        # Get smallest z of the top slab segment
+        _bottom_region, top_region = sorted(slab_regions, key=lambda r: r[1])
+        top_lowest_z = top_region[0]
+
+        # Get all sites of the top region (z >= top_lowest_z)
+        top_region_indices = [i for i, site in enumerate(slab) if site.frac_coords[2] >= top_lowest_z]
+
+        # Shift the top region down by one full unit cell
+        slab.translate_sites(
+            indices=top_region_indices,
+            vector=[0, 0, -1.0],
+            to_unit_cell=False,
+        )
+
+        # Check if two regions have been combined (in case distance too large)
+        if len(get_slab_regions(slab, blength=cutoff_radius)) != 1:
+            raise RuntimeError("Unable to combine slab regions (distance too large or cutoff too small).")
+
+    else:
+        raise RuntimeError(f"Unable to handle cases where there're {len(slab_regions)} slab regions")
 
+    weights: list[float] = [site.species.weight for site in slab]
+    center_of_mass: float = np.average(slab.frac_coords, weights=weights, axis=0)
+
+    shift: float = 0.5 - center_of_mass[2]
+    slab.translate_sites(indices=list(range(len(slab))), vector=[0, 0, shift])
     return slab
 
 
@@ -972,10 +973,7 @@ def calculate_surface_normal() -> np.ndarray:
             return normal
 
         def calculate_scaling_factor() -> np.ndarray:
-            """Calculate scaling factor.
-
-            # TODO (@DanielYang59): revise docstring to add more details.
-            """
+            """Calculate scaling factor."""
             slab_scale_factor = []
             non_orth_ind = []
             eye = np.eye(3, dtype=np.int64)
@@ -1068,7 +1066,7 @@ def calculate_scaling_factor() -> np.ndarray:
         self.min_slab_size = min_slab_size
         self.in_unit_planes = in_unit_planes
         self.primitive = primitive
-        self._normal = normal  # TODO (@DanielYang59): used only in unit test
+        self._normal = normal  # used only in unit test
         self.reorient_lattice = reorient_lattice
 
         _a, _b, c = self.oriented_unit_cell.lattice.matrix
@@ -1857,7 +1855,6 @@ def build_recon_json() -> dict:
                 # loaded "recon_json", use condition to avoid this
                 recon_json = copy.deepcopy(RECONSTRUCTIONS_ARCHIVE[recon_json["base_reconstruction"]])
 
-                # TODO (@DanielYang59): use "site" over "point" for consistency?
                 if "points_to_add" in recon_json:
                     del recon_json["points_to_add"]
                 if new_points_to_add: