Improved orientation spline updates for ssDNA

mrdna/segmentmodel.py

@@ -1749,27 +1749,93 @@ class SegmentModel(ArbdModel):
 
             s.position_spline_params = (tck,u)
 
-            """ Get twist """
-            cb = [e for e in zip(contours,beads) if 'orientation_bead' in e[1].__dict__]
-            beads = [b for c,b in cb] 
-            contours = [c for c,b in cb] 
-            ids = [b.idx for b in beads]
-            # if 'orientation_bead' in beads[0].__dict__:
-            if len(beads) > 3:
-                tangents = s.contour_to_tangent(contours)
+            """ Get orientation """
+            def get_orientation_vector(bead,tangent):
+                if 'orientation_bead' in bead.__dict__:
+                    o = bead.orientation_bead
+                    oVec = bead_coordinates[o.idx,:] - bead_coordinates[bead.idx,:]
+                    oVec = oVec - oVec.dot(tangent)*tangent
+                    oVec = oVec/np.linalg.norm(oVec)
+                else:
+                    oVec = None
+                return oVec
+
+            def remove_tangential_projection(vector, tangent):
+                """ Assume tangent is normalized """
+                v = vector - vector.dot(tangent)*tangent
+                return v/np.linalg.norm(v)
+
+            def get_orientation_vector(bead,tangent):
+                if 'orientation_bead' in bead.__dict__:
+                    o = bead.orientation_bead
+                    oVec = bead_coordinates[o.idx,:] - bead_coordinates[bead.idx,:]
+                    oVec = remove_tangential_projection(oVec,tangent)
+                else:
+                    oVec = None
+                return oVec
+
+
+            def get_previous_idx_if_none(list_):
+                previous = None
+                result = []
+                i = 0
+                for e in list_:
+                    if e is None:
+                        result.append(previous)
+                    else:
+                        previous = i
+                    i+=1
+                return result
+            def get_next_idx_if_none(list_):
+                tmp = get_previous_idx_if_none(list_[::-1])[::-1]
+                return [ len(list_)-1-idx if idx is not None else idx for idx in tmp ]
+
+            def fill_in_orientation_vectors(contours,orientation_vectors,tangents):
+                result = []
+                last_idx = get_previous_idx_if_none( orientation_vectors )
+                next_idx = get_next_idx_if_none( orientation_vectors )
+                none_idx = 0
+                for c,ov,t in zip(contours,orientation_vectors,tangents):
+                    if ov is not None:
+                        result.append(ov)
+                    else:
+                        p = last_idx[none_idx]
+                        n = next_idx[none_idx]
+                        none_idx += 1
+                        if p is None:
+                            if n is None:
+                                ## Should be quite rare; give something random if it happens
+                                print("WARNING: unable to interpolate orientation")
+                                o = np.array((1,0,0))
+                                result.append( remove_tangential_projection(o,t) )
+                            else:
+                                o = orientation_vectors[n]
+                                result.append( remove_tangential_projection(o,t) )
+                        else:
+                            if n is None:
+                                o = orientation_vectors[p]
+                                result.append( remove_tangential_projection(o,t) )
+                            else:
+                                cp,cn = [contours[i] for i in (p,n)]
+                                op,on = [orientation_vectors[i] for i in (p,n)]
+                                if (cn-cp) > 1e-6:
+                                    o = ((cn-c)*op+(c-cp)*on)/(cn-cp)
+                                else:
+                                    o = op+on
+                                result.append( remove_tangential_projection(o,t) )
+                return result
+
+            tangents = s.contour_to_tangent(contours)
+            orientation_vectors = [get_orientation_vector(b,t) for b,t in zip(beads,tangents)]
+            if len(beads) > 3 and any([e is not None for e in orientation_vectors] ):
+                orientation_vectors = fill_in_orientation_vectors(contours, orientation_vectors, tangents)
+
                 quats = []
                 lastq = None
-                for b,t in zip(beads,tangents):
-                    o = b.orientation_bead
-                    # positions
-                    # angleVec = o.position - b.position
-                    angleVec = bead_coordinates[o.idx,:] - bead_coordinates[b.idx,:]
-                    angleVec = angleVec - angleVec.dot(t)*t
-                    angleVec = angleVec/np.linalg.norm(angleVec)
-                    y = np.cross(t,angleVec)
+                for b,t,oVec in zip(beads,tangents,orientation_vectors):
+                    y = np.cross(t,oVec)
                     assert( np.abs(np.linalg.norm(y) - 1) < 1e-2 )
-                    q = quaternion_from_matrix( np.array([angleVec,y,t]).T)
-                    # q = quaternion_from_matrix( np.array([angleVec,y,t]))
+                    q = quaternion_from_matrix( np.array([oVec,y,t]).T)
                     
                     if lastq is not None:
                         if q.dot(lastq) < 0: