Add AtomCell.periodic_duplicate

atomlib/atomcell.py

@@ -296,6 +296,24 @@ def repeat_to_aspect(self: HasAtomCellT, plane: t.Literal['xy', 'xz', 'yz'] = 'x
         repeat[indices] = na.flatten()[min_i], nb.flatten()[min_i]
         return self.repeat(repeat)
 
+    def periodic_duplicate(self: HasAtomCellT, eps: float = 1e-5) -> HasAtomCellT:
+        """
+        Add duplicate copies of atoms near periodic boundaries.
+
+        For instance, an atom at a corner will be duplicated into 8 copies.
+        This is mostly only useful for visualization.
+        """
+        frame_save = self.get_frame()
+        self = self.to_frame('cell_box').wrap(eps=eps)
+
+        for i, ax in enumerate(('x', 'y', 'z')):
+            self = self.concat((self,
+                self.filter(polars.col(ax).abs() <= eps, frame='cell_box')
+                    .transform_atoms(AffineTransform3D.translate([1. if i == j else 0. for j in range(3)]), frame='cell_box')
+            ))
+
+        return self.to_frame(frame_save)
+
     # add frame to some HasAtoms methods
 
     @_fwd_atoms_transform

atomlib/test_atomcell.py

@@ -0,0 +1,74 @@
+
+import numpy
+from numpy.testing import assert_array_almost_equal
+import pytest
+
+from . import AtomCell, Cell
+from . import make
+from .transform import AffineTransform3D
+
+
+@pytest.fixture(scope='module')
+def cu_conv() -> AtomCell:
+    # translate to alternative basis
+    return make.fcc('Cu', 3.615, cell='conv') \
+        .transform_atoms(AffineTransform3D.translate(x=0.5), frame='cell_frac') \
+        .wrap().round_near_zero()
+
+
+@pytest.mark.parametrize(('in_coords', 'out_coords'), [
+    # corner
+    ([[0., 0., 0.]], [
+        [0., 0., 0.], [1., 0., 0.], [0., 1., 0.], [1., 1., 0.],
+        [0., 0., 1.], [1., 0., 1.], [0., 1., 1.], [1., 1., 1.],
+    ]),
+    # edge
+    ([[0.5, 0., 0.]], [
+        [0.5, 0., 0.], [0.5, 1., 0.], [0.5, 0., 1.], [0.5, 1., 1.],
+    ]),
+    # face, wrap + tolerancing
+    ([[0.5, 0.5, 1.-1e-4], [0.5, 0.5, 1.1e-3]], [
+        [0.5, 0.5, -1e-4], [0.5, 0.5, 1.1e-3], [0.5, 0.5, 1.-1e-4],
+    ]),
+])
+def test_periodic_duplicate(in_coords, out_coords):
+    cell = Cell.from_unit_cell([3.5, 4.0, 4.5], n_cells=(2, 3, 4), pbc=True)
+
+    xs, ys, zs = numpy.array(in_coords).T
+    elems = [29] * len(xs)
+
+    atomcell = AtomCell({
+        'x': xs, 'y': ys, 'z': zs, 'elem': elems
+    }, cell=cell, frame='cell_box')
+
+    duplicated = atomcell.periodic_duplicate(eps=1e-3)
+
+    assert_array_almost_equal(
+        duplicated.coords(frame='cell_box'), numpy.array(out_coords),
+        decimal=6
+    )
+
+
+def test_periodic_duplicate_cu(cu_conv):
+    duplicated = cu_conv.periodic_duplicate()
+
+    assert_array_almost_equal(
+        duplicated.coords(frame='cell_frac'), numpy.array([
+            [0.5, 0.0, 0.0],
+            [0.5, 0.5, 0.5],
+            [0.0, 0.0, 0.5],
+            [0.0, 0.5, 0.0],
+
+            [1.0, 0.0, 0.5],  # x
+            [1.0, 0.5, 0.0],
+
+            [0.5, 1.0, 0.0],  # y
+            [0.0, 1.0, 0.5],
+            [1.0, 1.0, 0.5],
+
+            [0.5, 0.0, 1.0],  # z
+            [0.0, 0.5, 1.0],
+            [1.0, 0.5, 1.0],
+            [0.5, 1.0, 1.0],
+        ])
+    )