From 26eee69be7159cd887b505c5e0cced21ae5b45a8 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/trajbase.py   | 36 +++++++++++++++++++++++++++++++++-
 prody/trajectory/trajectory.py |  1 +
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/prody/trajectory/trajbase.py b/prody/trajectory/trajbase.py
index 42dd37639..56c0248d0 100644
--- a/prody/trajectory/trajbase.py
+++ b/prody/trajectory/trajbase.py
@@ -2,10 +2,12 @@
 """This module defines base class for trajectory handling."""
 
 from numbers import Integral
-from numpy import ndarray, unique
+from numpy import ndarray, unique, array
 
 from prody.ensemble import Ensemble
 from prody.utilities import checkCoords, checkWeights
+from prody.measure import wrapAtoms
+from prody import LOGGER
 
 from .frame import Frame
 
@@ -366,3 +368,35 @@ def hasUnitcell(self):
         """Returns **True** if trajectory has unitcell data."""
 
         pass
+
+
+    def wrap(self, unitcell=None, center=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 = Ensemble(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
\ No newline at end of file
diff --git a/prody/trajectory/trajectory.py b/prody/trajectory/trajectory.py
index 503ac1933..ef4d18e4e 100644
--- a/prody/trajectory/trajectory.py
+++ b/prody/trajectory/trajectory.py
@@ -311,3 +311,4 @@ def numFixed(self):
         """Returns a list of fixed atom numbers, one from each file."""
 
         return [traj.numFixed() for traj in self._trajectories]
+