Clarify argument shift for SlabGenerator.get_slab

pymatgen/core/surface.py

@@ -62,12 +62,12 @@
 
 class Slab(Structure):
     """Class to hold information for a Slab, with additional
-    attributes pertaining to slabs, but the init method does not
-    actually create a slab. Also has additional methods that returns other information
-    about a Slab such as the surface area, normal, and atom adsorption.
+    attributes pertaining to slabs, but does not actually create a slab.
+    Also has additional methods for a Slab such as the surface area,
+    normal, and adsorbate atoms.
 
-    Note that all Slabs have the surface normal oriented perpendicular to the a
-    and b lattice vectors. This means the lattice vectors a and b are in the
+    Note that all Slabs have the surface normal oriented perpendicular to the
+    a and b lattice vectors. This means the lattice vectors a and b are in the
     surface plane and the c vector is out of the surface plane (though not
     necessarily perpendicular to the surface).
     """
@@ -221,7 +221,7 @@ def from_dict(cls, dct: dict[str, Any]) -> Self:  # type: ignore[override]
             dct: dict.
 
         Returns:
-            Creates slab from dict.
+            Slab: Created from dict.
         """
         lattice = Lattice.from_dict(dct["lattice"])
         sites = [PeriodicSite.from_dict(sd, lattice) for sd in dct["sites"]]
@@ -329,8 +329,8 @@ def get_surface_sites(self, tag: bool = False) -> dict[str, list]:
         site as well. This will only work for single-element systems for now.
 
         Args:
-            tag (bool): Option to adds site attribute "is_surfsite" (bool)
-                to all sites of slab. Defaults to False
+            tag (bool): Add attribute "is_surf_site" (bool)
+                to all sites of the Slab. Defaults to False.
 
         Returns:
             A dictionary grouping sites on top and bottom of the slab together.
@@ -405,8 +405,6 @@ def get_symmetric_site(
         symmetric properties of a slab when creating adsorbed
         structures or symmetric reconstructions.
 
-        TODO (@DanielYang59): use "site" over "point" as arg name for consistency
-
         Args:
             point (ArrayLike): Fractional coordinate of the original site.
             cartesian (bool): Use Cartesian coordinates.
@@ -424,7 +422,7 @@ def get_symmetric_site(
         for op in ops:
             slab = self.copy()
             site_other = op.operate(point)
-            if f"{site_other[2]:.6f}" == f"{point[2]:.6f}":
+            if isclose(site_other[2], point[2], abs_tol=1e-6):
                 continue
 
             # Add dummy sites to check if the overall structure is symmetric
@@ -642,19 +640,15 @@ def symmetrically_add_atom(
         specie: str | Element | Species | None = None,
         coords_are_cartesian: bool = False,
     ) -> None:
-        """Add a species at a specified site in a slab. Will also add an
-        equivalent site on the other side of the slab to maintain symmetry.
-
-        TODO (@DanielYang59): use "site" over "point" as arg name for consistency
+        """Add a species at a selected site in a Slab. Will also add an
+        equivalent site on the other side to maintain symmetry.
 
         Args:
             species (str | Element | Species): The species to add.
             point (ArrayLike): The coordinate of the target site.
             specie: Deprecated argument name in #3691. Use 'species' instead.
             coords_are_cartesian (bool): If the site is in Cartesian coordinates.
         """
-        # For now just use the species of the surface atom as the element to add
-
         # Check if deprecated argument is used
         if specie is not None:
             warnings.warn("The argument 'specie' is deprecated. Use 'species' instead.", DeprecationWarning)
@@ -736,23 +730,12 @@ def get_equi_sites(slab: Slab, sites: list[int]) -> list[int]:
 
 
 def center_slab(slab: Slab) -> Slab:
-    """Relocate the Slab to the center such that its center
-    (the slab region) is close to z=0.5.
+    """Relocate the Slab region to the center.
 
     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.
+    TODO (@DanielYang59): need to check if the Slab is continuous
 
     TODO (@DanielYang59): this should be a method for `Slab`?
 
@@ -762,34 +745,19 @@ def center_slab(slab: Slab) -> Slab:
     Returns:
         Slab: 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)
+    # Calculate shift vector
+    center_of_mass = np.average(
+        slab.frac_coords,
+        weights=[site.species.weight for site in slab],
+        axis=0,
+    )
     shift = 0.5 - center_of_mass[2]
 
-    slab.translate_sites(all_indices, [0, 0, shift])
+    # Move the slab to the center
+    slab.translate_sites(
+        indices=list(range(len(slab))),  # all sites
+        vector=[0, 0, shift],
+    )
 
     return slab
 
@@ -805,7 +773,7 @@ def get_slab_regions(
 
     TODO (@DanielYang59): this should be a method for `Slab`?
 
-    TODO (@DanielYang59): maybe project all z coordinates to 1D?
+    TODO (@DanielYang59): project all z coordinates to 1D?
 
     Args:
         slab (Slab): The Slab to analyse.
@@ -822,18 +790,18 @@ def get_slab_regions(
         neighbors = slab.get_neighbors(site, blength)
         for nn in neighbors:
             # TODO (@DanielYang59): use z coordinate (z<0) to check
-            # if a Slab is contiguous is suspicious (Slab could locate
+            # whether a Slab being continuous is suspicious (Slab could locate
             # entirely below z=0)
 
-            # Find sites with z < 0 (sites noncontiguous within cell)
+            # Find sites with z < 0 (sites non-continuous within cell)
             if nn[0].frac_coords[2] < 0:
                 frac_coords.append(nn[0].frac_coords[2])
                 indices.append(nn[-2])
 
                 if nn[-2] not in all_indices:
                     all_indices.append(nn[-2])
 
-    # If slab is noncontiguous
+    # If slab is non-continuous
     if frac_coords:
         # Locate the lowest site within the upper Slab
         last_fcoords = []
@@ -847,7 +815,7 @@ def get_slab_regions(
             frac_coords, indices = [], []
             for nn in neighbors:
                 if 1 > nn[0].frac_coords[2] > 0 and nn[0].frac_coords[2] < site.frac_coords[2]:
-                    # Sites are noncontiguous within cell
+                    # Sites are non-continuous within cell
                     frac_coords.append(nn[0].frac_coords[2])
                     indices.append(nn[-2])
                     if nn[-2] not in all_indices:
@@ -1083,7 +1051,8 @@ def get_slab(self, shift: float = 0, tol: float = 0.1, energy: float | None = No
             intended for other generation methods.
 
         Args:
-            shift (float): The shift value along the lattice c direction in Angstrom.
+            shift (float): The shift value along the lattice c direction
+                in fractional coordinates.
             tol (float): Tolerance to determine primitive cell.
             energy (float): The energy to assign to the slab.
 
@@ -1112,10 +1081,8 @@ def get_slab(self, shift: float = 0, tol: float = 0.1, energy: float | None = No
         species = self.oriented_unit_cell.species_and_occu
 
         # Shift all atoms
-        # DEBUG(@DanielYang59): shift value in Angstrom inconsistent with frac_coordis
         frac_coords = self.oriented_unit_cell.frac_coords
-        # DEBUG(@DanielYang59): suspicious shift direction (positive for downwards shift)
-        frac_coords = np.array(frac_coords) + np.array([0, 0, -shift])[None, :]
+        frac_coords = np.array(frac_coords) + np.array([0, 0, shift])[None, :]
         frac_coords -= np.floor(frac_coords)  # wrap frac_coords to the [0, 1) range
 
         # Scale down z-coordinate by the number of layers
@@ -1239,8 +1206,7 @@ def gen_possible_shifts(ftol: float) -> list[float]:
 
             # Clustering does not work when there is only one atom
             if n_atoms == 1:
-                # TODO (@DanielYang59): why this magic number 0.5?
-                shift = frac_coords[0][2] + 0.5
+                shift = -frac_coords[0][2] + 0.5  # shift to center
                 return [shift - math.floor(shift)]
 
             # Compute a Cartesian z-coordinate distance matrix
@@ -1281,7 +1247,9 @@ def gen_possible_shifts(ftol: float) -> list[float]:
 
             return sorted(shifts)
 
-        def get_z_ranges(bonds: dict[tuple[Species | Element, Species | Element], float]) -> list[tuple[float, float]]:
+        def get_z_ranges(
+            bonds: dict[tuple[Species | Element, Species | Element], float], tol: float
+        ) -> list[tuple[float, float]]:
             """Collect occupied z ranges where each z_range is a (lower_z, upper_z) tuple.
 
             This method examines all sites in the oriented unit cell (OUC)
@@ -1314,15 +1282,13 @@ def get_z_ranges(bonds: dict[tuple[Species | Element, Species | Element], float]
                                     z_ranges.extend([(0, z_range[1]), (z_range[0] + 1, 1)])
 
                                 # Neglect overlapping positions
-                                elif z_range[0] != z_range[1]:
-                                    # TODO (@DanielYang59): use the following for equality check
-                                    # elif not isclose(z_range[0], z_range[1], abs_tol=tol):
+                                elif not isclose(z_range[0], z_range[1], abs_tol=tol):
                                     z_ranges.append(z_range)
 
             return z_ranges
 
         # Get occupied z_ranges
-        z_ranges = [] if bonds is None else get_z_ranges(bonds)
+        z_ranges = [] if bonds is None else get_z_ranges(bonds, tol=tol)
 
         slabs = []
         for shift in gen_possible_shifts(ftol=ftol):
@@ -1472,7 +1438,7 @@ def move_to_other_side(
 
         # Calculate the moving distance as the fractional height
         # of the Slab inside the cell
-        # DEBUG(@DanielYang59): the use actually sizes for slab/vac
+        # DEBUG(@DanielYang59): use actual sizes for slab/vac
         # instead of the input arg (min_slab/vac_size)
         n_layers_slab: int = math.ceil(self.min_slab_size / height)
         n_layers_vac: int = math.ceil(self.min_vac_size / height)
@@ -1691,17 +1657,13 @@ def generate_all_slabs(
 
 
 def get_d(slab: Slab) -> float:
-    """Determine the z-spacing between the bottom two layers for a Slab.
-
-    TODO (@DanielYang59): this should be private/internal to ReconstructionGenerator?
-    """
+    """Determine the z-spacing between the bottom two layers for a Slab."""
     # Sort all sites by z-coordinates
     sorted_sites = sorted(slab, key=lambda site: site.frac_coords[2])
 
     distance = None
     for site, next_site in zip(sorted_sites, sorted_sites[1:]):
         if not isclose(site.frac_coords[2], next_site.frac_coords[2], abs_tol=1e-6):
-            # DEBUG (@DanielYang59): code will break if no distinguishable layers found
             distance = next_site.frac_coords[2] - site.frac_coords[2]
             break