Added rotate and translate operations to segments

mrdna/segmentmodel.py

@@ -398,11 +398,63 @@ class Segment(ConnectableElement, Group):
 
     def set_splines(self, contours, coords):
         tck, u = interpolate.splprep( coords.T, u=contours, s=0, k=1)
-        self.position_spline_params = tck
+        self.position_spline_params = (tck,u)
 
     def set_orientation_splines(self, contours, quaternions):
         tck, u = interpolate.splprep( quaternions.T, u=contours, s=0, k=1)
-        self.quaternion_spline_params = tck
+        self.quaternion_spline_params = (tck,u)
+
+    def get_center(self):
+        tck, u = self.position_spline_params
+        return np.mean(self.contour_to_position(u), axis=0)
+
+    def __filter_contours(contours, positions, position_filter, contour_filter):
+        u = contours
+        r = positions
+
+        ## Filter
+        ids = list(range(len(u)))
+        if contour_filter is not None:
+            ids = list(filter(lambda i: contour_filter(u[i]), ids))
+        if position_filter is not None:
+            ids = list(filter(lambda i: position_filter(r[i,:]), ids))
+        return ids
+
+    def translate(self, translation_vector, about=None, position_filter=None, contour_filter=None):
+        dr = np.array(translation_vector)
+        tck, u = self.position_spline_params
+        r = self.contour_to_position(u)
+
+        ids = Segment.__filter_contours(u, r, position_filter, contour_filter)
+
+        ## Translate
+        r[ids,:] = r[ids,:] + dr[np.newaxis,:]
+        self.set_position_splines(r,u)
+
+    def rotate(self, rotation_matrix, about=None, position_filter=None, contour_filter=None):
+        tck, u = self.position_spline_params
+        r = self.contour_to_position(u)
+
+        ids = Segment.__filter_contours(u, r, position_filter, contour_filter)
+
+        if about is None:
+            ## TODO: do this more efficiently
+            r[ids,:] = np.array([rotation_matrix.dot(r[i,:]) for i in range(r.shape[0])])
+        else:
+            dr = np.array(about)
+            ## TODO: do this more efficiently
+            r[ids,:] = np.array([rotation_matrix.dot(r[i,:]-dr) + dr for  i in range(r.shape[0])])
+
+        self.set_position_splines(r,u)
+
+        if self.quaternion_spline_params is not None:
+            ## TODO: performance: don't shift between quaternion and matrix representations so much
+            tck, u = self.quaternion_spline_params
+            orientations = [self.contour_to_orientation(v)) for v in u]
+            for i in ids:
+                orientations[i,:] = rotation_matrix.dot(orientations[i])
+            quats = [quaternion_from_matrix(o) for o in orientations]
+            self.set_orientation_splines(quats,u)
 
     def _set_splines_from_ends(self):
         self.quaternion_spline_params = None
@@ -426,12 +478,12 @@ class Segment(ConnectableElement, Group):
         return (nt_pos+0.5)/(self.num_nt)
 
     def contour_to_position(self,s):
-        p = interpolate.splev( s, self.position_spline_params )
+        p = interpolate.splev( s, self.position_spline_params[0] )
         if len(p) > 1: p = np.array(p).T
         return p
 
     def contour_to_tangent(self,s):
-        t = interpolate.splev( s, self.position_spline_params, der=1 )
+        t = interpolate.splev( s, self.position_spline_params[0], der=1 )
         t = (t / np.linalg.norm(t,axis=0))
         return t.T
         
@@ -456,7 +508,7 @@ class Segment(ConnectableElement, Group):
             orientation = rotationAboutAxis( axis, self.twist_per_nt*self.contour_to_nt_pos(s), normalizeAxis=False )
             orientation = orientation.dot(orientation0)
         else:
-            q = interpolate.splev( s, self.quaternion_spline_params )
+            q = interpolate.splev( s, self.quaternion_spline_params[0] )
             if len(q) > 1: q = np.array(q).T # TODO: is this needed?
             orientation = quaternion_to_matrix(q)
 
@@ -885,7 +937,7 @@ class DoubleStrandedSegment(Segment):
         ## Set up interpolation for azimuthal angles 
         a = np.array([self.start_position,self.end_position]).T
         tck, u = interpolate.splprep( a, u=[0,1], s=0, k=1)
-        self.position_spline_params = tck
+        self.position_spline_params = (tck,u)
         
         ## TODO: initialize sensible spline for orientation
 
@@ -1600,7 +1652,7 @@ class SegmentModel(ArbdModel):
             #         if ret[2] > 0:
             #             pdb.set_trace()
 
-            s.position_spline_params = tck
+            s.position_spline_params = (tck,ret[1][0])
 
             """ Get twist """
             cb = [e for e in zip(contours,beads) if 'orientation_bead' in e[1].__dict__]
@@ -1634,7 +1686,7 @@ class SegmentModel(ArbdModel):
                 quats = np.array(quats)
                 # tck, u = interpolate.splprep( quats.T, u=contours, s=3, k=3 ) ;# cubic spline not as good
                 tck, u = interpolate.splprep( quats.T, u=contours, s=0, k=1 )
-                s.quaternion_spline_params = tck
+                s.quaternion_spline_params = (tck,u)
 
     def _generate_bead_model(self,
                              max_basepairs_per_bead = 7,