Commit e6d4a78d authored by cmaffeo2's avatar cmaffeo2
Browse files

Added rotate and translate operations to segments

parent 89279732
......@@ -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,
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment