Recreate structure fresh during get_primitive_cell

structuretoolkit/analyse/symmetry.py

@@ -345,25 +345,36 @@ def get_primitive_cell(
         >>> len(symmetry.get_primitive_cell()) == len(basis)
         True
         """
+        if not all(self._structure.pbc):
+            raise ValueError("Can only symmetrize periodic structures.")
         ret = spglib.standardize_cell(
             self._get_spglib_cell(use_elements=use_elements, use_magmoms=use_magmoms),
             to_primitive=not standardize,
         )
         if ret is None:
             raise SymmetryError(spglib.spglib.spglib_error.message)
-        cell, positions, indices = ret
-        positions = (cell.T @ positions.T).T
-        new_structure = self._structure.copy()
-        new_structure.cell = cell
-        new_structure = new_structure[: len(indices)]
+        cell, scaled_positions, indices = ret
         indices_dict = {
             v: k
             for k, v in structuretoolkit.common.helper.get_species_indices_dict(
                 structure=self._structure
             ).items()
         }
-        new_structure.symbols = [indices_dict[i] for i in indices]
-        new_structure.positions = positions
+        symbols = [indices_dict[i] for i in indices]
+        arrays = {
+            k: self._structure.arrays[k]
+            for k in self._structure.arrays
+            if k not in ("numbers", "positions")
+        }
+        new_structure = type(self._structure)(
+            symbols=symbols,
+            scaled_positions=scaled_positions,
+            cell=cell,
+            pbc=[True, True, True],
+        )
+        for k, a in arrays.items():
+            new_structure.arrays[k] = a
+
         return new_structure
 
     def get_ir_reciprocal_mesh(

tests/test_symmetry.py

@@ -143,7 +143,7 @@ def test_get_ir_reciprocal_mesh(self):
 
     def test_get_primitive_cell(self):
         cell = 2.2 * np.identity(3)
-        basis = Atoms("AlFe", scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)], cell=cell)
+        basis = Atoms("AlFe", scaled_positions=[(0, 0, 0), (0.5, 0.5, 0.5)], cell=cell, pbc=True)
         structure = basis.repeat([2, 2, 2])
         sym = stk.analyse.get_symmetry(structure=structure)
         self.assertEqual(len(basis), len(sym.get_primitive_cell(standardize=True)))
@@ -169,7 +169,7 @@ def test_get_primitive_cell_hex(self):
             [0.77, 1.57, 5.74],
         ]
         cell = [[2.519, 1.454, 4.590], [-2.519, 1.454, 4.590], [0.0, -2.909, 4.590]]
-        structure = Atoms(symbols=elements, positions=positions, cell=cell)
+        structure = Atoms(symbols=elements, positions=positions, cell=cell, pbc=True)
         structure_repeat = structure.repeat([2, 2, 2])
         sym = stk.analyse.get_symmetry(structure=structure_repeat)
         structure_prim_base = sym.get_primitive_cell()