def cryo_epsds(imstack, samples_idx, max_d):
p = imstack.shape[0]
if max_d >= p:
max_d = p - 1
print('max_d too large. Setting max_d to {}'.format(max_d))
r, x, _ = cryo_epsdr(imstack, samples_idx, max_d)
r2 = np.zeros((2 * p - 1, 2 * p - 1))
dsquare = np.square(x)
for i in range(-max_d, max_d + 1):
for j in range(-max_d, max_d + 1):
d = i**2 + j**2
if d <= max_d**2:
idx, _ = bsearch(dsquare, d * (1 - 1e-13), d * (1 + 1e-13))
r2[i + p - 1, j + p - 1] = r[idx - 1]
w = gwindow(p, max_d)
p2 = common.fast_cfft2(r2 * w)
p2 = p2.real
e = 0
for i in range(imstack.shape[2]):
im = imstack[:, :, i]
e += np.sum(np.square(im[samples_idx] - np.mean(im[samples_idx])))
mean_e = e / (len(samples_idx[0]) * imstack.shape[2])
p2 = (p2 / p2.sum()) * mean_e * p2.size
neg_idx = np.where(p2 < 0)
p2[neg_idx] = 0
return p2, r, r2, x
def im_filter(im, filter_f):
n_im = im.shape[2]
n_filter = filter_f.shape[2]
if n_filter != 1 and n_filter != n_im:
raise ValueError('The number of filters must be 1 or match the number of images')
if im.shape[:2] != filter_f.shape[:2]:
raise ValueError('The size of the images and filters must match')
im_f = fast_cfft2(im, axes=(0, 1))
im_filtered_f = im_f * filter_f
im_filtered = np.real(fast_icfft2(im_filtered_f, axes=(0, 1)))
return im_filtered
def im_backproject(im, rot_matrices, half_pixel=False):
# TODO - test for even resolution
resolution = im.shape[1]
n = im.shape[2]
if im.shape[0] != resolution:
raise ValueError('im must be squared - LxLxN')
if rot_matrices.shape[2] != n:
raise ValueError('The number of rotation matrices must match the number of images')
pts_rot = rotated_grids(resolution, rot_matrices, half_pixel)
pts_rot = pts_rot.reshape((3, -1), order='F')
if resolution % 2 == 0 and half_pixel:
grid = np.arange(-resolution / 2, resolution / 2)
y, x = np.meshgrid(grid, grid)
phase_shift = 2 * np.pi * (x + y) / (2 * resolution)
im = np.einsum('ijk, ij -> ijk', im, np.exp(1j * phase_shift))
im = im.transpose((1, 0, 2))
im_f = fast_cfft2(im, axes=(0, 1)) / resolution ** 2
if resolution % 2 == 0:
if half_pixel:
grid = np.arange(-resolution / 2, resolution / 2)
y, x = np.meshgrid(grid, grid)
phase_shift = 2 * np.pi * (x + y) / (2 * resolution)
phase_shift += - np.reshape(pts_rot.sum(0), (resolution, resolution, n)) / 2
im_f = np.einsum('ijk, ij -> ijk', im_f, np.exp(1j * phase_shift))
else:
im_f[0] = 0
im_f[:, 0] = 0
im_f = im_f.flatten('F')
vol = anufft3(im_f, pts_rot, [resolution] * 3)
vol = vol.real
return vol
def cryo_prewhiten(proj, noise_response, rel_threshold=None):
"""
Pre-whiten a stack of projections using the power spectrum of the noise.
:param proj: stack of images/projections
:param noise_response: 2d image with the power spectrum of the noise. If all
images are to be whitened with respect to the same power spectrum,
this is a single image. If each image is to be whitened with respect
to a different power spectrum, this is a three-dimensional array with
the same number of 2d slices as the stack of images.
:param rel_threshold: The relative threshold used to determine which frequencies
to whiten and which to set to zero. If empty (the default)
all filter values less than 100*eps(class(proj)) are
zeroed out, while otherwise, all filter values less than
threshold times the maximum filter value for each filter
is set to zero.
:return: Pre-whitened stack of images.
"""
delta = np.finfo(proj.dtype).eps
L1, L2, num_images = proj.shape
l1 = L1 // 2
l2 = L2 // 2
K1, K2 = noise_response.shape
k1 = int(np.ceil(K1 / 2))
k2 = int(np.ceil(K2 / 2))
filter_var = np.sqrt(noise_response)
filter_var /= np.linalg.norm(filter_var)
filter_var = (filter_var + np.flipud(filter_var)) / 2
filter_var = (filter_var + np.fliplr(filter_var)) / 2
if rel_threshold is None:
nzidx = np.where(filter_var > 100 * delta)
else:
raise NotImplementedError('not implemented for rel_threshold != None')
fnz = filter_var[nzidx]
one_over_fnz = 1 / fnz
# matrix with 1/fnz in nzidx, 0 elsewhere
one_over_fnz_as_mat = np.zeros(
(noise_response.shape[0], noise_response.shape[1]))
one_over_fnz_as_mat[nzidx] += one_over_fnz
pp = np.zeros((noise_response.shape[0], noise_response.shape[1]))
p2 = np.zeros((num_images, L1, L2), dtype='complex128')
proj = proj.transpose((2, 0, 1)).copy()
row_start_idx = k1 - l1 - 1
row_end_idx = k1 + l1
col_start_idx = k2 - l2 - 1
col_end_idx = k2 + l2
if L1 % 2 == 0 and L2 % 2 == 0:
row_end_idx -= 1
col_end_idx -= 1
for i in range(num_images):
pp[row_start_idx:row_end_idx, col_start_idx:col_end_idx] = proj[i]
fp = common.fast_cfft2(pp)
fp *= one_over_fnz_as_mat
pp2 = common.fast_icfft2(fp)
p2[i] = np.real(pp2[row_start_idx:row_end_idx,
col_start_idx:col_end_idx])
# change back to x,y,z convention
proj = p2.real.transpose((1, 2, 0)).copy()
return proj, filter_var, nzidx