def bessel_ns_radial(bandlimit, support_size, x): bessel = get_bessel() bessel = bessel[bessel[:, 3] <= 2 * np.pi * bandlimit * support_size, :] angular_freqs = bessel[:, 0] max_ang_freq = int(np.max(angular_freqs)) n_theta = int(np.ceil(16 * bandlimit * support_size)) if n_theta % 2 == 1: n_theta += 1 radian_freqs = bessel[:, 1] r_ns = bessel[:, 2] phi_ns = np.zeros((len(x), len(angular_freqs))) phi = {} angular_freqs_length = len(angular_freqs) for i in range(angular_freqs_length): r0 = x * r_ns[i] / bandlimit f = sp.jv(angular_freqs[i], r0) # probably the square and the sqrt not needed tmp = np.pi * np.square(sp.jv(angular_freqs[i] + 1, r_ns[i])) phi_ns[:, i] = f / (bandlimit * np.sqrt(tmp)) for i in range(max_ang_freq + 1): phi[i] = phi_ns[:, angular_freqs == i] struct_dict = {'phi_ns': phi, 'angular_freqs': angular_freqs, 'radian_freqs': radian_freqs, 'n_theta': n_theta} return common.create_struct(struct_dict)
def fast_rotate_precomp(szx, szy, phi): phi, mult90 = adjust_rotate(phi) phi = np.pi * phi / 180 phi = -phi if szy % 2: cy = (szy + 1) // 2 sy = 0 else: cy = szy // 2 + 1 sy = 0.5 if szx % 2: cx = (szx + 1) // 2 sx = 0 else: cx = szx // 2 + 1 sx = 0.5 my = np.zeros((szy, szx), dtype='complex128') r = np.arange(cy) r_t = np.arange(szy, cy, -1) - 1 u = (1 - np.cos(phi)) / np.sin(phi + np.finfo(float).eps) alpha1 = 2 * np.pi * 1j * r / szy for x in range(szx): ux = u * (x + 1 - cx + sx) my[r, x] = np.exp(alpha1 * ux) my[r_t, x] = np.conj(my[1:cy - 2 * sy, x]) my = my.T mx = np.zeros((szx, szy), dtype='complex128') r = np.arange(cx) r_t = np.arange(szx, cx, -1) - 1 u = -np.sin(phi) alpha2 = 2 * np.pi * 1j * r / szx for y in range(szy): uy = u * (y + 1 - cy + sy) mx[r, y] = np.exp(alpha2 * uy) mx[r_t, y] = np.conj(mx[1:cx - 2 * sx, y]) # because I am using real fft I take only part of mx and my return common.create_struct({ 'phi': phi, 'mx': mx[:szx // 2 + 1].copy(), 'my': my[:, :szy // 2 + 1].copy(), 'mult90': mult90 })
def fbcoeff_nfft(images, support_size, basis, sample_points, num_threads): split_images = np.array_split(images, num_threads, axis=2) image_size = split_images[0].shape[0] orig = int(np.floor(image_size / 2)) start_pixel = orig - support_size end_pixel = orig + support_size new_image_size = int(2 * support_size) # unpacking input phi_ns = basis.phi_ns angular_freqs = basis.angular_freqs max_angular_freqs = int(np.max(angular_freqs)) n_theta = basis.n_theta x = sample_points.x w = sample_points.w w = w * x # sampling points in the fourier domain freqs = pft_freqs(x, n_theta) precomp = common.create_struct({ 'n_theta': n_theta, 'n_r': len(x), 'resolution': new_image_size, 'freqs': freqs }) scale = 2 * np.pi / n_theta coeff_pos_k = [] pos_k = [] for i in range(num_threads): curr_images = split_images[i] # can work with odd images as well curr_images = curr_images[start_pixel:end_pixel, start_pixel:end_pixel, :] tmp = cryo_pft_nfft(curr_images, precomp) pf_f = scale * np.fft.fft(tmp, axis=1) pos_k.append(pf_f[:, :max_angular_freqs + 1, :]) pos_k = np.concatenate(pos_k, axis=2) for i in range(max_angular_freqs + 1): coeff_pos_k.append( np.einsum('ki, k, kj -> ij', phi_ns[i], w, pos_k[:, i])) return coeff_pos_k
def organize_star_records(star_records): stacks_info = {} for i, rec in enumerate(star_records): pos, path = rec.rlnImageName.split('@') pos = int(pos) - 1 if path in stacks_info.keys(): stack_struct = stacks_info[path] stack_struct.pos_in_stack.append(pos) stack_struct.pos_in_records.append(i) else: stack_struct = create_struct({'pos_in_stack': [pos], 'pos_in_records': [i]}) stacks_info[path] = stack_struct for path in stacks_info: stack_struct = stacks_info[path] stack_struct.pos_in_stack = np.array(stack_struct.pos_in_stack) stack_struct.pos_in_records = np.array(stack_struct.pos_in_records) return stacks_info
def cryo_abinitio_c1_worker(stack, algo, n_theta=360, n_r=0.5, max_shift=0.15, shift_step=1): resolution = stack.shape[1] num_projections = stack.shape[2] # n_r = int(np.ceil(n_r * resolution)) max_shift = int(np.ceil(max_shift * resolution)) mask_radius = resolution * 0.45 # mask_radius is of the form xxx.5 if mask_radius * 2 == int(mask_radius * 2): mask_radius = np.ceil(mask_radius) # mask is not of the form xxx.5 else: mask_radius = int(round(mask_radius)) # mask projections center = (resolution + 1) / 2 m = fuzzy_mask(resolution, mask_radius, origin=(center, center)) masked_projs = stack.copy() masked_projs = masked_projs.transpose((2, 0, 1)) masked_projs *= m masked_projs = masked_projs.transpose((1, 2, 0)).copy() # compute polar fourier transform pf, _ = cryo_pft(masked_projs, n_r, n_theta) # find common lines from projections print( 'Finding common lines between pairs of imagges with maximum shift of {} pixels and shift step of {}' .format(max_shift, shift_step)) tic = time.time() clstack, _, _, _, _ = cryo_clmatrix_cpu(pf, num_projections, max_shift, shift_step) toc = time.time() print('Finished in {} seconds'.format(toc - tic)) if algo == 2: s = cryo_syncmatrix_vote(clstack, n_theta) rotations = cryo_sync_rotations(s) else: raise NotImplementedError('algo currently support only "2"!') est_shifts, _ = cryo_estimate_shifts(pf, rotations, max_shift, shift_step) # reconstruct downsampled volume with no CTF correction n = stack.shape[1] params = create_struct({ 'rot_matrices': rotations, 'ctf': np.ones((n, n)), 'ampl': np.ones(num_projections), 'ctf_idx': np.array([True] * num_projections), 'shifts': est_shifts }) print('Estimating mean') tic = time.time() v1, _ = cryo_estimate_mean(stack, params) toc = time.time() print('Finished in {} seconds'.format(toc - tic)) v1 = v1.real return v1
def compute_spca(images, noise_v_r, adaptive_support=False): num_images = images.shape[2] resolution = images.shape[0] if adaptive_support: raise NotImplementedError('Adaptive support was not implemented yet') # energy_thresh = 0.99 # # # Estimate bandlimit and compact support size # [bandlimit, support_size] = choose_support_v6(common.fast_cfft2(images), energy_thresh) # # Rescale between 0 and 0.5 # bandlimit = bandlimit * 0.5 / np.floor(resolution / 2.0) else: bandlimit = 0.5 support_size = resolution // 2 n_r = int(np.ceil(4 * bandlimit * support_size)) basis, sample_points = precompute_fb(n_r, support_size, bandlimit) _, coeff, mean_coeff, spca_coeff, u, d = jobscript_ffbspca(images, support_size, noise_v_r, basis, sample_points) ang_freqs = [] rad_freqs = [] vec_d = [] for i in range(len(d)): if len(d[i]) != 0: ang_freqs.extend(np.ones(len(d[i]), dtype='int') * i) rad_freqs.extend(np.arange(len(d[i])) + 1) vec_d.extend(d[i]) ang_freqs = np.array(ang_freqs) rad_freqs = np.array(rad_freqs) d = np.array(vec_d) k = min(len(d), 400) # keep the top 400 components sorted_indices = np.argsort(-d) sorted_indices = sorted_indices[:k] d = d[sorted_indices] ang_freqs = ang_freqs[sorted_indices] rad_freqs = rad_freqs[sorted_indices] s_coeff = np.zeros((len(d), num_images), dtype='complex128') for i in range(len(d)): s_coeff[i] = spca_coeff[ang_freqs[i]][rad_freqs[i] - 1] fn = ift_fb(support_size, bandlimit) eig_im = np.zeros((np.square(2 * support_size), len(d)), dtype='complex128') for i in range(len(d)): tmp = fn[ang_freqs[i]] tmp = tmp.reshape((int(np.square(2 * support_size)), tmp.shape[2]), order='F') eig_im[:, i] = np.dot(tmp, u[ang_freqs[i]][:, rad_freqs[i] - 1]) fn0 = fn[0].reshape((int(np.square(2 * support_size)), fn[0].shape[2]), order='F') spca_data_struct = {'eigval': d, 'freqs': ang_freqs, 'radial_freqs': rad_freqs, 'coeff': s_coeff, 'mean': mean_coeff, 'c': bandlimit, 'r': support_size, 'eig_im': eig_im, 'fn0': fn0} spca_data = common.create_struct(spca_data_struct) return spca_data