Made segmentmodel_from_lists reader create splines from more than two points

mrdna/readers/segmentmodel_from_lists.py

@@ -81,6 +81,22 @@ def SegmentModelFromPdb(*args,**kwargs):
 
     ## Build map from residue index to helix index
 
+def set_splines(seg, coordinates, hid, hmap, hrank):
+    maxrank = np.max( hrank[hmap==hid] )
+    if maxrank == 0:
+        ids = np.where((hmap == hid))[0]
+        pos = np.mean( [coordinates[r,:] for r in ids ], axis=0 )
+        coords = [pos,pos]
+        contours = [0,1]
+    else:
+        coords,contours = [[],[]]
+        for rank in range(int(maxrank)+1):
+            ids = np.where((hmap == hid) * (hrank == rank))[0]
+            coords.append(np.mean( [coordinates[r,:] for r in ids ], axis=0 ))
+            contours.append( float(rank)/maxrank )
+    coords = np.array(coords)
+    seg.set_splines(coords,contours)
+
 
 def model_from_basepair_stack_3prime(coordinates, basepair, stack, three_prime, sequence=None):
     """ 
@@ -160,32 +176,20 @@ def model_from_basepair_stack_3prime(coordinates, basepair, stack, three_prime,
     ## Create double-stranded segments
     doubleSegments = []
     for hid in double_stranded_helices:
-        maxrank = np.max( hrank[hmap==hid] )
-        start = np.where((hmap == hid) * (hrank == 0))[0]
-        end = np.where((hmap == hid) * (hrank == maxrank))[0]
-
-        start_pos = np.mean( [coordinates[r,:] for r in start], axis=0 )
-        end_pos = np.mean( [coordinates[r,:] for r in end], axis=0 )
         seg = DoubleStrandedSegment(name=str(hid),
-                                    num_bp = np.sum(hmap==hid)//2,
-                                    start_position = start_pos,
-                                    end_position = end_pos)
+                                    num_bp = np.sum(hmap==hid)//2)
+        set_splines(seg, coordinates, hid, hmap, hrank)
+
         assert(hid == len(doubleSegments))
         doubleSegments.append(seg)
 
     ## Create single-stranded segments
     singleSegments = []
     for hid in single_stranded_helices:
-        maxrank = np.max( hrank[hmap==hid] )
-        start = np.where((hmap == hid) * (hrank == 0))[0]
-        end = np.where((hmap == hid) * (hrank == maxrank))[0]
-
-        start_pos = np.mean( [coordinates[r,:] for r in start], axis=0 )
-        end_pos = np.mean( [coordinates[r,:] for r in end], axis=0 )
         seg = SingleStrandedSegment(name=str(hid),
-                                    num_nt = np.sum(hmap==hid),
-                                    start_position = start_pos,
-                                    end_position = end_pos)
+                                    num_nt = np.sum(hmap==hid))
+        set_splines(seg, coordinates, hid, hmap, hrank)
+
         assert(hid == len(doubleSegments) + len(singleSegments))
         singleSegments.append(seg)
 
@@ -230,6 +234,7 @@ def model_from_basepair_stack_3prime(coordinates, basepair, stack, three_prime,
                           local_twist = False,
                           dimensions=(5000,5000,5000))
 
+
     if sequence is None:
         model.randomize_unset_sequence()
 

mrdna/segmentmodel.py

@@ -385,12 +385,14 @@ class Segment(ConnectableElement, Group):
     def __repr__(self):
         return "<{} {}[{:d}]>".format( type(self), self.name, self.num_nt )
 
+    def set_splines(self, coords, contours):
+        tck, u = interpolate.splprep( coords.T, u=contours, s=0, k=1)
+        self.position_spline_params = tck
 
     def _set_splines_from_ends(self):
-        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.quaternion_spline_params = None
+        self.set_splines(np.array([self.start_position,
+                                   self.end_position]), [0,1])
 
 
     def clear_all(self):