python: added folded pos to ase

Author: reinaual

src/python/espressomd/plugins/ase.py

@@ -45,10 +45,10 @@ def register_system(self, system):
     def __getstate__(self):
         return {"type_mapping": self.type_mapping}
 
-    def get(self) -> ase.Atoms:
+    def get(self, folded=False) -> ase.Atoms:
         """Export the ESPResSo system particle data to an ASE atoms object."""
         particles = self._system.part.all()
-        positions = np.copy(particles.pos)
+        positions = np.copy(particles.pos_folded if folded else particles.pos)
         types = np.copy(particles.type)
         forces = np.copy(particles.f)
         unknown_types = set(types) - set(self.type_mapping)

src/python/espressomd/zn.py

@@ -66,7 +66,7 @@ def encode(self, system) -> zndraw.utils.ASEDict:
 
         self.numbers = self.num_particles * [1]
 
-        if self.params["folded"] is True:
+        if self.params["folded"]:
             self.positions = self.particles.pos_folded
         else:
             self.positions = self.particles.pos
@@ -81,7 +81,7 @@ def encode(self, system) -> zndraw.utils.ASEDict:
         else:
             self.radii = self.set_radii(self.params["radii"])
 
-        if self.params["bonds"] is True:
+        if self.params["bonds"]:
             bonds = self.get_bonds()
         else:
             bonds = []
@@ -542,7 +542,7 @@ def update(self):
         Asedata = ase.ASEInterface(
             {x: "X" for x in set(all_types)})
         Asedata.register_system(self.system)
-        data = Asedata.get()
+        data = Asedata.get(folded=self.params["folded"])
         if self.params["colors"] is not None:
             data.arrays['colors'] = [self.params["colors"].get(
                 z, "white") for z in all_types]

testsuite/python/ase_interface.py

@@ -33,6 +33,7 @@ class ASEInterfaceTest(ut.TestCase):
     def setUp(self):
         self.system.part.add(pos=[0., 0., 0.], f=[1., -1., 0.], type=0)
         self.system.part.add(pos=[0., 0., 1.], f=[0., 12., 0.], type=1)
+        self.system.part.add(pos=[11., 13., 12.], f=[0., 0., -8.], type=1)
         self.system.ase = espressomd.plugins.ase.ASEInterface(
             type_mapping={0: "H", 1: "O"},
         )
@@ -43,15 +44,18 @@ def tearDown(self):
     def test_ase_get(self):
         """Test the ``ASEInterface.get()`` method."""
         # Create a simple ASE atoms object
-        atoms = self.system.ase.get()
-        self.assertIsInstance(atoms, ase.Atoms)
-        self.assertEqual(set(atoms.get_chemical_symbols()), {"H", "O"})
-        np.testing.assert_equal(atoms.pbc, np.copy(self.system.periodicity))
-        np.testing.assert_allclose(atoms.cell, np.diag(self.system.box_l))
-        np.testing.assert_allclose(atoms.get_positions(),
-                                   [[0., 0., 0.], [0., 0., 1.]])
-        np.testing.assert_allclose(atoms.get_forces(),
-                                   [[1., -1., 0.], [0., 12., 0.]])
+        for folded in [True, False]:
+            atoms = self.system.ase.get(folded=folded)
+            self.assertIsInstance(atoms, ase.Atoms)
+            self.assertEqual(set(atoms.get_chemical_symbols()), {"H", "O"})
+            np.testing.assert_equal(
+                atoms.pbc, np.copy(self.system.periodicity))
+            np.testing.assert_allclose(atoms.cell, np.diag(self.system.box_l))
+            positions_ref = self.system.part.all(
+            ).pos_folded if folded else self.system.part.all().pos
+            np.testing.assert_allclose(atoms.get_positions(), positions_ref)
+            np.testing.assert_allclose(atoms.get_forces(),
+                                       [[1., -1., 0.], [0., 12., 0.], [0., 0., -8.]])
 
     @utx.skipIfMissingFeatures("VIRTUAL_SITES_RELATIVE")
     def test_exceptions(self):