def normalize_background(cls, stack, radius=None):
"""
Normalize background to mean 0 and std 1.
Estimate the mean and std of each image in the stack using pixels
outside radius r (pixels), and normalize the image such that the
background has mean 0 and std 1. Each image in the stack is corrected
separately.
:param radius: radius for normalization (default is half the side of image)
Example:
normalized_stack = cryo_normalize_background(stack,55);
"""
validate_square_projections(stack)
num_images = stack.shape[0] # assuming C-contiguous array
side = stack.shape[1]
if radius is None:
radius = np.floor(side / 2)
# find indices of backgruond pixels in the images
ctr = (side + 1) / 2
x_axis, y_axis = np.meshgrid(range(1, side + 1), range(1, side + 1))
radiisq = (x_axis.flatten() - ctr)**2 + (y_axis.flatten() - ctr)**2
background_pixels_idx = radiisq > radius * radius
if AspireConfig.verbosity == 1:
pb = ProgressBar(total=100,
prefix='normalizing background',
suffix='completed',
decimals=0,
length=100,
fill='%')
else:
pb = None
normalized_stack = np.ones(stack.shape)
for i in range(num_images):
if pb:
pb.print_progress_bar((i + 1) / num_images * 100)
proj = stack[i, :, :]
background_pixels = proj.flatten() * background_pixels_idx
background_pixels = background_pixels[background_pixels != 0]
# compute mean and standard deviation of background pixels
proj_mean = np.mean(background_pixels)
std = np.std(background_pixels, ddof=1)
# normalize the projections
if std < 1.0e-5:
logger.warning(
f'Variance of background of image {i} is too small (std={std}). '
'Cannot normalize!')
normalized_stack[i, :, :] = (proj - proj_mean) / std
return normalized_stack
def cryo_epsds(imstack, samples_idx, max_d, verbose=None):
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, verbose)
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))
if idx is None:
logger.warning('something went wrong in bsearch')
r2[i + p - 1, j + p - 1] = r[idx - 1]
w = gwindow(p, max_d)
p2 = cfft2(r2 * w)
err = np.linalg.norm(p2.imag) / np.linalg.norm(p2)
if err > 1e-12:
logger.warning('Large imaginary components in P2 = {}'.format(err))
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)
if len(neg_idx[0]) != 0:
max_neg_err = np.max(np.abs(p2[neg_idx]))
if max_neg_err > 1e-2:
neg_norm = np.linalg.norm(p2[neg_idx])
logger.warning(
'Power specrtum P2 has negative values with energy {}'.format(
neg_norm))
p2[neg_idx] = 0
return p2, r, r2, x
def phaseflip_star_file(cls, star_file, pixel_size=None):
"""
todo add verbosity
"""
# star is a list of star lines describing projections
star_records = read_star(star_file)['__root__']
num_projections = len(star_records)
projs_init = False # has the stack been initialized already
last_processed_stack = None
for idx in range(num_projections):
# Get the identification string of the next image to process.
# This is composed from the index of the image within an image stack,
# followed by '@' and followed by the filename of the MRC stack.
image_id = star_records[idx].rlnImageName
image_parts = image_id.split('@')
image_idx = int(image_parts[0]) - 1
stack_name = image_parts[1]
# Read the image stack from the disk, if different from the current one.
# TODO can we revert this condition to positive? what they're equal?
if stack_name != last_processed_stack:
mrc_path = os.path.join(os.path.dirname(star_file), stack_name)
stack = load_stack_from_file(mrc_path)
logger.info(
f"flipping stack in {yellow(os.path.basename(mrc_path))}"
f" - {stack.shape}")
last_processed_stack = stack_name
if image_idx > stack.shape[2]:
raise DimensionsIncompatible(
f'projection {image_idx} in '
f'stack {stack_name} does not exist')
proj = stack[image_idx]
validate_square_projections(proj)
side = proj.shape[1]
if not projs_init: # TODO why not initialize before loop (maybe b/c of huge stacks?)
# projections was "PFprojs" originally
projections = np.zeros((num_projections, side, side),
dtype='float32')
projs_init = True
star_record_data = Box(
cryo_parse_Relion_CTF_struct(star_records[idx]))
if pixel_size is None:
if star_record_data.tmppixA != -1:
pixel_size = star_record_data.tmppixA
else:
raise WrongInput(
"Pixel size not provided and does not appear in STAR file"
)
h = cryo_CTF_Relion(side, star_record_data)
imhat = fftshift(fft2(proj))
pfim = ifft2(ifftshift(imhat * np.sign(h)))
if side % 2 == 1:
# This test is only vali for odd n
# images are single precision
imaginery_comp = np.norm(np.imag(pfim[:])) / np.norm(pfim[:])
if imaginery_comp > 5.0e-7:
logger.warning(
f"Large imaginary components in image {image_idx}"
f" in stack {stack_name} = {imaginery_comp}")
pfim = np.real(pfim)
projections[idx, :, :] = pfim.astype('float32')
return projections
def compare_stacks(stack1, stack2, verbose=None, max_error=None):
""" Calculate the difference between two projection-stacks.
Return the relative error between them.
:param stack1: first stack to compare
:param stack2: second stack to compare
:param verbose: level of verbosity
verbose=0 silent
verbose=1 show progress bar
verbose=2 print progress every 1000 images
verbose=3 print message for each processed image
:param max_error: when given, raise an exception if difference between stacks is too big
:return: returns the accumulative error between the two stacks
"""
if max_error is not None:
try:
max_error = np.longdouble(max_error)
except (TypeError, ValueError):
raise WrongInput("max_error must be either a float or an integer!")
if verbose is None:
verbose = AspireConfig.verbosity
# check the dimensions of the stack are compatible
if stack1.shape != stack2.shape:
raise DimensionsIncompatible("Can't compare stacks of different sizes!"
f" {stack1.shape} != {stack2.shape}")
num_of_images = stack1.shape[0]
if num_of_images == 0:
logger.warning('stacks are empty!')
if verbose == 1:
pb = ProgressBar(total=100,
prefix='comparing:',
suffix='completed',
decimals=0,
length=100,
fill='%')
relative_err = 0
accumulated_err = 0
for i in range(num_of_images):
err = np.linalg.norm(stack1[i] - stack2[i]) / np.linalg.norm(stack1[i])
accumulated_err += err
relative_err = accumulated_err / (i + 1)
# if we already reached a relatively big error, we can stop here
# we can't ask "if max_error" as max_error is so small and treated as 0 (False)
if max_error is not None and relative_err > max_error:
raise ErrorTooBig(
'Stacks comparison failed! error is too big: {}'.format(
relative_err))
if verbose == 0:
continue
elif verbose == 1:
pb.print_progress_bar((i + 1) / num_of_images * 100)
elif verbose == 2 and (i + 1) % 100 == 0:
logger.info(
f'Finished comparing {i+1}/{num_of_images} projections. '
f'Relative error so far: {relative_err}')
elif verbose == 3:
logger.info(
f'Difference between projections ({i+1}) <> ({i+1}): {err}')
if verbose == 2:
logger.info(
f'Finished comparing {num_of_images}/{num_of_images} projections. '
f'Relative error: {relative_err}')
return relative_err
def cryo_epsdr(vol, samples_idx, max_d, verbose):
p = vol.shape[0]
k = vol.shape[2]
i, j = np.meshgrid(np.arange(max_d + 1), np.arange(max_d + 1))
dists = np.square(i) + np.square(j)
dsquare = np.sort(np.unique(dists[np.where(dists <= max_d**2)]))
corrs = np.zeros(len(dsquare))
corr_count = np.zeros(len(dsquare))
x = np.sqrt(dsquare)
dist_map = np.zeros(dists.shape)
for i in range(max_d + 1):
for j in range(max_d + 1):
d = i**2 + j**2
if d <= max_d**2:
idx, _ = bsearch(dsquare, d - 1e-13, d + 1e-13)
if idx is None:
logger.warning('something went wrong in bsearch')
dist_map[i, j] = idx
dist_map = dist_map.astype('int') - 1
valid_dists = np.where(dist_map != -1)
mask = np.zeros((p, p))
mask[samples_idx] = 1
tmp = np.zeros((2 * p + 1, 2 * p + 1))
tmp[:p, :p] = mask
ftmp = np.fft.fft2(tmp)
c = np.fft.ifft2(ftmp * np.conj(ftmp))
c = c[:max_d + 1, :max_d + 1]
c = np.round(c.real).astype('int')
r = np.zeros(len(corrs))
# optimized version
vol = vol.transpose((2, 0, 1)).copy()
input_fft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
output_fft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
input_ifft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
output_ifft2 = np.zeros((2 * p + 1, 2 * p + 1), dtype='complex128')
flags = ('FFTW_MEASURE', 'FFTW_UNALIGNED')
fft2 = pyfftw.FFTW(input_fft2,
output_fft2,
axes=(0, 1),
direction='FFTW_FORWARD',
flags=flags)
ifft2 = pyfftw.FFTW(input_ifft2,
output_ifft2,
axes=(0, 1),
direction='FFTW_BACKWARD',
flags=flags)
sum_s = np.zeros(output_ifft2.shape, output_ifft2.dtype)
sum_c = c * vol.shape[0]
for i in range(k):
proj = vol[i]
input_fft2[samples_idx] = proj[samples_idx]
fft2()
np.multiply(output_fft2, np.conj(output_fft2), out=input_ifft2)
ifft2()
sum_s += output_ifft2
for curr_dist in zip(valid_dists[0], valid_dists[1]):
dmidx = dist_map[curr_dist]
corrs[dmidx] += sum_s[curr_dist].real
corr_count[dmidx] += sum_c[curr_dist]
idx = np.where(corr_count != 0)[0]
r[idx] += corrs[idx] / corr_count[idx]
cnt = corr_count[idx]
idx = np.where(corr_count == 0)[0]
r[idx] = 0
x[idx] = 0
return r, x, cnt