Merge pull request #210 from pyiron/symmetrize_tensor

Symmetrize tensor

Author: samwaseda

structuretoolkit/analyse/symmetry.py

@@ -9,6 +9,8 @@
 import spglib
 from scipy.spatial import cKDTree
 from typing import Optional
+import string
+from functools import cached_property
 
 import structuretoolkit.common.helper
 from structuretoolkit.common.error import SymmetryError
@@ -230,6 +232,33 @@ def symmetrize_vectors(
             np.einsum("ijk->jki", v_reshaped)[self.permutations],
         ).reshape(np.shape(vectors)) / len(self["rotations"])
 
+    def symmetrize_tensor(self, tensor: np.ndarray) -> np.ndarray:
+        """
+        Symmetrization of any tensor. The tensor is defined by a matrix with a
+        shape of `n * (n_atoms, 3)`. For example, if the structure has 100
+        atoms, the vector can have a shape of (100, 3), (100, 3, 100, 3),
+        (100, 3, 100, 3, 100, 3) etc. Additionally, you can also have an array
+        of tensors, i.e. in this example you can have a shape like (4, 100, 3)
+        or (2, 100, 3, 100, 3). When the shape is (n, n_atoms, 3), the function
+        works in the same way as `symmetrize_vectors`, which might be somewhat
+        faster.
+
+        This function can be useful for the symmetrization of Hessian tensors,
+        or any other tensors which should be symmetric.
+
+        Args:
+            tensors (ndarray): n * (n_atoms, 3) tensor to symmetrize
+
+        Returns
+            (np.ndarray) symmetrized tensor of the same shape
+        """
+        return _SymmetrizeTensor(
+            tensor=tensor,
+            structure=self._structure,
+            rotations=self.rotations,
+            permutations=self.permutations,
+        ).result
+
     def _get_spglib_cell(
         self, use_elements: Optional[bool] = None, use_magmoms: Optional[bool] = None
     ) -> tuple:
@@ -389,3 +418,73 @@ def get_ir_reciprocal_mesh(
         if mesh is None:
             raise SymmetryError(spglib.spglib.spglib_error.message)
         return mesh
+
+
+class _SymmetrizeTensor:
+    def __init__(self, tensor, structure, rotations, permutations):
+        self._tensor = np.array(tensor)
+        self._structure = structure
+        self._rotations = rotations
+        self._permutations = permutations
+
+    @cached_property
+    def order(self):
+        order = len(self._tensor.shape) // 2
+        if self._tensor.shape[-2 * order :] != order * self._structure.positions.shape:
+            raise ValueError(
+                "Tensor must have a shape of a multiple of (n_atoms, 3). See"
+                " docstring for more info"
+            )
+        return order
+
+    @cached_property
+    def ij(self):
+        return string.ascii_lowercase[: 2 * self.order]
+
+    @property
+    def IJ(self):
+        return self.ij.upper()
+
+    @property
+    def ij_reorder(self):
+        return "".join(
+            [self.ij[ii] for ii in np.arange(2 * self.order).reshape(-1, 2).T.flatten()]
+        )
+
+    @property
+    def IJ_reorder(self):
+        return "".join(
+            [self.IJ[ii] for ii in np.arange(2 * self.order).reshape(2, -1).T.flatten()]
+        )
+
+    @cached_property
+    def t_t(self):
+        return np.einsum("...{}->{}...".format(self.ij, self.ij_reorder), self._tensor)
+
+    @cached_property
+    def str_einsum(self):
+        return (
+            ",".join(
+                [I + i for i, I in zip(self.ij[-self.order :], self.IJ[-self.order :])]
+            )
+            + ","
+            + self.IJ[: self.order]
+            + self.ij[self.order :]
+            + "...->..."
+            + self.IJ_reorder
+        )
+
+    @property
+    def result(self):
+        return np.mean(
+            [
+                np.einsum(
+                    self.str_einsum,
+                    *self.order * (rot,),
+                    self.t_t[tuple(np.meshgrid(*self.order * (perm,), indexing="ij"))],
+                    optimize=True,
+                )
+                for rot, perm in zip(self._rotations, self._permutations)
+            ],
+            axis=0,
+        )

tests/test_symmetry.py

@@ -18,6 +18,7 @@
 
 try:
     import spglib
+    from structuretoolkit.analyse.symmetry import _SymmetrizeTensor
 
     skip_spglib_test = False
 except ImportError:
@@ -108,20 +109,22 @@ def test_get_symmetry(self):
             "AlAl", scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)], cell=cell, pbc=True
         )
         v = np.random.rand(6).reshape(-1, 3)
+        sym = stk.analyse.get_symmetry(structure=Al)
         self.assertAlmostEqual(
-            np.linalg.norm(
-                stk.analyse.get_symmetry(structure=Al).symmetrize_vectors(v)
-            ),
+            np.linalg.norm(sym.symmetrize_vectors(v)),
             0,
         )
         vv = np.random.rand(12).reshape(2, 2, 3)
-        for vvv in stk.analyse.get_symmetry(structure=Al).symmetrize_vectors(vv):
+        for vvv in sym.symmetrize_vectors(vv):
             self.assertAlmostEqual(np.linalg.norm(vvv), 0)
         Al.positions[0, 0] += 0.01
-        w = stk.analyse.get_symmetry(structure=Al).symmetrize_vectors(v)
+        w = sym.symmetrize_vectors(v)
         self.assertAlmostEqual(
             np.absolute(w[:, 0]).sum(), np.linalg.norm(w, axis=-1).sum()
         )
+        self.assertAlmostEqual(
+            np.linalg.norm(sym.symmetrize_vectors(v) - sym.symmetrize_tensor(v)), 0
+        )
 
     def test_get_symmetry_dataset(self):
         cell = 2.2 * np.identity(3)
@@ -155,7 +158,7 @@ def test_get_primitive_cell(self):
         )
 
     def test_get_primitive_cell_hex(self):
-        elements = ['Fe', 'Fe', 'Fe', 'Fe', 'O', 'O', 'O', 'O', 'O', 'O']
+        elements = ["Fe", "Fe", "Fe", "Fe", "O", "O", "O", "O", "O", "O"]
         positions = [
             [0.0, 0.0, 4.89],
             [0.0, 0.0, 11.78],
@@ -174,8 +177,7 @@ def test_get_primitive_cell_hex(self):
         sym = stk.analyse.get_symmetry(structure=structure_repeat)
         structure_prim_base = sym.get_primitive_cell()
         self.assertEqual(
-            structure_prim_base.get_chemical_symbols(),
-            structure.get_chemical_symbols()
+            structure_prim_base.get_chemical_symbols(), structure.get_chemical_symbols()
         )
 
     def test_get_equivalent_points(self):
@@ -284,5 +286,57 @@ def test_error(self):
             stk.analyse.get_symmetry(structure=structure)
 
 
+@unittest.skipIf(
+    skip_spglib_test, "spglib is not installed, so the spglib tests are skipped."
+)
+class TestSymmetrizeTensors(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.structure = bulk("Al", cubic=True, a=4.0).repeat(2)
+        cls.dataset = {
+            "structure": cls.structure,
+            "rotations": np.eye(3),
+            "permutations": np.arange(len(cls.structure)),
+        }
+
+    def test_order(self):
+        with self.assertRaises(ValueError):
+            _SymmetrizeTensor(
+                tensor=np.array([1]), **self.dataset
+            ).order
+        self.assertEqual(
+            _SymmetrizeTensor(
+                tensor=np.random.randn(*self.structure.positions.shape), **self.dataset
+            ).order,
+            1,
+        )
+        self.assertEqual(
+            _SymmetrizeTensor(
+                tensor=np.random.randn(*2 * self.structure.positions.shape),
+                **self.dataset,
+            ).order,
+            2,
+        )
+
+    def test_indexing(self):
+        st = _SymmetrizeTensor(
+            tensor=np.random.randn(*2 * self.structure.positions.shape), **self.dataset
+        )
+        self.assertEqual(st.ij, "abcd")
+        self.assertEqual(st.ij_reorder, "acbd")
+        self.assertEqual(st.IJ, "ABCD")
+        self.assertEqual(st.IJ_reorder, "ACBD")
+
+    def test_str_einsum(self):
+        st = _SymmetrizeTensor(
+            tensor=np.random.randn(*2 * self.structure.positions.shape), **self.dataset
+        )
+        self.assertEqual(st.str_einsum, "Cc,Dd,ABcd...->...ACBD")
+        st = _SymmetrizeTensor(
+            tensor=np.random.randn(*self.structure.positions.shape), **self.dataset
+        )
+        self.assertEqual(st.str_einsum, "Bb,Ab...->...AB")
+
+
 if __name__ == "__main__":
     unittest.main()