fix: copy AtomicSystem arrays and carry energy, virial, stress

PQAnalysis/atomic_system/atomic_system.py

@@ -561,10 +561,13 @@ def copy(self) -> "AtomicSystem":
             A copy of the AtomicSystem.
         """
         return AtomicSystem(
-            pos=self.pos,
-            vel=self.vel,
-            forces=self.forces,
-            charges=self.charges,
+            pos=self.pos.copy(),
+            vel=self.vel.copy(),
+            forces=self.forces.copy(),
+            charges=self.charges.copy(),
+            energy=self.energy,
+            virial=None if self.virial is None else self.virial.copy(),
+            stress=None if self.stress is None else self.stress.copy(),
             cell=self.cell,
             topology=self.topology
         )

tests/atomicSystem/test_atomic_system.py

@@ -758,6 +758,62 @@ def test_stress(self):
         assert not system.has_stress
         assert system.stress is None
 
+    def test_copy(self):
+        """
+        Test that copy carries all fields and does not alias the arrays.
+        """
+        system = AtomicSystem(
+            atoms=[Atom('C'), Atom('H')],
+            pos=np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]),
+            vel=np.array([[0.1, 0.0, 0.0], [0.2, 0.0, 0.0]]),
+            forces=np.array([[1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]),
+            charges=np.array([-1.0, 1.0]),
+            energy=-123.456,
+            virial=np.eye(3),
+            stress=2 * np.eye(3),
+            cell=Cell(10, 10, 10),
+        )
+
+        copy = system.copy()
+
+        assert np.isclose(copy.energy, -123.456)
+        assert np.allclose(copy.virial, np.eye(3))
+        assert np.allclose(copy.stress, 2 * np.eye(3))
+        assert copy.cell == system.cell
+        assert copy.topology == system.topology
+
+        assert copy.pos is not system.pos
+        assert copy.vel is not system.vel
+        assert copy.forces is not system.forces
+        assert copy.charges is not system.charges
+        assert copy.virial is not system.virial
+        assert copy.stress is not system.stress
+
+        copy.center(np.array([5.0, 0.0, 0.0]), image=False)
+        copy.vel[0] = [9.0, 9.0, 9.0]
+        copy.forces[0] = [9.0, 9.0, 9.0]
+        copy.charges[0] = 9.0
+        copy.virial[0, 0] = 9.0
+        copy.stress[0, 0] = 9.0
+
+        assert np.allclose(
+            system.pos, np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
+        )
+        assert np.allclose(
+            system.vel, np.array([[0.1, 0.0, 0.0], [0.2, 0.0, 0.0]])
+        )
+        assert np.allclose(
+            system.forces, np.array([[1.0, 0.0, 0.0], [2.0, 0.0, 0.0]])
+        )
+        assert np.allclose(system.charges, np.array([-1.0, 1.0]))
+        assert np.allclose(system.virial, np.eye(3))
+        assert np.allclose(system.stress, 2 * np.eye(3))
+
+        empty_copy = AtomicSystem().copy()
+        assert empty_copy.energy is None
+        assert empty_copy.virial is None
+        assert empty_copy.stress is None
+
     def test_center_of_mass_resiudes(self):
         system = AtomicSystem()