# Author: Carsten Sachse
# Copyright: EMBL (2010 - 2018), Forschungszentrum Juelich (2019 - 2021)
# License: see license.txt for details
from collections import namedtuple
import os
from spring.csinfrastr.csdatabase import SpringDataBase, base, CtfMicrographTable, SegmentTable, HelixTable
from spring.csinfrastr.csproductivity import OpenMpi
from spring.segment2d.segment_prep import SegmentPreparationFileRead
from EMAN2 import Util, Transform
from scipy import interpolate
from sparx import model_blank, rot_shift2D
from tabulate import tabulate
import numpy as np
[docs]class SegmentInterpolation(SegmentPreparationFileRead):
[docs] def make_helixinfo_named_tuple(self):
helixinfo = namedtuple('helixinfo', 'helix_id micrograph segment_list directory coordinates ' + \
'picked_coordinates inplane_angle curvature')
return helixinfo
[docs] def convert_coordinate_line_to_xy_coordinates(self, each_coordinate_line):
col = each_coordinate_line.split()
xcoord = float(col[0])
ycoord = float(col[1])
boxsize = float(col[2])
xcenter = xcoord + boxsize / 2
ycenter = ycoord + boxsize / 2
return xcenter, ycenter
[docs] def compute_cumulative_distances_from_start_of_helix(self, x_coordinates, y_coordinates):
"""
>>> from spring.segment2d.segment import Segment
>>> x_coordinates = y_coordinates = np.arange(10)
>>> s = Segment()
>>> s.compute_cumulative_distances_from_start_of_helix(x_coordinates, y_coordinates)
array([ 0. , 1.41421356, 2.82842712, 4.24264069, 5.65685425,
7.07106781, 8.48528137, 9.89949494, 11.3137085 , 12.72792206])
>>> x_coordinates = 10 * np.ones(10)
>>> y_coordinates = np.arange(10)
>>> s.compute_cumulative_distances_from_start_of_helix(x_coordinates, y_coordinates)
array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
"""
previous_x_coord = np.roll(np.array(x_coordinates), 1)
previous_y_coord = np.roll(np.array(y_coordinates), 1)
distances = np.sqrt((x_coordinates - previous_x_coord)**2 + (y_coordinates - previous_y_coord)**2)
distances[0] = 0
distances_from_start = np.cumsum(distances)
return distances_from_start
[docs] def enter_helixinfo_into_helices_and_write_boxfiles(self, helices, micfile, overlap, helixfile, ipangle, curvature,
picked_coordinates, int_xy_coordinates):
total_helix_number = len(helices) + 1
helices = self.enter_helixinfo_into_helices(helices, micfile, overlap, helixfile, ipangle, curvature,
picked_coordinates, int_xy_coordinates)
interpolated_xcoord, interpolated_ycoord = list(zip(*int_xy_coordinates))
self.write_boxfile(np.array(interpolated_xcoord), np.array(interpolated_ycoord), self.segsizepix,
filename=helixfile)
data = zip([total_helix_number] + (len(int_xy_coordinates) - 1) * [None],
[os.path.basename(helixfile)] + (len(int_xy_coordinates) - 1) * [None],
interpolated_xcoord, interpolated_ycoord)
log_info = '\n' + tabulate(data, ['helix_id', 'helix_file', 'x-coordinate', 'y-coordinate'])
self.log.ilog(log_info)
return helices
[docs] def enter_helixinfo_into_helices(self, helices, micfile, overlap, helixfile, ipangle, curvature, picked_coordinates,
int_xy_coordinates):
helixinfo = self.make_helixinfo_named_tuple()
total_helix_number = len(helices) + 1
helices.append(helixinfo(total_helix_number, micfile, helixfile, overlap, int_xy_coordinates,
picked_coordinates, ipangle, curvature))
return helices
[docs] def fill_in_helix_info_from_coordinates(self, helices, segfile, stepsize, pixelsize, helixwidth, micfile, overlap,
mic_hindex, x_coordinates, y_coordinates):
helixfile = '{0}{1}{2}_{3:03}{4}box'.format(os.path.abspath(overlap), os.sep,
os.path.basename(os.path.splitext(segfile)[0]), mic_hindex, os.extsep)
helix_length = pixelsize * np.sqrt(
(y_coordinates[-1] - y_coordinates[0]) ** 2 + (x_coordinates[-1] - x_coordinates[0]) ** 2)
if helix_length < 3 * stepsize:
log_info = 'Helix from {0} was not further used because length of helix is '.format(helixfile) + \
'shorter than specified 3 * stepsize ' + \
'({0} Angstrom < {1} Angstrom). '.format(helix_length, 3 * stepsize) + \
'No sensible segmentation possible. '
self.log.ilog(log_info)
else:
helices = self.perform_interpolation_perturbation_removal_if_required(helices, segfile, stepsize, pixelsize,
helixwidth, micfile, overlap, helixfile, x_coordinates, y_coordinates)
return helices
[docs] def compute_new_coordinates_with_given_stepsize(self, helices, segfile, stepsize, pixelsize, helixwidth,
micfile, overlap, helixcoord):
mic_hindex = 0
for helix in helixcoord:
xycentercoord = []
for each_coordinate_line in helix.splitlines():
xcenter, ycenter = self.convert_coordinate_line_to_xy_coordinates(each_coordinate_line)
xycentercoord.append((xcenter, ycenter))
if each_coordinate_line.endswith('-2'):
x_coordinates, y_coordinates = list(zip(*xycentercoord))
helices = self.fill_in_helix_info_from_coordinates(helices, segfile, stepsize, pixelsize,
helixwidth, micfile, overlap, mic_hindex, x_coordinates, y_coordinates)
mic_hindex += 1
return helices
[docs] def get_picked_coordinates_from_database(self, micfile, db_name, session):
if session is None:
session = SpringDataBase().setup_sqlite_db(base, db_name)
this_mic = session.query(CtfMicrographTable).\
filter(CtfMicrographTable.micrograph_name.startswith(os.path.basename(os.path.splitext(micfile)[0]))).first()
helices = session.query(HelixTable).filter(HelixTable.mic_id == this_mic.id).order_by(HelixTable.id)
mic_coord = []
for each_helix in helices:
segments = session.query(SegmentTable).\
filter(SegmentTable.helix_id == each_helix.id).all()
coordinates = [(each_segment.picked_x_coordinate_A / self.pixelsize,
each_segment.picked_y_coordinate_A / self.pixelsize) for each_segment in segments]
mic_coord.append(coordinates)
return session, mic_coord
[docs] def get_picked_coordinates_from_file(self, segfile, overlap):
ifile = open(os.path.abspath(overlap) + os.sep + os.path.basename(segfile), 'r')
coordinate_lines = ifile.readlines()
ifile.close()
if segfile.endswith('box') or segfile.endswith('txt'):
helixcoord = self.read_coordinate_lines_from_eman_format_and_separate_helices(coordinate_lines)
elif segfile.endswith('star'):
helixcoord = \
self.read_coordinate_lines_from_bsoft_particle_format_and_convert_to_eman_format(coordinate_lines)
if helixcoord == []:
helixcoord = \
self.read_coordinate_lines_from_bsoft_filament_format_and_convert_to_eman_format(coordinate_lines)
return helixcoord
[docs] def assign_ids_based_on_new_frames(self, inp_stack_ids, inp_mics, mic_coord):
"""
>>> from spring.segment2d.segment import Segment
>>> s = Segment()
>>> mic_coord = [list(zip(np.arange(3), np.arange(3))), list(zip(np.arange(5), np.arange(5)))]
>>> first = s.assign_ids_based_on_new_frames([], list(range(5)), mic_coord)
>>> first #doctest: +NORMALIZE_WHITESPACE
[0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7,
0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7]
>>> s.assign_ids_based_on_new_frames(first, list(range(5)), mic_coord) #doctest: +NORMALIZE_WHITESPACE
[0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7,
0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
15, 8, 9, 10, 11, 12, 13, 14, 15, 8, 9, 10, 11, 12, 13, 14, 15, 8, 9, 10,
11, 12, 13, 14, 15, 8, 9, 10, 11, 12, 13, 14, 15]
>>> helices = [1, 2]
>>> s.assign_ids_based_on_new_frames([], list(range(5)), helices)
[0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
"""
all_coord = np.vstack(mic_coord)
stack_ids = []
if inp_stack_ids == []:
max_inp_stack_id = -1
else:
max_inp_stack_id = max(inp_stack_ids)
for each_mic_id, each_inp in enumerate(inp_mics):
for each_coord_id, each_coord in enumerate(all_coord):
stack_ids.append(each_coord_id + max_inp_stack_id + 1)
if len(inp_stack_ids) == 0:
merged_stack_ids = stack_ids
else:
merged_stack_ids = inp_stack_ids + stack_ids
return merged_stack_ids
[docs] def single_out(self, pair, stepsize, pixelsize, assigned_mics):
"""
* Function to single out individual helices from entire micrograph, re-interpolates coordinates\
according to specified stepsize,
#. Input: list of 3-tuples (micrograph file, segment file, overlapping letters), stepsize, pixelsize
#. Output: list of helices containing micrograph, segment_list, directory, coordinates, inplane_angle
#. Usage: self.helices = single_out(pair, stepsize, pixelsize)
"""
self.log.fcttolog()
helices = []
session = None
if assigned_mics is not None:
assigned_mics = dict(assigned_mics)
assigned_stack_ids = []
assigned_helix_ids = []
for (micfile, segfile, overlap) in pair:
if segfile.endswith('db'):
session, mic_coord = self.get_picked_coordinates_from_database(micfile, segfile, session)
base_micfile = os.path.basename(micfile)
if assigned_mics is not None:
inp_mics = assigned_mics[base_micfile]
else:
inp_mics = [micfile]
assigned_stack_ids = self.assign_ids_based_on_new_frames(assigned_stack_ids, inp_mics, mic_coord)
assigned_helix_ids = self.assign_ids_based_on_new_frames(assigned_helix_ids, inp_mics,
list(range(len(mic_coord))))
for each_inp_mic in inp_mics:
for mic_hindex, each_coord in enumerate(mic_coord):
segfile = os.path.splitext(each_inp_mic)[0] + os.extsep + 'box'
each_x_coord, each_y_coord = list(zip(*each_coord))
helixfile = '{0}{1}{2}_{3:03}{4}box'.format(os.path.abspath(overlap), os.sep,
os.path.basename(os.path.splitext(segfile)[0]), mic_hindex, os.extsep)
helices = self.perform_interpolation_perturbation_removal_if_required(helices, segfile,
stepsize, pixelsize, self.helixwidth, each_inp_mic, overlap, helixfile, each_x_coord,
each_y_coord)
else:
helixcoord = self.get_picked_coordinates_from_file(segfile, overlap)
helices = self.compute_new_coordinates_with_given_stepsize(helices, segfile, stepsize,
pixelsize, self.helixwidth, micfile, overlap, helixcoord)
assigned_stack_ids = None
if helices == []:
msg = 'No helices were segmented because picked helices were too short to be segmented. You ' + \
'can provide longer helices, choose a shorter step size and/or disable the \'Remove helix ends option \''
raise ValueError(msg)
return helices, assigned_stack_ids, assigned_helix_ids
[docs] def prepare_segmentation(self):
OpenMpi().setup_and_start_mpi_version_if_demanded(self.mpi_option, self.feature_set, self.cpu_count)
assigned_mics = self.validate_input()
pair = self.assign_reorganize(self.micrograph_files, self.coordinate_files)
self.helices, assigned_stack_ids, assigned_helix_ids = self.single_out(pair, self.stepsize, self.pixelsize,
assigned_mics)
return self.helices, assigned_stack_ids, assigned_helix_ids
[docs] def determine_angle_to_rotate_coordinates_to_minimal_slope(self, xcoord, ycoord, helixwidth):
if np.std(xcoord) < helixwidth/3 or np.std(ycoord) < helixwidth/3:
if np.std(xcoord) < np.std(ycoord):
mangle = 90.0
else:
mangle = 0
else:
polyline = np.polyfit((xcoord[0], xcoord[-1]), [ycoord[0], ycoord[-1]], 1)
# self.log.dlog('xcoord: {0}, ycoord: {1}, polyline: {2}'.format(xcoord, ycoord, polyline[-2]))
mangle = np.rad2deg(np.arctan(polyline[-2]))
return mangle
[docs] def rotate_coordinates_by_angle(self, xcoord, ycoord, angle_in_deg, center_x=0, center_y=0):
"""
* Function to rotate coordinates clockwise by angle_in_deg around origin\
(in correspondence to EMAN2 convention)
#. Input: x- and y-array, angle_in_deg to be rotated
#. Output: rotated x- and y-array
#. Usage: rotx, roty = rotate_coordinates_by_angle(xcoord, ycoord, angle_in_deg)
>>> import numpy as np
>>> np.arange(0, 10)
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.zeros(10)
array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
>>> from spring.segment2d.segment import Segment
>>> s = Segment()
>>> rot_x, rot_y = s.rotate_coordinates_by_angle(np.arange(10), np.zeros(10), -90.0)
>>> np.rint(rot_x)
array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
>>> np.rint(rot_y)
array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
>>> s.rotate_coordinates_by_angle(np.array([5]), np.array([0]), 180)
(array([-5.]), array([-6.123234e-16]))
>>> s.rotate_coordinates_by_angle(0, 0, 180, 4, 3)
(8.0, 6.0)
"""
c_xcoord = xcoord - center_x
c_ycoord = ycoord - center_y
angle_in_rad = np.deg2rad(angle_in_deg)
rotx = np.cos(angle_in_rad)*c_xcoord + np.sin(angle_in_rad)*c_ycoord
roty = -np.sin(angle_in_rad)*c_xcoord + np.cos(angle_in_rad)*c_ycoord
rotx += center_x
roty += center_y
return rotx, roty
[docs] def write_boxfile(self, interpolated_xcoord=None, interpolated_ycoord=None, segsizepix=None, filename=None):
"""
* Function to write out new boxparameters
#. Input: interpolated x- and y-array, segment size (pixel), filename
#. Output: written file filename
#. Usage: write_boxfile(interpolated_xcoord, interpolated_ycoord, \
segsizepix, filename)
"""
interpolated_xcoord = interpolated_xcoord - segsizepix/2
interpolated_ycoord = interpolated_ycoord - segsizepix/2
ofile = open(filename, 'a')
for index, xcoord in enumerate(interpolated_xcoord):
interpolated_xcoord[index], interpolated_ycoord[index]
fivecol = 0
if index == 0:
fivecol = -1
elif index == (len(interpolated_xcoord) - 1):
fivecol = -2
ofile.write('{0}\t{1}\t{2}\t{3}\t{4}\n'.format(int(round(interpolated_xcoord[index])),
int(round(interpolated_ycoord[index])), segsizepix, segsizepix, fivecol))
ofile.close()
return filename
[docs] def compute_persistence_length_m_from_coordinates_A(self, coord_x_A, coord_y_A):
"""
>>> from spring.segment2d.segment import Segment
>>> s = Segment()
>>> coord_x_A = np.arange(20)
>>> coord_y_A = coord_x_A ** 2
>>> s.compute_persistence_length_m_from_coordinates_A(coord_x_A, coord_y_A)
5.146314847399703e-06
>>> coord_y_A = coord_x_A ** 1
>>> s.compute_persistence_length_m_from_coordinates_A(coord_x_A, coord_y_A)
1.0
>>> coord_y_A = 10 * np.sin(coord_x_A)
>>> s.compute_persistence_length_m_from_coordinates_A(coord_x_A, coord_y_A)
0.0
"""
length = 1e-10 * np.sqrt((coord_x_A[0] - coord_x_A[-1]) ** 2 + (coord_y_A[0] - coord_y_A[-1]) ** 2)
contour = 1e-10 * self.compute_cumulative_distances_from_start_of_helix(coord_x_A, coord_y_A)[-1]
ratio = length / contour
if ratio <= 1 / np.sqrt(2):
pers_length_m = 0.0
elif 1 / np.sqrt(2) < ratio and ratio < 1.0:
pers_length_m = 1 / ((-np.log(2 * (length / contour) ** 2 - 1) / contour))
pers_length_m = np.min([1.0, pers_length_m])
elif ratio >= 1.0:
pers_length_m = 1.0
return pers_length_m
[docs] def interpolate_coordinates(self, xcoord, ycoord, pixelsize, stepsize, helixwidth, segfile, new_stepsize=True):
"""
* Function to interpolate x and y coordinates according to new stepsize
#. Input: xarray and yarray of segment coordinates, pixelsize, stepsize \
and helixwidth in Angstrom
#. Output: interpolated x and y-array, computed in-plane angle
#. Usage: interpolated_xcoord, interpolated_ycoord, ipangle = \
interpolate_coordinates(xcoord, ycoord, pixelsize, stepsize, helixwidth, segfile)
>>> import numpy as np
>>> from spring.segment2d.segment import Segment
>>> s = Segment()
>>> i_xcoord, i_ycoord, in_plane_angle, curvature = \
s.interpolate_coordinates(np.arange(10), np.arange(10), 1, 0.5, 3, 'test.box')
>>> np.round(i_xcoord, 2)
array([-0. , 0.35, 0.71, 1.06, 1.41, 1.77, 2.12, 2.47, 2.83,
3.18, 3.54, 3.89, 4.24, 4.6 , 4.95, 5.3 , 5.66, 6.01,
6.36, 6.72, 7.07, 7.42, 7.78, 8.13, 8.49, 8.84])
>>> np.round(i_ycoord, 2)
array([0. , 0.35, 0.71, 1.06, 1.41, 1.77, 2.12, 2.47, 2.83, 3.18, 3.54,
3.89, 4.24, 4.6 , 4.95, 5.3 , 5.66, 6.01, 6.36, 6.72, 7.07, 7.42,
7.78, 8.13, 8.49, 8.84])
>>> np.round(in_plane_angle, 2)
array([135., 135., 135., 135., 135., 135., 135., 135., 135., 135., 135.,
135., 135., 135., 135., 135., 135., 135., 135., 135., 135., 135.,
135., 135., 135., 135.])
>>> interpolated_x_coordinates, interpolated_y_coordinates, in_plane_rotation_angle, curvature = \
Segment().interpolate_coordinates(np.arange(10), 2*np.arange(10), 1, 3, 3, 'test.box')
>>> np.round(interpolated_x_coordinates, 1)
array([-0. , 1.3, 2.7, 4. , 5.4, 6.7, 8. ])
>>> np.round(interpolated_y_coordinates, 1)
array([ 0. , 2.7, 5.4, 8. , 10.7, 13.4, 16.1])
>>> in_plane_rotation_angle
array([153.43494882, 153.43494882, 153.43494882, 153.43494882,
153.43494882, 153.43494882, 153.43494882])
>>> # helix parallel to image rows
...
>>> interpolated_x_coordinates, interpolated_y_coordinates, in_plane_rotation_angle, curvature = \
Segment().interpolate_coordinates(np.arange(10), np.zeros(10), 1, 3, 3, 'test.box')
>>> interpolated_x_coordinates
array([0., 3., 6.])
>>> np.round(interpolated_y_coordinates, 2)
array([0., 0., 0.])
>>> in_plane_rotation_angle
array([90., 90., 90.])
>>> # helix parallel to image columns
...
>>> interpolated_x_coordinates, interpolated_y_coordinates, in_plane_rotation_angle, curvature = \
Segment().interpolate_coordinates(np.zeros(10), np.arange(10), 10, 20, 30, 'test.box')
>>> np.round(interpolated_y_coordinates, 0)
array([0., 2., 4., 6., 8.])
>>> in_plane_rotation_angle
array([180., 180., 180., 180., 180.])
>>> np.round(interpolated_x_coordinates, 2)
array([-0., 0., 0., 0., 0.])
"""
xcoord_A = xcoord * pixelsize
ycoord_A = ycoord * pixelsize
flat_angle = self.determine_angle_to_rotate_coordinates_to_minimal_slope(xcoord_A, ycoord_A, helixwidth)
rotated_xcoord, rot_ycoord = self.rotate_coordinates_by_angle(xcoord_A, ycoord_A, flat_angle)
if len(xcoord_A) > 4 and np.std(rot_ycoord) < 100:
poly = np.polyfit(rotated_xcoord, rot_ycoord, 3)
rotated_ycoord = np.polyval(poly, rotated_xcoord)
else:
rotated_ycoord = rot_ycoord
# self.log.ilog('Coordinates are rotated by {0} degrees'.format(np.degrees(flat_angle)))
interpolated_rotated_xcoord = self.recalculate_xcoordinate_stepsize_to_perform_interpolation(rotated_xcoord,
stepsize, new_stepsize)
interpolated_rotated_ycoord = self.perform_interpolation_of_ycoordinates(rotated_xcoord, rotated_ycoord,
interpolated_rotated_xcoord, segfile)
interpolated_rotated_xcoord, interpolated_rotated_ycoord, fitted_1st_derivative, fitted_2nd_derivative = \
self.perform_interpolation_derivative_for_in_plane_angle_rotation(interpolated_rotated_xcoord,
interpolated_rotated_ycoord, rotated_xcoord, rotated_ycoord, segfile)
# backrotate coordinates
interpolated_xcoord_A, interpolated_ycoord_A = self.rotate_coordinates_by_angle(interpolated_rotated_xcoord,
interpolated_rotated_ycoord, -flat_angle)
ipangle = 90 + np.rad2deg(np.arctan(fitted_1st_derivative)) + flat_angle
pers_length = self.compute_persistence_length_m_from_coordinates_A(interpolated_xcoord_A, interpolated_ycoord_A)
curvature = len(interpolated_rotated_xcoord) * [pers_length]
interpolated_xcoord = interpolated_xcoord_A / pixelsize
interpolated_ycoord = interpolated_ycoord_A / pixelsize
return interpolated_xcoord, interpolated_ycoord, ipangle, curvature
[docs]class SegmentStraighten(SegmentInterpolation):
[docs] def select_coordinates_within_segment_dimensions(self, coordinates, current_center, segment_size_pix):
"""
>>> from spring.segment2d.segment import Segment
>>> Segment().select_coordinates_within_segment_dimensions(list(zip(np.arange(20.), np.arange(20., 40.))), (10, 30), 10.0)
(array([ 6., 7., 8., 9., 10., 11., 12., 13., 14.]), array([26., 27., 28., 29., 30., 31., 32., 33., 34.]))
"""
x_coordinates, y_coordinates = list(zip(*coordinates))
x_coordinates = np.array(x_coordinates)
y_coordinates = np.array(y_coordinates)
x_center, y_center = current_center
x_left_corner = x_center - segment_size_pix/2
y_left_corner = y_center - segment_size_pix/2
x_right_corner = x_center + segment_size_pix/2
y_right_corner = y_center + segment_size_pix/2
x_indices = np.arange(x_coordinates.size)
y_indices = np.arange(y_coordinates.size)
within_xlimits = x_indices[(x_coordinates > x_left_corner) & (x_coordinates < x_right_corner)]
within_ylimits = y_indices[(y_coordinates > y_left_corner) & (y_coordinates < y_right_corner)]
within_x_and_y = np.intersect1d(within_xlimits, within_ylimits)
filtered_x = np.array([x_coordinates[each_index] for each_index in within_x_and_y])
filtered_y = np.array([y_coordinates[each_index] for each_index in within_x_and_y])
return filtered_x, filtered_y
[docs] def compute_bending_path_row_wise(self, length, polyfit, angle):
"""
>>> from spring.segment2d.segment import Segment
>>> Segment().compute_bending_path_row_wise(20.0, [1/20.0, -1.0, 15.0], 0)
(array([15. , 14.05, 13.2 , 12.45, 11.8 , 11.25, 10.8 , 10.45, 10.2 ,
10.05, 10. , 10.05, 10.2 , 10.45, 10.8 , 11.25, 11.8 , 12.45,
13.2 , 14.05]), array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12.,
13., 14., 15., 16., 17., 18., 19.]), array([315. , 318.0127875 , 321.34019175, 325.0079798 ,
329.03624347, 333.43494882, 338.19859051, 343.30075577,
348.69006753, 354.28940686, 360. , 5.71059314,
11.30993247, 16.69924423, 21.80140949, 26.56505118,
30.96375653, 34.9920202 , 38.65980825, 41.9872125 ]))
>>> Segment().compute_bending_path_row_wise(20.0, [1/20.0, -1.0, 15.0], 90)
(array([20., 19., 18., 17., 16., 15., 14., 13., 12., 11., 10., 9., 8.,
7., 6., 5., 4., 3., 2., 1.]), array([15. , 14.05, 13.2 , 12.45, 11.8 , 11.25, 10.8 , 10.45, 10.2 ,
10.05, 10. , 10.05, 10.2 , 10.45, 10.8 , 11.25, 11.8 , 12.45,
13.2 , 14.05]), array([225. , 228.0127875 , 231.34019175, 235.0079798 ,
239.03624347, 243.43494882, 248.19859051, 253.30075577,
258.69006753, 264.28940686, 270. , 275.71059314,
281.30993247, 286.69924423, 291.80140949, 296.56505118,
300.96375653, 304.9920202 , 308.65980825, 311.9872125 ]))
"""
straight_rows = np.arange(float(length))
straight_cols = np.polyval(polyfit, straight_rows)
spline_coefficients = interpolate.splrep(straight_rows, straight_cols, k=3, s=10)
fitted_first_derivative = interpolate.splev(straight_rows, spline_coefficients, der=1)
inplane_angles = (np.rad2deg(np.arctan(fitted_first_derivative)) - angle)%360
cols, rows = self.rotate_coordinates_by_angle(straight_cols, straight_rows, -angle, center_x=length/2,
center_y=length/2)
return cols, rows, inplane_angles
[docs] def shift_and_rotate_image_linear(self, image, inplane_angle, each_x, each_y):
t_shift = Transform({'type':'SPIDER', 'psi':0.0, 'tx':each_x, 'ty':each_y, 'scale':1.0})
t_rot = Transform({'type':'SPIDER', 'psi':inplane_angle, 'tx':0.0, 'ty':0.0, 'scale':1.0})
composite_transform = t_rot * t_shift
trans_rot_image = image.rot_scale_trans_background(composite_transform)
return trans_rot_image
[docs] def convert_rotate_shift_to_shift_rotate_order(self, inplane_angle, each_x, each_y):
"""
>>> from spring.segment2d.segment import Segment
>>> a, b, c = Segment().convert_shift_rotate_to_rotate_shift_order(50,22,2)
>>> Segment().convert_rotate_shift_to_shift_rotate_order(a, b, c)
(49.99999868188684, 22.0, 2.0)
"""
psi, tx, ty = self.convert_shift_rotate_to_rotate_shift_order(-inplane_angle, each_x, each_y)
return (-psi) % 360, tx, ty
[docs] def convert_shift_rotate_to_rotate_shift_order(self, inplane_angle, each_x, each_y):
"""
>>> from spring.segment2d.segment import Segment
>>> Segment().convert_shift_rotate_to_rotate_shift_order(50,22,2)
(49.99999868188683, 15.673417091369629, -15.567402839660645)
"""
t_shift = Transform({'type':'SPIDER', 'psi':0.0, 'tx':each_x, 'ty':each_y, 'scale':1.0})
t_rot = Transform({'type':'SPIDER', 'psi':inplane_angle, 'tx':0.0, 'ty':0.0, 'scale':1.0})
composite_transform = t_rot * t_shift
t = composite_transform.get_params('spider')
return t['psi'], t['tx'], t['ty']
[docs] def shift_and_rotate_image(self, image, inplane_angle, each_x, each_y):
angle, tx, ty = self.convert_shift_rotate_to_rotate_shift_order(inplane_angle, each_x, each_y)
trans_rot_image = image.rot_scale_trans2D_background(angle, tx, ty, 1.0)
return trans_rot_image
#===========================================================================
# def shift_and_rotate_image_gridding(self, image, inplane_angle, each_x, each_y):
# trans_rot_image, kb = prepi(image)
# angle, tx, ty = self.convert_shift_rotate_to_rotate_shift_order(inplane_angle, each_x, each_y)
# trans_rot_image = trans_rot_image.rot_scale_conv_new_background(angle, tx, ty, kb, 1.0)
#
# return trans_rot_image
#===========================================================================
[docs] def straighten_segment(self, helix_img, img_size, cols, rows, inplane_angles, mode='straighten'):
if mode == 'bend':
pd_img_size = 2*img_size
else:
pd_img_size = int(1.2*img_size)
helix_pd = Util.pad(helix_img, pd_img_size, pd_img_size, 1, 0, 0, 0, '0')
bent_helix = model_blank(img_size, img_size)
for each_index, each_row in enumerate(rows):
each_col = cols[each_index]
each_angle = inplane_angles[each_index]
if mode == 'bend':
helix_wi = Util.window(helix_pd, img_size, img_size, 1, 0, int(each_row - img_size / 2.0), 0)
helix_rot = rot_shift2D(helix_wi, each_angle, int(each_col - img_size / 2.0), 0)
elif mode == 'straighten':
helix_rot = helix_pd.rot_scale_trans2D_background(-each_angle, int(each_col - img_size / 2.0),
int(each_row - img_size / 2.0), 1.0)
#===============================================================
# helix_rot = self.shift_and_rotate_image(helix_pd, -each_angle, -(each_col - img_size / 2.0),
# -(each_row - img_size / 2.0))
#===============================================================
helix_rot = Util.window(helix_rot, img_size, img_size, 1, 0, 0, 0)
bent_row = helix_rot.get_row(int(img_size / 2.0))
bent_helix.set_row(bent_row, each_index)
return bent_helix
[docs] def compute_rotated_helix_path_on_segment_locally(self, current_coordinates, cut_coordinates, current_inplane_angle,
large_segsizepix, filtered_x, filtered_y):
"""
>>> from spring.segment2d.segment import Segment
>>> Segment().compute_rotated_helix_path_on_segment_locally((17.0, 87.0), (17, 87), 180.0, 10.0, np.arange(10, 20), np.arange(80, 90))
(array([12., 11., 10., 9., 8., 7., 6., 5., 4., 3.]), array([12., 11., 10., 9., 8., 7., 6., 5., 4., 3.]))
"""
center_x, center_y = current_coordinates
cut_center_x, cut_center_y = cut_coordinates
local_x = filtered_x - cut_center_x + large_segsizepix / 2
local_y = filtered_y - cut_center_y + large_segsizepix / 2
x_offset = center_x - cut_center_x
y_offset = center_y - cut_center_y
rot_local_x, rot_local_y = self.rotate_coordinates_by_angle(local_x, local_y, current_inplane_angle,
large_segsizepix // 2+ x_offset, large_segsizepix // 2+ y_offset)
return rot_local_x, rot_local_y
[docs] def set_polyfit_variable(self, large_segsizepix, current_inplane_angle, second_order_fit):
first_order_from_ip_angle = np.tan(np.deg2rad(-current_inplane_angle - 90))
zero_order = large_segsizepix // 2 - (first_order_from_ip_angle * large_segsizepix // 2)
polyfit = [second_order_fit, first_order_from_ip_angle, zero_order]
return polyfit
[docs] def fit_square_function_path_of_coordinates(self, large_segsizepix, helix_coordinates, current_coordinates,
cut_coordinates, current_inplane_angle, second_order_fit=None):
filtered_x, filtered_y = self.select_coordinates_within_segment_dimensions(helix_coordinates,
current_coordinates, 1.5 * large_segsizepix)
rot_local_x, rot_local_y = self.compute_rotated_helix_path_on_segment_locally(current_coordinates,
cut_coordinates, current_inplane_angle, large_segsizepix, filtered_x, filtered_y)
if second_order_fit is None:
if rot_local_y != []:
polyfit = np.polyfit(rot_local_y, rot_local_x, 2)
else:
polyfit = self.set_polyfit_variable(large_segsizepix, current_inplane_angle, second_order_fit=0)
elif second_order_fit is not None:
corrected_local_x = rot_local_x - (second_order_fit * rot_local_y**2)
if rot_local_y != [] or corrected_local_x != []:
polyfit = np.polyfit(rot_local_y, corrected_local_x, 1)
polyfit = np.insert(polyfit, 0, second_order_fit)
else:
polyfit = self.set_polyfit_variable(large_segsizepix, current_inplane_angle, second_order_fit)
return polyfit, filtered_x, filtered_y
[docs] def unbend_segment_using_coordinates(self, segment, large_segsizepix, helix_coordinates, current_coordinates,
cut_coordinates, current_inplane_angle, second_order_fit=None):
polyfit, filtered_x, filtered_y = self.fit_square_function_path_of_coordinates(large_segsizepix,
helix_coordinates, current_coordinates, cut_coordinates, current_inplane_angle, second_order_fit)
cols, rows, inplane_angles = self.compute_bending_path_row_wise(large_segsizepix, polyfit,
current_inplane_angle)
straightened_segment = self.straighten_segment(segment, large_segsizepix, cols, rows, inplane_angles)
central_inplane_angle = inplane_angles[int(len(inplane_angles) / 2.0)]
return straightened_segment, polyfit, central_inplane_angle
