def subtract(p1, submap_ft, refmap_ft, sx, sy, s, a, apix, def1, def2, angast,
phase, kv, ac, cs, az, el, sk, xshift, yshift, coefs_method, r,
nr, pfac):
c = ctf.eval_ctf(s / apix,
a,
def1,
def2,
angast,
phase,
kv,
ac,
cs,
bf=0,
lp=2 * apix)
orient = euler2rot(np.deg2rad(az), np.deg2rad(el), np.deg2rad(sk))
pshift = np.exp(-2 * np.pi * 1j * (-xshift * sx + -yshift * sy))
p2 = vop.interpolate_slice_numba(submap_ft, orient, pfac=pfac)
p2 *= pshift
if coefs_method < 1:
# p1s = p1 - p2 * c
p1s = p2 * c
elif coefs_method == 1:
p3 = vop.interpolate_slice_numba(refmap_ft, orient, pfac=pfac)
p3 *= pshift
frc = np.abs(algo.bincorr_nb(p1, p3 * c, r, nr))
coefs = np.take(frc, r)
# p1s = p1 - p2 * c * coefs
p1s = p2 * c * coefs
return p1s
def project(f3d,
p,
s,
sx,
sy,
a,
apply_ctf=False,
size=None,
flip_phase=False):
orient = util.euler2rot(np.deg2rad(p[star.Relion.ANGLEROT]),
np.deg2rad(p[star.Relion.ANGLETILT]),
np.deg2rad(p[star.Relion.ANGLEPSI]))
pshift = np.exp(
-2 * np.pi * 1j *
(-p[star.Relion.ORIGINX] * sx + -p[star.Relion.ORIGINY] * sy))
f2d = vop.interpolate_slice_numba(f3d, orient, size=size)
f2d *= pshift
if apply_ctf or flip_phase:
apix = star.calculate_apix(p)
c = ctf.eval_ctf(s / apix,
a,
p[star.Relion.DEFOCUSU],
p[star.Relion.DEFOCUSV],
p[star.Relion.DEFOCUSANGLE],
p[star.Relion.PHASESHIFT],
p[star.Relion.VOLTAGE],
p[star.Relion.AC],
p[star.Relion.CS],
bf=0,
lp=2 * apix)
if flip_phase:
c = np.sign(c)
f2d *= c
return f2d
def particle_xcorr(ptcl, refmap_ft):
r = util.euler2rot(*np.deg2rad(ptcl[star.Relion.ANGLES]))
proj = vop.interpolate_slice_numba(refmap_ft, r)
c = ctf.eval_ctf(s / apix,
a,
def1[i],
def2[i],
angast[i],
phase[i],
kv[i],
ac[i],
cs[i],
bf=0,
lp=2 * apix)
pshift = np.exp(-2 * np.pi * 1j * (-xshift[i] * sx + -yshift * sy))
proj_ctf = proj * pshift * c
with mrc.ZSliceReader(ptcl[star.Relion.IMAGE_NAME]) as f:
exp_image_fft = rfft2(fftshift(f.read(i)))
xcor_fft = exp_image_fft * proj_ctf
xcor = fftshift(irfft2(xcor_fft))
return xcor
def process_snr(nn,
mic_file,
metadata,
freqs,
angles,
apix,
cutoff,
hp=.005,
phaseflip=True,
augment=False):
""" Denoise a cryoEM image for SNR calculation
The following steps are performed:
(1) The micrograph is loaded, phaseflipped, and Fourier cropped
(2) A bandpass filter is applied with pass-band from cutoff to 1/200A
(3) The inverse FT is calculated to return to real-space
(4) The micrograph is padded do a dimension divisible by 32
(5) The padded is passed through the CNN to denoise
(6) The Fourier cropped denoised and raw images are returned
"""
# Load the micrograph and phase-flip to correct the CTF
mic = load_mic(mic_file)
mic_ft = rfft2(mic)
if phaseflip:
m = metadata
try:
phase_shift = m[star.relion.PHASESHIFT]
except:
phase_shift = 0.
ctf_img = ctf.eval_ctf(freqs,
angles,
m[star.Relion.DEFOCUSU],
m[star.Relion.DEFOCUSV],
m[star.Relion.DEFOCUSANGLE],
phase_shift,
m[star.Relion.VOLTAGE],
m[star.Relion.AC],
m[star.Relion.CS],
lp=2 * apix)
mic_ft *= np.sign(ctf_img)
# Fourier crop the micrograph and bandpass filter
mic_ft_bin = fourier_crop(mic_ft, freqs, cutoff)
freqs_bin = fourier_crop(freqs, freqs, cutoff)
bp_filter = ((1. - 1. / (1. + (freqs_bin / hp)**(10))) + 1. /
(1. + (freqs_bin / cutoff)**(10))) / 2.
mic_ft_bin *= bp_filter
mic_bin = normalize(irfft2(mic_ft_bin).real)
# Pad the image so the dimension is divisible by 32 (2**5)
n_x, n_y = mic_bin.shape
x_p, y_p = (next32(n_x) - n_x) // 2, (next32(n_y) - n_y) // 2
mic_bin = np.pad(mic_bin, ((x_p, x_p), (y_p, y_p)), mode='mean')
n_x2, n_y2 = mic_bin.shape
# Denoise and unpad the image
# Optionally, generate augmented copies of the data by rotation
# and reflection (reflection not implemented yet). The idea here is
# to average over the rotation/reflection-covariant kernels and maybe
# further reduce the noise/bias? We will test this directly and add
# to denoise.py if the results are favorable...
if augment:
m = mic_bin.reshape((1, n_x2, n_y2, 1))
denoised = [np.rot90(m, i + 1, axes=(1, 2)) for i in range(4)]
denoised = [nn.predict(denoised[i]) for i in range(4)]
denoised = [
np.rot90(denoised[i], -(i + 1), axes=(1, 2)) for i in range(4)
]
denoised = np.mean(np.array(denoised), axis=0)
denoised = denoised.reshape((n_x2, n_y2))
else:
denoised = nn.predict(mic_bin.reshape((1, n_x2, n_y2, 1)))
denoised = denoised.reshape((n_x2, n_y2))
# We don't normalize the outputs for calculate covariances
denoised = denoised[x_p:n_x + x_p, y_p:n_y + y_p]
raw = mic_bin[x_p:n_x + x_p, y_p:n_y + y_p]
return raw, denoised
def process(metadata,
cutoff,
window,
mic_file,
freqs,
angles,
bandpass=True,
hp=.005,
phaseflip=True):
""" Process a training micrograph.
The following steps are performed:
(1) The micrograph is loaded, Fourier transformed and Fourier cropped
(2) A bandpass filter is applied with pass-band from cutoff to 1/200A
(3) The FT is multiplied by the sign of the CTF (phase-flipping)
This is a crude form of CTF correction.
(4) The inverse FT is calculated to return to real-space
(5) The binned, filtered image is divided into patches, which are
normalized (Z-score normalization) and returned
"""
mic = load_mic(mic_file)
mic_ft = rfft2(mic)
mic_ft_bin = fourier_crop(mic_ft, freqs, cutoff)
freqs_bin = fourier_crop(freqs, freqs, cutoff)
angs_bin = fourier_crop(angles, freqs, cutoff)
apix_bin = 0.5 / freqs_bin[0, -1]
if bandpass:
bp_filt = ((1. - 1. / (1. + (freqs_bin / hp)**10)) + 1. /
(1. + (freqs_bin / cutoff)**10) / 2.)
mic_ft_bin *= bp_filt
if phaseflip:
m = metadata
try:
phase_shift = m[star.relion.PHASESHIFT]
except:
phase_shift = 0.
ctf_img = ctf.eval_ctf(freqs_bin,
angs_bin,
m[star.Relion.DEFOCUSU],
m[star.Relion.DEFOCUSV],
m[star.Relion.DEFOCUSANGLE],
phase_shift,
m[star.Relion.VOLTAGE],
m[star.Relion.AC],
m[star.Relion.CS],
bf=0,
lp=2 * apix_bin)
mic_ft_bin *= np.sign(ctf_img)
mic = irfft2(mic_ft_bin).real.astype('float32')
patches = [normalize(p) for p in get_patches(mic, window)]
n_patches = len(patches)
return np.array(patches).reshape((n_patches, window, window, 1)), apix_bin
def process(nn,
mic_file,
metadata,
freqs,
angles,
apix,
cutoff,
softmask,
merge_band,
hp=.005,
phaseflip=True,
flipback=False,
merge_noisy=False):
""" Denoise a cryoEM image
The following steps are performed:
(1) The micrograph is loaded, phaseflipped, and Fourier cropped
(2) A bandpass filter is applied with pass-band from cutoff to 1/200A
(3) The inverse FT is calculated to return to real-space
(4) The micrograph is padded do a dimension divisible by 32
(5) The padded is passed through the CNN to denoise then unpadded
(6) The micrograph is upsampled by padding the Fourier transform
with zeros. This procedure creates a sharp edge between components
with finite amplitudes and zero amplitude, which manifests as
high-frequency 'ringing' artefacts in real space. To reduce these,
we apply a soft mask to the denoised FT.
(7) Optional: the low-pass filtered denoised image is combined with
the complementary high-resolution noisy image.
"""
# Load the micrograph and phase-flip to correct the CTF
mic = normalize(load_mic(mic_file))
mic_ft = rfft2(mic)
if phaseflip:
m = metadata
try:
phase_shift = m[star.relion.PHASESHIFT]
except:
phase_shift = 0.
ctf_img = ctf.eval_ctf(freqs,
angles,
m[star.Relion.DEFOCUSU],
m[star.Relion.DEFOCUSV],
m[star.Relion.DEFOCUSANGLE],
phase_shift,
m[star.Relion.VOLTAGE],
m[star.Relion.AC],
m[star.Relion.CS],
lp=2 * apix)
mic_ft *= np.sign(ctf_img)
# Fourier crop the micrograph and bandpass filter
mic_ft_bin = fourier_crop(mic_ft, freqs, cutoff)
freqs_bin = fourier_crop(freqs, freqs, cutoff)
bp_filter = ((1. - 1. / (1. + (freqs_bin / hp)**(10))) + 1. /
(1. + (freqs_bin / cutoff)**(10))) / 2.
mic_ft_bin *= bp_filter
mic_bin = normalize(irfft2(mic_ft_bin).real)
# Pad the image so the dimension is divisible by 32 (2**5)
n_x, n_y = mic_bin.shape
x_p, y_p = (next32(n_x) - n_x) // 2, (next32(n_y) - n_y) // 2
mic_bin = np.pad(mic_bin, ((x_p, x_p), (y_p, y_p)), mode='mean')
n_x2, n_y2 = mic_bin.shape
# Denoise and unpad the image
denoised = nn.predict(mic_bin.reshape((1, n_x2, n_y2, 1)))
denoised = denoised.reshape((n_x2, n_y2))
denoised = denoised[x_p:n_x + x_p, y_p:n_y + y_p]
# Upsample by Fourier padding
denoised_ft = rfft2(normalize(denoised))
denoised_ft_full = fourier_pad_to_shape(denoised_ft, mic_ft.shape)
# Merge images or apply final low-pass filter
if merge_noisy:
# Sample every nth Fourier amplitude for normalization
n = 50
denoised_amp_band = np.abs(denoised_ft_full[merge_band][::n]).ravel()
noisy_amp_band = np.abs(mic_ft[merge_band][::n]).ravel()
merge_factor = (np.std(denoised_amp_band) / np.std(noisy_amp_band))
# Separate images into amplitude and phase
denoised_amp = np.abs(denoised_ft_full)
denoised_phase = np.angle(denoised_ft_full)
noisy_amp = np.abs(mic_ft)
noisy_phase = np.angle(mic_ft)
# Weighted amplitude average, unweighted phase average, reconstruct FT
merge_amp = (denoised_amp *
softmask) + merge_factor * noisy_amp * (1. - softmask)
merge_phase = (denoised_phase *
softmask) + noisy_phase * (1. - softmask)
merge_ft = merge_amp * (np.cos(merge_phase) + 1j * np.sin(merge_phase))
denoised_ft_full = merge_ft
else:
denoised_ft_full = denoised_ft_full * softmask
# Flip phases back (multiply FT again by sign of the CTF) if requested
if phaseflip and flipback:
denoised_ft_full *= np.sign(ctf_img)
denoised_full = irfft2(denoised_ft_full).real.astype(np.float32)
new_mic = denoised_full
return new_mic