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
