Add repulse() to push close atoms apart

src/structuretoolkit/build/geometry.py

@@ -0,0 +1,89 @@
+"""Utilities that operate purely geometric aspects of structures."""
+
+import numpy as np
+from ase.atoms import Atoms
+
+from structuretoolkit.analyse import get_neighbors
+
+
+def repulse(
+    structure: Atoms,
+    min_dist: float = 1.5,
+    step_size: float = 0.2,
+    axis: int | None = None,
+    iterations: int = 100,
+    inplace: bool = False,
+) -> Atoms:
+    """Iteratively displace atoms apart until all interatomic distances exceed a minimum threshold.
+
+    For each pair of atoms closer than ``min_dist``, the atom is displaced away from its nearest
+    neighbour by up to ``step_size`` along the direction of the interatomic vector.  The loop
+    repeats until all nearest-neighbour distances satisfy the minimum criterion or the iteration
+    limit is reached.
+
+    Args:
+        structure (:class:`ase.Atoms`):
+            Structure to modify.
+        min_dist (float):
+            Minimum interatomic distance (in Å) to enforce between every pair of atoms.
+            Defaults to 1.5.
+        step_size (float):
+            Maximum displacement (in Å) applied to a single atom per iteration.
+            Smaller values give smoother convergence but require more iterations.
+            Defaults to 0.2.
+        axis (int or None):
+            Cartesian axis index (0, 1, or 2) along which displacements are restricted.
+            When *None* (default) displacements are applied in all three directions.
+        iterations (int):
+            Maximum number of displacement steps before raising a :class:`RuntimeError`.
+            Defaults to 100.
+        inplace (bool):
+            If *True*, the positions of ``structure`` are modified directly.
+            If *False* (default), a copy is made and the original is left unchanged.
+
+    Returns:
+        :class:`ase.Atoms`: The structure with adjusted atomic positions.  This is the
+        same object as ``structure`` when ``inplace=True``, or a new copy otherwise.
+
+    Raises:
+        RuntimeError: If the minimum distance criterion is not satisfied within
+            ``iterations`` steps.
+    """
+    if not inplace:
+        structure = structure.copy()
+    if axis is None:
+        axis = slice(None)
+    for _ in range(iterations):
+        neigh = get_neighbors(structure, num_neighbors=1)
+        dd = neigh.distances[:, 0]
+        if dd.min() >= min_dist:
+            break
+
+        I = dd < min_dist
+
+        dd_I = dd[I]
+        vv = neigh.vecs[I, 0, :].copy()
+        # Avoid division by zero for coincident atoms (distance == 0).
+        # Assign opposite fallback directions based on atom-index ordering so
+        # the two coincident atoms separate rather than move together.
+        atom_indices = np.where(I)[0]
+        zero_mask = dd_I == 0
+        if np.any(zero_mask):
+            neighbor_indices = neigh.indices[atom_indices[zero_mask], 0]
+            sign = np.where(atom_indices[zero_mask] < neighbor_indices, 1.0, -1.0)
+            vv[zero_mask] = sign[:, None] * np.array([1.0, 0.0, 0.0])
+        safe_dd = np.where(dd_I > 0, dd_I, 1.0)
+        vv /= safe_dd[:, None]
+
+        disp = np.clip(min_dist - dd[I], 0, step_size)
+
+        displacement = disp[:, None] * vv  # (N_close, 3)
+        structure.positions[I, axis] -= displacement[:, axis]
+
+    else:
+        raise RuntimeError(f"repulse did not converge within {iterations} iterations")
+
+    return structure
+
+
+__all__ = ["repulse"]

tests/test_geometry.py

@@ -0,0 +1,61 @@
+import unittest
+
+import numpy as np
+from ase.build import bulk
+
+from structuretoolkit.analyse import get_neighbors
+from structuretoolkit.build.geometry import repulse
+
+
+class TestRepulse(unittest.TestCase):
+    def setUp(self):
+        self.atoms = bulk("Cu", cubic=True).repeat(5)
+        self.atoms.rattle(1)
+
+    def test_noop(self):
+        """If no atoms are violating min_dist, atoms should be unchanged."""
+        atoms = bulk("Cu", cubic=True).repeat(5)  # perfect FCC, no rattle
+        original_positions = atoms.positions.copy()
+        # Cu nearest-neighbor ~2.55 Å, so min_dist=2.0 triggers no displacement
+        result = repulse(atoms, min_dist=2.0)
+        np.testing.assert_array_equal(result.positions, original_positions)
+
+    def test_inplace(self):
+        """Input atoms should be copied depending on `inplace`."""
+        original_positions = self.atoms.positions.copy()
+
+        # inplace=False (default): original must be untouched, result is a copy
+        result = repulse(self.atoms, inplace=False)
+        np.testing.assert_array_equal(self.atoms.positions, original_positions)
+        self.assertIsNot(result, self.atoms)
+
+        # inplace=True: result is the same object as the input
+        result2 = repulse(self.atoms, inplace=True)
+        self.assertIs(result2, self.atoms)
+
+    def test_iterations(self):
+        """Should raise error if iterations exhausted."""
+        # min_dist=5.0 is far larger than any achievable spacing; step_size tiny
+        # → convergence is impossible, so iterations will be exhausted
+        with self.assertRaises(RuntimeError):
+            repulse(self.atoms, min_dist=5.0, step_size=0.001, iterations=2)
+
+    def test_axis(self):
+        """If axis given, the other axes should be exactly untouched."""
+        atoms = self.atoms.copy()
+        original_y = atoms.positions[:, 1].copy()
+        original_z = atoms.positions[:, 2].copy()
+        # Modify inplace so we can inspect positions even if convergence fails
+        try:
+            repulse(atoms, axis=0, inplace=True)
+        except RuntimeError:
+            pass  # convergence irrelevant; we only care which axes were touched
+        np.testing.assert_array_equal(atoms.positions[:, 1], original_y)
+        np.testing.assert_array_equal(atoms.positions[:, 2], original_z)
+
+    def test_min_dist(self):
+        """min_dist must be respected after a call to repulse."""
+        min_dist = 1.5
+        result = repulse(self.atoms, min_dist=min_dist)
+        neigh = get_neighbors(result, num_neighbors=1)
+        self.assertGreaterEqual(neigh.distances[:, 0].min(), min_dist)