def test_write_star_file_len_4(self): output_file = 'test_output.star' output_array = np.array([(0, '0', 0.0, np.nan, 0.0), (1, '1', 1.0, np.nan, 1.0), (2, '2', 2.0, np.nan, 2.0), (3, '3', 3.0, np.nan, 3.0), (0, '0', 0.0, 0.0, 0.0), (1, '1', 1.0, 1.0, 1.0), (2, '2', 2.0, 2.0, 2.0), (3, '3', 3.0, 3.0, 3.0), (0, '0', 0.0, 0, np.nan), (1, '1', 1.0, 1, np.nan), (2, '2', 2.0, 2, np.nan), (3, '3', 3.0, 3, np.nan)], dtype=[('int', '<i8'), ('str', '|S100'), ('float', '<f8'), ('_rlnAnglePsiPrior', '<f8'), ('_rlnAngleTiltPrior', '<f8')]) header_string = '\ndata_\n\nloop_\nint #1\nstr #2\nfloat #3\n_rlnAnglePsiPrior #4\n_rlnAngleTiltPrior #5\n' return_string, return_array = write_star.write_star_file( output_array=output_array, header_string=header_string, output_file=output_file) assert (len(return_array) == 4) os.remove(output_file)
def main(file_name, tolerance_psi, tolerance_theta, tolerance_filament, window_size, plot=False, typ='sphire', params=None, index=None, output_dir=None): """Start calculation""" if output_dir is None: output_dir = '.' # Check angle range, import arrays if typ == 'sphire': angle_max = 360 angle_min = 0 original_stack = read_sphire.get_sphire_stack(file_name) parameter = read_sphire.get_sphire_file('params', params) indices = read_sphire.get_sphire_file('index', index) substack = create_substack(original_stack, indices) array = combine_arrays(substack, parameter) id_name = 'filament' particle_name = 'data_n' micrograph_name = 'ptcl_source_image' angle_name = ['psi', 'theta'] angle_name_new = ['psi_prior', 'theta_prior'] output_names = list(parameter.dtype.names) elif typ == 'relion': angle_max = 180 angle_min = -180 array, header, path = read_star.import_star_file(file_name) id_name = '_rlnHelicalTubeID' micrograph_name = '_rlnMicrographName' particle_name = '_rlnImageName' angle_name = ['_rlnAnglePsi', '_rlnAngleTilt'] angle_name_new = ['_rlnAnglePsiPrior', '_rlnAngleTiltPrior'] output_names = list(array.dtype.names) + angle_name_new else: print('Unreachable code!') return 'Unreachable code!' # Add the particle number to the array array_order = 'particle_order' order_numbers = np.arange(len(array)) array = add_column(array, order_numbers, array_order) # Sort the array array = np.sort(array, order=[micrograph_name, id_name, particle_name]) for angle in angle_name: # Add rotated data column to the array data_rotated_name = 'data_rotated_{0}'.format(angle) array = add_column(array, array[angle], data_rotated_name) print(np.shape(array)) filament_array = calculations.get_filaments(array, id_name) array_modified = None if plot: do_plot = True else: do_plot = False for angle, angle_new, tolerance in zip(angle_name, angle_name_new, [tolerance_psi, tolerance_theta]): print(angle) data_rotated_name = 'data_rotated_{0}'.format(angle) for idx, filament in enumerate(filament_array): if do_plot: if idx % 10000 == 0: plot = True else: plot = False if idx % 10000 == 0: print('{0}%'.format(idx * 100 / len(filament_array)), len(filament_array)) plot = True else: plot = False if plot: calculations.plot_polar('raw_data', filament[data_rotated_name], 0, angle_max, 0, output=output_dir) filament[ data_rotated_name] = calculations.subtract_and_adjust_angles( filament[data_rotated_name], 0, 180, -180) filament[ data_rotated_name], rotate_angle = calculations.rotate_angles( filament[data_rotated_name], plot, output=output_dir) is_outlier, filament[ data_rotated_name], rotate_angle, inside_tol_idx, outside_tol_idx = calculations.get_filament_outliers( data_rotated=filament[data_rotated_name], rotate_angle=rotate_angle, tolerance=tolerance, tolerance_filament=tolerance_filament, plot=plot, output=output_dir) mean_array = calculations.calculate_mean_prior( input_array=filament[data_rotated_name], window_size=window_size, inside_tolerance_idx=inside_tol_idx, outside_tolerance_idx=outside_tol_idx, plot=plot, output=output_dir) if plot: calculations.plot_polar('mean_array', mean_array, rotate_angle, 180, -180, output=output_dir) mean_array = calculations.subtract_and_adjust_angles( mean_array, -rotate_angle, angle_max, angle_min) if plot: calculations.plot_polar('mean_array', mean_array, 0, angle_max, angle_min, output=output_dir) filament = add_column(filament, mean_array, angle_new) if idx == 0: new_array = filament else: new_array = np.append(new_array, filament) array = add_column(array, new_array[angle_new], angle_new) if typ == 'sphire': array = swap_columns(array, angle, angle_new) else: pass # Sort the array back to the original order array = np.sort(array, order=array_order) if not os.path.exists(output_dir): os.mkdir(output_dir) if typ == 'relion': header_string = write_star.create_header_string(output_names) write_star.write_star_file(array[output_names], header_string, '{0}/TEST.star'.format(output_dir)) elif typ == 'sphire': write_sphire.write_params_file( array, output_names, '{0}/{1}.txt'.format(output_dir, params)) else: pass
def main(file_name, tolerance_psi, tolerance_theta, tolerance_filament, window_size, plot=False, typ='sphire', params=None, index=None, output_dir=None): """Start calculation""" if output_dir is None: output_dir = '.' # Check angle range, import arrays if typ == 'sphire': angle_max = 360 angle_min = 0 original_stack = read_sphire.get_sphire_stack(file_name) parameter = read_sphire.get_sphire_file('params', params) indices = read_sphire.get_sphire_file('index', index) substack = create_substack(original_stack, indices) array = combine_arrays(substack, parameter) id_name = 'filament' particle_name = 'data_n' micrograph_name = 'ptcl_source_image' angle_name = ['theta', 'psi'] angle_name_new = [['theta_prior', 'theta_std'], ['psi_prior', 'psi_std']] output_names = list(parameter.dtype.names) for entry in angle_name_new: output_names += entry elif typ == 'relion': angle_max = 180 angle_min = -180 array, header, path = read_star.import_star_file(file_name) id_name = '_rlnHelicalTubeID' micrograph_name = '_rlnMicrographName' particle_name = '_rlnImageName' angle_name = ['_rlnAngleTilt', '_rlnAnglePsi'] angle_name_new = [['_rlnAngleTiltPrior', '_rlnAngleTiltStd'], ['_rlnAnglePsiPrior', '_rlnAnglePsiStd']] output_names = list(array.dtype.names) for entry in angle_name_new: output_names.append(entry[0]) else: print('Unreachable code!') return 'Unreachable code!' # Add the particle number to the array and sort it by filament and particle order_idx = 'particle_order' ordered_numbers = np.arange(len(array)) array = add_column(array, ordered_numbers, order_idx) array = np.sort(array, order=[micrograph_name, id_name, particle_name]) filament_array = calculations.get_filaments(array, id_name) if plot: do_plot = True else: do_plot = False for angle, angle_new, tolerance in zip(angle_name, angle_name_new, [tolerance_psi, tolerance_theta]): for idx, filament in enumerate(filament_array): if do_plot: if idx % 1000 == 0: plot = True else: plot = False mean_list, std_list = wrapped_distribution(filament[angle]) for name, entry in zip(angle_new, [mean_list, std_list]): filament = add_column(filament, entry, name) if idx == 0: new_array = filament else: new_array = np.append(new_array, filament) if plot: calculations.plot_polar('raw_data', filament[angle], 0, angle_max, 0, output=output_dir) calculations.plot_polar('mean', filament[angle], 0, angle_max, 0, mean=np.degrees(mean_list[0]), output=output_dir) calculations.plot_polar('sigma1', filament[angle], 0, angle_max, 0, mean=np.degrees(mean_list[0]), output=output_dir, tol=np.degrees(std_list[0])) calculations.plot_polar('sigma2', filament[angle], 0, angle_max, 0, mean=np.degrees(mean_list[0]), output=output_dir, tol=2*np.degrees(std_list[0])) for idx, name in enumerate(angle_new): array = add_column(array, new_array[name], name) # Sort the array back to the original order array = np.sort(array, order=order_idx) if not os.path.exists(output_dir): os.mkdir(output_dir) if typ == 'relion': header_string = write_star.create_header_string(output_names) write_star.write_star_file(array[output_names], header_string, '{0}/TEST.star'.format(output_dir)) elif typ == 'sphire': write_sphire.write_params_file(array, output_names, '{0}/TEST.txt'.format(output_dir)) else: pass
def main(file_name, data, tolerance_psi, tolerance_theta, tolerance_filament, window_size, plot=False, typ='sphire', output_dir=None, params=None): """Start calculation""" # Extract data array, stack_id, mic_name, filament_name, particle_id, angle_name, output_names, angle_max, angle_min = data if output_dir is None: output_dir = '.' if plot: do_plot = True else: do_plot = False # Add new angle names to the output dtype IDX_OLD = 0 IDX_NEW = 1 IDX_ROT = 2 dtype_temp = array.dtype.descr for angle in angle_name: dtype_temp.append((angle[IDX_NEW], '<f8')) dtype_temp.append((angle[IDX_ROT], '<f8')) output_names.append(angle[IDX_NEW]) # Create a new combined array array_temp = np.empty(len(array), dtype=dtype_temp) for name in array.dtype.names: array_temp[name] = array[name] for angle in angle_name: array_temp[angle[IDX_ROT]] = np.copy(array[angle[IDX_OLD]]) # Sort the array and remove array_temp from scope array = np.sort(array_temp, order=[mic_name, filament_name, particle_id]) del array_temp # Split the array into filaments filament_array = calculations.get_filaments(array=array, filament_name=filament_name) # Loop over both angles and their tolerance for angle, tolerance in zip(angle_name, [tolerance_psi, tolerance_theta]): # Array names angle_old = angle[IDX_OLD] angle_prior = angle[IDX_NEW] angle_rot = angle[IDX_ROT] print(angle_old) # Loop over all filaments for idx, filament in enumerate(filament_array): # Do plot if necessary if do_plot: if idx % 1000 == 0: plot = True else: plot = False if plot: calculations.plot_polar('raw_data', filament[angle_rot], 0, angle_max, 0, output=output_dir) # Shift the angle range from 180 to -180 calculations.subtract_and_adjust_angles(filament[angle_rot], 0, 180, -180) # Rotate the angle range, so that the median is the new center rotate_angle = calculations.rotate_angles(filament[angle_rot], plot, output=output_dir) # Calculate the indices of outliers is_outlier, rotate_angle, inside_tol_idx, outside_tol_idx = calculations.get_filament_outliers( data_rotated=filament[angle_rot], rotate_angle=rotate_angle, tolerance=tolerance, tolerance_filament=tolerance_filament, plot=plot, output=output_dir) # Calculate the new prior values calculations.calculate_mean_prior( input_array=filament[angle_rot], mean_array=filament[angle_prior], window_size=window_size, inside_tolerance_idx=inside_tol_idx, outside_tolerance_idx=outside_tol_idx, plot=plot, output=output_dir) if plot: calculations.plot_polar('mean_array', filament[angle_prior], rotate_angle, 180, -180, output=output_dir) # Adjust the angles to match the aimed angle range calculations.subtract_and_adjust_angles(filament[angle_prior], -rotate_angle, angle_max, angle_min) if plot: calculations.plot_polar('mean_array', filament[angle_prior], 0, angle_max, angle_min, output=output_dir) # Combine the filaments to one array array_out = np.empty(len(array), dtype=array.dtype.descr) index = 0 for entry in filament_array: for row in entry: array_out[index] = row index += 1 # Sort the array back to the original order array_out = np.sort(array_out, order=stack_id) # Write output if not os.path.exists(output_dir): os.mkdir(output_dir) if typ == 'relion': header_string = write_star.create_header_string(output_names) write_star.write_star_file( array_out[output_names], header_string, '{0}_prior.star'.format(file_name.split('.star')[0])) elif typ == 'sphire': write_sphire.write_params_file( array_out, output_names, '{0}_prior.txt'.format(params.split('.txt')[0])) else: assert (False)