From be6141cce2cbdedb4b8ce5bf2830e54525df7fe7 Mon Sep 17 00:00:00 2001
From: James Krieger <jamesmkrieger@gmail.com>
Date: Thu, 11 Mar 2021 09:50:22 +0000
Subject: [PATCH] added trajectory wrap method

---
 prody/trajectory/trajectory.py | 35 ++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/prody/trajectory/trajectory.py b/prody/trajectory/trajectory.py
index 503ac1933..0747ecfbd 100644
--- a/prody/trajectory/trajectory.py
+++ b/prody/trajectory/trajectory.py
@@ -10,6 +10,9 @@
 from .frame import Frame
 
 from prody.trajectory import openTrajFile
+from prody.measure import wrapAtoms
+
+from prody import LOGGER
 
 __all__ = ['Trajectory']
 
@@ -311,3 +314,35 @@ def numFixed(self):
         """Returns a list of fixed atom numbers, one from each file."""
 
         return [traj.numFixed() for traj in self._trajectories]
+
+    def wrap(self, unitcell=None, center=np.array([0., 0., 0.])):
+        """Wrap atoms into an image of the system simulated under periodic boundary
+        conditions for all frames and returns a new Trajectory.
+
+        .. note::
+        This function will wrap all atoms into the specified periodic image, so
+        covalent bonds will be broken.
+
+        :arg unitcell: orthorhombic unitcell array with shape (3,)
+        :type unitcell: :class:`numpy.ndarray`
+
+        :arg center: coordinates of the center of the wrapping cell, default is
+            the origin of the Cartesian coordinate system
+        :type center: :class:`numpy.ndarray`"""
+
+        if unitcell is None:
+            unitcell = self[0].getUnitcell()[:3]
+
+        wrapped = Trajectory(self.getTitle() + '_wrapped')
+
+        coordsets = self.getCoordsets()
+
+        LOGGER.progress('Wrapping trajectory ...', len(coordsets))
+        for i, coordset in enumerate(coordsets):
+            wrapped.addCoordset(wrapAtoms(coordset), unitcell=unitcell, center=center)
+            LOGGER.update(i)
+        
+        LOGGER.finish()
+
+        return wrapped
+