def window(v, func='hanning', params=None):
""" applies a windowing function to the 3D volume v (inplace, as reference) """
N = v.shape[0]
D = v.ndim
if any([d != N for d in list(v.shape)]) or D != 3:
raise Exception("Error: Volume is not Cube.")
def apply_seperable_window(v, w):
v *= np.reshape(w, (-1, 1, 1))
v *= np.reshape(w, (1, -1, 1))
v *= np.reshape(w, (1, 1, -1))
if func == "hanning":
w = np.hanning(N)
apply_seperable_window(v, w)
elif func == 'hamming':
w = np.hamming(N)
apply_seperable_window(v, w)
elif func == 'gaussian':
raise Exception('Unimplimented')
elif func == 'circle':
c = gencoords(N, 3)
if params == None:
r = N / 2 - 1
else:
r = params[0] * (N / 2 * 1)
v *= (np.sum(c**2, 1) < (r ** 2)).reshape((N, N, N))
elif func == 'box':
v[:, 0, 0] = 0.0
v[0, :, 0] = 0.0
v[0, 0, :] = 0.0
else:
raise Exception("Error: Window Type Not Supported")
def correlation_noise(num_images=1000, N=128, rad=0.6, stack_noise=False):
noise_stack = np.require(np.random.randn(num_images, N, N), dtype=density.real_t)
real_image = noise_stack[np.random.randint(num_images)]
corr_real_image = correlation.calc_full_ac(real_image, rad=rad)
fourier_noise = density.real_to_fspace(real_image)
fourier_corr_image = correlation.calc_full_ac(fourier_noise, rad=rad)
_, _, mask = geometry.gencoords(N, 2, rad, True)
plot_noise_histogram(corr_real_image, fourier_corr_image, mask, mask)
if stack_noise:
center = int(N/2)
x_shift, y_shift = np.random.randint(-center, center, size=2)
fourier_noise_stack = density.real_to_fspace(noise_stack, axes=(1, 2))
corr_noise_stack = np.zeros_like(noise_stack, dtype=density.real_t)
fourier_corr_noise_stack = np.zeros_like(fourier_noise_stack, dtype=density.complex_t)
for i in range(num_images):
corr_noise_stack[i] = correlation.calc_full_ac(noise_stack[i], rad=rad)
fourier_corr_noise_stack[i] = correlation.calc_full_ac(fourier_noise_stack[i], rad=rad)
noise_zoom = noise_stack[:, x_shift, y_shift]
fourier_noise_zoom = fourier_noise_stack[:, x_shift, y_shift]
plot_stack_noise(noise_zoom, fourier_noise_zoom)
def generate_phantom_density(N, window, sigma, num_blobs, seed=None):
if seed is not None:
np.random.seed(seed)
M = np.zeros((N, N, N), dtype=np.float32)
coords = gencoords(N, 3).reshape((N**3, 3))
inside_window = np.sum(coords**2, axis=1).reshape((N, N, N)) < window**2
curr_c = np.array([0.0, 0.0, 0.0])
curr_n = 0
while curr_n < num_blobs:
csigma = sigma * np.exp(0.25 * np.random.randn())
radM = np.sum((coords - curr_c.reshape((1, 3)))
** 2, axis=1).reshape((N, N, N))
inside = np.logical_and(radM < (3 * csigma)**2, inside_window)
# M[inside] = 1
M[inside] += np.exp(-0.5 * (radM[inside] / csigma**2))
curr_n += 1
curr_dir = np.random.randn(3)
curr_dir /= np.sum(curr_dir**2)
curr_c += 2.0 * csigma * curr_dir
curr_w = np.sqrt(np.sum(curr_c**2))
while curr_w > window:
curr_n_dir = curr_c / curr_w
curr_r_dir = (2 * np.dot(curr_dir, curr_n_dir)) * \
curr_n_dir - curr_dir
curr_c = curr_n_dir + (curr_w - window) * curr_r_dir
curr_w = np.sqrt(np.sum(curr_c**2))
return M
def generate_phantom_density(N, window, sigma, num_blobs, seed=None):
if seed is not None:
np.random.seed(seed)
M = np.zeros((N, N, N), dtype=np.float32)
coords = gencoords(N, 3).reshape((N**3, 3))
inside_window = np.sum(coords**2, axis=1).reshape((N, N, N)) < window**2
curr_c = np.array([0.0, 0.0, 0.0])
curr_n = 0
while curr_n < num_blobs:
csigma = sigma * np.exp(0.25 * np.random.randn())
radM = np.sum((coords - curr_c.reshape((1, 3)))**2, axis=1).reshape(
(N, N, N))
inside = np.logical_and(radM < (3 * csigma)**2, inside_window)
# M[inside] = 1
M[inside] += np.exp(-0.5 * (radM[inside] / csigma**2))
curr_n += 1
curr_dir = np.random.randn(3)
curr_dir /= np.sum(curr_dir**2)
curr_c += 2.0 * csigma * curr_dir
curr_w = np.sqrt(np.sum(curr_c**2))
while curr_w > window:
curr_n_dir = curr_c / curr_w
curr_r_dir = (2 * np.dot(curr_dir, curr_n_dir)) * \
curr_n_dir - curr_dir
curr_c = curr_n_dir + (curr_w - window) * curr_r_dir
curr_w = np.sqrt(np.sum(curr_c**2))
return M
def window(v, func='hanning', params=None):
""" applies a windowing function to the 3D volume v (inplace, as reference) """
N = v.shape[0]
D = v.ndim
if any([d != N for d in list(v.shape)]) or D != 3:
raise Exception("Error: Volume is not Cube.")
def apply_seperable_window(v, w):
v *= np.reshape(w, (-1, 1, 1))
v *= np.reshape(w, (1, -1, 1))
v *= np.reshape(w, (1, 1, -1))
if func == "hanning":
w = np.hanning(N)
apply_seperable_window(v, w)
elif func == 'hamming':
w = np.hamming(N)
apply_seperable_window(v, w)
elif func == 'gaussian':
raise Exception('Unimplimented')
elif func == 'circle':
c = gencoords(N, 3)
if params == None:
r = N / 2 - 1
else:
r = params[0] * (N / 2 * 1)
v *= (np.sum(c**2, 1) < (r**2)).reshape((N, N, N))
elif func == 'box':
v[:, 0, 0] = 0.0
v[0, :, 0] = 0.0
v[0, 0, :] = 0.0
else:
raise Exception("Error: Window Type Not Supported")
def get_envelope_map(self,sigma2,rho,env_lb=None,env_ub=None,minFreq=None,bfactor=None,rotavg=True):
N = self.cryodata.N
N_D = float(self.cryodata.N_D_Train)
num_batches = float(self.cryodata.num_batches)
psize = self.params['pixel_size']
beamstop_freq = self.params.get('beamstop_freq', None)
mean_corr = self.correlation_history.get_mean().reshape((N,N))
mean_power = self.power_history.get_mean().reshape((N,N))
mean_mask = self.mask_history.get_mean().reshape((N,N))
mask_w = self.mask_history.get_wsum() * (N_D / num_batches)
if rotavg:
mean_corr = cryoem.rotational_average(mean_corr,normalize=True,doexpand=True)
mean_power = cryoem.rotational_average(mean_power,normalize=True,doexpand=True)
mean_mask = cryoem.rotational_average(mean_mask,normalize=False,doexpand=True)
if isinstance(sigma2,np.ndarray):
sigma2 = sigma2.reshape((N,N))
if bfactor is not None:
coords = gencoords(N,2).reshape((N**2,2))
freqs = np.sqrt(np.sum(coords**2,axis=1))/(psize*N)
prior_envelope = ctf.envelope_function(freqs,bfactor).reshape((N,N))
else:
prior_envelope = 1.0
obsw = (mask_w * mean_mask / sigma2)
exp_env = (mean_corr * obsw + prior_envelope*rho) / (mean_power * obsw + rho)
if minFreq is not None:
# Only consider envelope parameters for frequencies above a threshold
minRad = minFreq*2.0*psize
if beamstop_freq is None:
_, _, minRadMask = gencoords(N, 2, minRad, True)
else:
maskRad = beamstop_freq * 2.0 * psize
_, _, minRadMask = gencoords_centermask(N, 2, minRad, maskRad, True)
exp_env[minRadMask.reshape((N,N))] = 1.0
if env_lb is not None or env_ub is not None:
np.clip(exp_env,env_lb,env_ub,out=exp_env)
return exp_env
def window_images(self, rad=0.99):
N = self.get_num_pixels()
coords = geometry.gencoords(N, 2).reshape((N**2, 2))
Cs = np.sum(coords**2, axis=1).reshape((N, N)
) > (rad * N / 2.0 - 1.5)**2
for img in self:
img[Cs] = 0
def window_images(self, rad=0.99):
N = self.get_num_pixels()
coords = geometry.gencoords(N, 2).reshape((N**2, 2))
Cs = np.sum(coords**2, axis=1).reshape((N, N)
) > (rad * N / 2.0 - 1.5)**2
for img in self:
img[Cs] = 0
def compute_shift_phases(pts, N, rad):
xy = geometry.gencoords(N, 2, rad)
N_T = xy.shape[0]
N_S = pts.shape[0]
shifts = np.empty((N_S, N_T), dtype=np.complex64)
for (i, (sx, sy)) in enumerate(pts):
shifts[i] = np.exp(2.0j * np.pi / N * (xy[:, 0] * sx + xy[:, 1] * sy))
return shifts
def compute_shift_phases(pts, N, rad):
xy = geometry.gencoords(N, 2, rad)
N_T = xy.shape[0]
N_S = pts.shape[0]
shifts = np.empty((N_S, N_T), dtype=np.complex64)
for (i, (sx, sy)) in enumerate(pts):
shifts[i] = np.exp(2.0j * np.pi / N * (xy[:, 0] * sx + xy[:, 1] * sy))
return shifts
def shift_vis(model):
N = model.shape[0]
rad = 0.8
kernel = 'lanczos'
kernsize = 4
xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True)
N_T = trunc_xy.shape[0]
premult = cryoops.compute_premultiplier(N,
kernel=kernel,
kernsize=kernsize)
TtoF = sincint.gentrunctofull(N=N, rad=rad)
fM = density.real_to_fspace(model)
prefM = density.real_to_fspace(
premult.reshape((1, 1, -1)) * premult.reshape(
(1, -1, 1)) * premult.reshape((-1, 1, 1)) * model)
pt = np.random.randn(3)
pt /= np.linalg.norm(pt)
psi = 2 * np.pi * np.random.rand()
ea = geometry.genEA(pt)[0]
ea[2] = psi
print('project model for Euler angel: ({:.2f}, {:.2f}, {:.2f}) degree'.
format(*np.rad2deg(ea)))
rot_matrix = geometry.rotmat3D_EA(*ea)[:, 0:2]
slop = cryoops.compute_projection_matrix([rot_matrix], N, kernel, kernsize,
rad, 'rots')
# trunc_slice = slop.dot(prefM.reshape((-1,)))
trunc_slice = cryoem.getslices(prefM, slop)
fourier_slice = TtoF.dot(trunc_slice).reshape(N, N)
real_proj = density.fspace_to_real(fourier_slice)
fig, axes = plt.subplots(4, 4, figsize=(12.8, 8))
im_real = axes[0, 0].imshow(real_proj)
im_fourier = axes[1, 0].imshow(np.log(np.abs(fourier_slice)))
for i, ax in enumerate(axes[:, 1:].T):
shift = np.random.randn(2) * (N / 4.0)
S = cryoops.compute_shift_phases(shift.reshape(1, 2), N, rad)[0]
shift_trunc_slice = S * trunc_slice
shift_fourier_slice = TtoF.dot(shift_trunc_slice).reshape(N, N)
shift_real_proj = density.fspace_to_real(shift_fourier_slice)
ax[0].imshow(shift_real_proj)
ax[1].imshow(np.log(np.abs(shift_fourier_slice)))
ax[2].imshow(np.log(shift_fourier_slice.real))
ax[3].imshow(np.log(shift_fourier_slice.imag))
fig.tight_layout()
plt.show()
def get_envelope_map(self,sigma2,rho,env_lb=None,env_ub=None,minFreq=None,bfactor=None,rotavg=True):
N = self.cryodata.N
N_D = float(self.cryodata.N_D_Train)
num_batches = float(self.cryodata.num_batches)
psize = self.params['pixel_size']
mean_corr = self.correlation_history.get_mean().reshape((N,N))
mean_power = self.power_history.get_mean().reshape((N,N))
mean_mask = self.mask_history.get_mean().reshape((N,N))
mask_w = self.mask_history.get_wsum() * (N_D / num_batches)
if rotavg:
mean_corr = cryoem.rotational_average(mean_corr,normalize=True,doexpand=True)
mean_power = cryoem.rotational_average(mean_power,normalize=True,doexpand=True)
mean_mask = cryoem.rotational_average(mean_mask,normalize=False,doexpand=True)
if isinstance(sigma2,np.ndarray):
sigma2 = sigma2.reshape((N,N))
if bfactor is not None:
coords = gencoords(N,2).reshape((N**2,2))
freqs = np.sqrt(np.sum(coords**2,axis=1))/(psize*N)
prior_envelope = ctf.envelope_function(freqs,bfactor).reshape((N,N))
else:
prior_envelope = 1.0
obsw = (mask_w * mean_mask / sigma2)
exp_env = (mean_corr * obsw + prior_envelope*rho) / (mean_power * obsw + rho)
if minFreq is not None:
# Only consider envelope parameters for frequencies above a threshold
minRad = minFreq*2.0*psize
_, _, minRadMask = gencoords(N, 2, minRad, True)
exp_env[minRadMask.reshape((N,N))] = 1.0
if env_lb is not None or env_ub is not None:
np.clip(exp_env,env_lb,env_ub,out=exp_env)
return exp_env
def rotational_expand(vals, N, D, interp_order=1):
interp_coords = np.sqrt(np.sum(gencoords(N, D).reshape(
(N**D, D))**2, axis=1)).reshape((1,) + D * (N,))
if np.iscomplexobj(vals):
rotexp = 1.0j * spinterp.map_coordinates(vals.imag, interp_coords,
order=interp_order, mode='nearest')
rotexp += spinterp.map_coordinates(vals.real, interp_coords,
order=interp_order, mode='nearest')
else:
rotexp = spinterp.map_coordinates(vals, interp_coords,
order=interp_order, mode='nearest')
return rotexp
def envelope_vis(N=128, rad=1):
trunc_xy = geometry.gencoords(N, 2, rad)
psize = 2.8
bfactor = 500
freqs = np.sqrt(np.sum(trunc_xy**2, axis=1)) / (psize * N)
envelope = ctf.envelope_function(freqs, bfactor)
fig, ax = plt.subplots()
ax.plot(freqs, envelope, '.', label='bfactor: %d' % bfactor)
ax.legend(frameon=False)
ax.set_title('envelope')
plt.show()
def compute_density_moments(M, mu=None):
N = M.shape[0]
absM = (M**2).reshape((N**3, 1))
absM /= np.sum(absM)
coords = gencoords(N, 3).reshape((N**3, 3))
if mu == None:
wcoords = coords.reshape((N**3, 3)) * absM
mu = np.sum(wcoords, axis=0).reshape((1, 3))
wccoords = np.sqrt(absM / N**3) * (coords - mu)
covar = np.dot(wccoords.T, wccoords)
return mu, covar
def rotational_average(M,
maxRadius=None,
doexpand=False,
normalize=True,
return_cnt=False):
N = M.shape[0]
D = len(M.shape)
assert D >= 2, 'Cannot rotationally average a 1D array'
pts = gencoords(N, D).reshape((N**D, D))
r = np.sqrt(np.sum(pts**2, axis=1)).reshape(M.shape)
ir = np.require(np.floor(r), dtype='uint32')
f = r - ir
if maxRadius is None:
maxRadius = np.ceil(np.sqrt(D) * N / D)
if maxRadius < np.max(ir) + 2:
valid_ir = ir + 1 < maxRadius
ir = ir[valid_ir]
f = f[valid_ir]
M = M[valid_ir]
if np.iscomplexobj(M):
raps = 1.0j * np.bincount(ir, weights=(1 - f) * M.imag, minlength=maxRadius) + \
np.bincount(ir + 1, weights=f * M.imag, minlength=maxRadius)
raps += np.bincount(ir, weights=(1 - f) * M.real, minlength=maxRadius) + \
np.bincount(ir + 1, weights=f * M.real, minlength=maxRadius)
else:
raps = np.bincount(ir, weights=(1 - f) * M, minlength=maxRadius) + \
np.bincount(ir + 1, weights=f * M, minlength=maxRadius)
raps = raps[0:maxRadius]
if normalize or return_cnt:
cnt = np.bincount(ir, weights=(1 - f), minlength=maxRadius) + \
np.bincount(ir + 1, weights=f, minlength=maxRadius)
cnt = cnt[0:maxRadius]
if normalize:
raps[cnt <= 0] = 0
raps[cnt > 0] /= cnt[cnt > 0]
if doexpand:
raps = rotational_expand(raps, N, D)
if return_cnt:
return raps, cnt
else:
return raps
def compute_density_moments(M, mu=None):
N = M.shape[0]
absM = (M**2).reshape((N**3, 1))
absM /= np.sum(absM)
coords = gencoords(N, 3).reshape((N**3, 3))
if mu == None:
wcoords = coords.reshape((N**3, 3)) * absM
mu = np.sum(wcoords, axis=0).reshape((1, 3))
wccoords = np.sqrt(absM / N**3) * (coords - mu)
covar = np.dot(wccoords.T, wccoords)
return mu, covar
def no_correlation(num_images=1000, N=128, rad=0.8):
center = int(N/2)
x_shift, y_shift = np.random.randint(-center, center, size=2)
noise_stack = np.require(np.random.randn(num_images, N, N), dtype=density.real_t)
real_image = noise_stack[np.random.randint(num_images)]
fourier_noise_stack = density.real_to_fspace(noise_stack, axes=(1, 2))
fourier_noise = density.real_to_fspace(real_image)
noise_zoom = noise_stack[:, x_shift, y_shift]
fourier_noise_zoom = fourier_noise_stack[:, x_shift, y_shift]
_, _, mask = geometry.gencoords(N, 2, rad, True)
plot_noise_histogram(real_image, fourier_noise, rmask=mask, fmask=mask)
plot_stack_noise(noise_zoom, fourier_noise_zoom)
def float_images(self, rad=0.99):
N = self.get_num_pixels()
coords = geometry.gencoords(N, 2).reshape((N**2, 2))
Cs = np.sum(coords**2, axis=1).reshape((N, N)) \
> (rad * N / 2.0 - 1.5)**2
vals = []
for img in self:
corner_pixels = img[Cs]
float_val = np.mean(corner_pixels)
img -= float_val
vals.append(float_val)
return vals
def float_images(self, rad=0.99):
N = self.get_num_pixels()
coords = geometry.gencoords(N, 2).reshape((N**2, 2))
Cs = np.sum(coords**2, axis=1).reshape((N, N)) \
> (rad * N / 2.0 - 1.5)**2
vals = []
for img in self:
corner_pixels = img[Cs]
float_val = np.mean(corner_pixels)
img -= float_val
vals.append(float_val)
return vals
def rotational_average(M, maxRadius=None, doexpand=False, normalize=True, return_cnt=False):
N = M.shape[0]
D = len(M.shape)
assert D >= 2, 'Cannot rotationally average a 1D array'
pts = gencoords(N, D).reshape((N**D, D))
r = np.sqrt(np.sum(pts**2, axis=1)).reshape(M.shape)
ir = np.require(np.floor(r), dtype='uint32')
f = r - ir
if maxRadius is None:
maxRadius = np.ceil(np.sqrt(D) * N / D)
if maxRadius < np.max(ir) + 2:
valid_ir = ir + 1 < maxRadius
ir = ir[valid_ir]
f = f[valid_ir]
M = M[valid_ir]
if np.iscomplexobj(M):
raps = 1.0j * np.bincount(ir, weights=(1 - f) * M.imag, minlength=maxRadius) + \
np.bincount(ir + 1, weights=f * M.imag, minlength=maxRadius)
raps += np.bincount(ir, weights=(1 - f) * M.real, minlength=maxRadius) + \
np.bincount(ir + 1, weights=f * M.real, minlength=maxRadius)
else:
raps = np.bincount(ir, weights=(1 - f) * M, minlength=maxRadius) + \
np.bincount(ir + 1, weights=f * M, minlength=maxRadius)
raps = raps[0:maxRadius]
if normalize or return_cnt:
cnt = np.bincount(ir, weights=(1 - f), minlength=maxRadius) + \
np.bincount(ir + 1, weights=f, minlength=maxRadius)
cnt = cnt[0:maxRadius]
if normalize:
raps[cnt <= 0] = 0
raps[cnt > 0] /= cnt[cnt > 0]
if doexpand:
raps = rotational_expand(raps, N, D)
if return_cnt:
return raps, cnt
else:
return raps
def rotate_density(M, R, t=None, upsamp=1.0):
assert len(M.shape) == 3
N = M.shape[0]
Nup = int(np.round(N * upsamp))
# print "Upsampling by", upsamp, "to", Nup, "^3"
coords = gencoords(Nup, 3).reshape((Nup**3, 3)) / float(upsamp)
if t is None:
interp_coords = np.transpose(np.dot(coords, R.T)).reshape(
(3, Nup, Nup, Nup)) + N / 2
else:
interp_coords = np.transpose(
np.dot(coords, R.T) + t).reshape((3, Nup, Nup, Nup)) + N / 2
out = spinterp.map_coordinates(M, interp_coords, order=1)
return out
def view_rad_range(N=128):
fig, axes = plt.subplots(4, 5, figsize=(12.8, 8)) # , sharex=True, sharey=True, squeeze=False)
rad_list = np.arange(0.1, 1.1, step=0.2)
for i, rad in enumerate(rad_list):
TtoF = sincint.gentrunctofull(N, rad)
xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True)
N_T = trunc_xy.shape[0]
trunc = np.arange(0, N_T)
image = TtoF.dot(trunc).reshape(N, N)
axes[0, i].imshow(image, origin='lower')
axes[0, i].set_title('radius: {:.2f}'.format(rad))
# outside of radius
xy_outside = xy[~truncmask]
image_outside_rad = np.zeros((N, N))
image_outside_rad[~truncmask.reshape(N, N)] = np.arange(xy_outside.shape[0])
axes[1, i].imshow(image_outside_rad, origin='lower')
# sort trunc_xy coordinates
pol_trunc_xy = correlation.cart2pol(trunc_xy)
sorted_idx = np.lexsort((pol_trunc_xy[:, 1], pol_trunc_xy[:, 0])) # lexsort; first, sort rho; second, sort theta
pol_trunc = trunc[sorted_idx.argsort()]
pol_image = TtoF.dot(pol_trunc).reshape(N, N)
axes[2, i].imshow(pol_image, origin='lower')
# pol coordinate in outside part
pol_xy_outside = correlation.cart2pol(xy_outside)
outside_sorted_idx = np.lexsort((pol_xy_outside[:, 1], pol_xy_outside[:, 0]))
pol_image_outside = np.zeros((N, N))
pol_image_outside[~truncmask.reshape(N, N)] = np.arange(pol_xy_outside.shape[0])[outside_sorted_idx.argsort()]
axes[3, i].imshow(pol_image_outside, origin='lower')
for i, ax in enumerate(axes.flat):
m, n = np.unravel_index(i, (4, 5))
if m != 3:
ax.set_xticks([])
else:
ax.set_xticks([0, int(N/4), int(N/2), int(N*3/4), int(N-1)])
if n != 0:
ax.set_yticks([])
else:
ax.set_yticks([0, int(N/4), int(N/2), int(N*3/4), int(N-1)])
# fig.savefig('cart_coordinate_view_to_polar_coordinate_view.png', dpi=300)
plt.show()
def rotate_density(M, R, t=None, upsamp=1.0):
assert len(M.shape) == 3
N = M.shape[0]
Nup = int(np.round(N * upsamp))
# print "Upsampling by", upsamp, "to", Nup, "^3"
coords = gencoords(Nup, 3).reshape((Nup**3, 3)) / float(upsamp)
if t is None:
interp_coords = np.transpose(np.dot(coords, R.T)).reshape(
(3, Nup, Nup, Nup)) + int(N / 2)
else:
interp_coords = np.transpose(np.dot(coords, R.T) + t).reshape(
(3, Nup, Nup, Nup)) + int(N / 2)
out = spinterp.map_coordinates(M, interp_coords, order=1)
return out
def dataset_loading_test(params, visualize=False):
imgpath = params['inpath']
psize = params['resolution']
imgstk = MRCImageStack(imgpath, psize)
# if params.get('float_images', True):
# imgstk.float_images()
ctfpath = params['ctfpath']
mscope_params = params['microscope_params']
ctfstk = CTFStack(ctfpath, mscope_params)
cryodata = CryoDataset(imgstk, ctfstk)
cryodata.compute_noise_statistics()
# if params.get('window_images',True):
# imgstk.window_images()
cryodata.divide_dataset(params['minisize'], params['test_imgs'],
params['partition'], params['num_partitions'], params['random_seed'])
# cryodata.set_datasign(params.get('datasign', 'auto'))
# if params.get('normalize_data',True):
# cryodata.normalize_dataset()
# voxel_size = cryodata.pixel_size
N = cryodata.imgstack.get_num_pixels()
fspace_stack = FourierStack(cryodata.imgstack,
caching = True, zeropad=1)
premult = cryoops.compute_premultiplier(N + 2 * int(1 * (N/2)), 'lanczos', 8)
premult = premult.reshape((-1,1)) * premult.reshape((1,-1))
fspace_stack.set_transform(premult, 1)
if visualize:
rad = 0.99
coords = geometry.gencoords(N, 2).reshape((N**2, 2))
Cs = np.sum(coords**2, axis=1).reshape((N, N)) > (rad * N / 2.0 - 1.5)**2
idx = np.random.randint(cryodata.imgstack.num_images)
normalized = cryodata.imgstack.get_image(1)
f_normalized = fspace_stack.get_image(1)
plot_noise_histogram(normalized, f_normalized,
rmask=~Cs, fmask=None, plot_unmask=False)
plt.show()
return cryodata, fspace_stack
def compute_symmetry_errs(rs,symop,gfV,gfVnopre,radwn,stream=None):
N = gfV[0].shape[0]
offRs = geometry.expmaps(rs)
Rsyms = symop.get_rotations(include_identity=False)
pts = geometry.gencoords(N, 3, radwn).reshape((-1,3))
N_R = offRs.shape[0]
N_S = 1
N_T = pts.shape[0]
N_T_aligned = int(n.ceil(float(N_T)/(2*cudaworker.cukrns.blocksize)))*(2*cudaworker.cukrns.blocksize)
print(' ')
print('NR, NS, NT ', N_R, N_S, N_T, N_T_aligned, (N_R*N_T_aligned)/1.e9)
with cudaworker.GPUContext():
gts = gpuarray.zeros((1,3),n.float32)
gret = gpuarray.zeros((N_R,N_S),n.float32)
gphis = gpuarray.empty((N_R,N_S),n.float32)
gP = (gpuarray.empty((N_R,N_T_aligned),n.float32),
gpuarray.empty((N_R,N_T_aligned),n.float32))
gQ = (gpuarray.empty((1,N_T_aligned),n.float32),
gpuarray.empty((1,N_T_aligned),n.float32))
gpts = gpuarray.to_gpu(pts.reshape((-1,3)).astype(n.float32))
gI = gpuarray.to_gpu(n.identity(3,dtype=n.float32).reshape((1,9)))
cudaworker.cukrns.slice('linear',2,gQ,gfVnopre,gpts, gI,None,1,stream=stream)
for Rsym in Rsyms:
cRsym = Rsym.reshape((3,3))
Rs = n.array([cR.reshape((3,3)).dot(cRsym.dot(cR.reshape((3,3)).T))
for cR in offRs],dtype=n.float32)
gRs = gpuarray.to_gpu(Rs.reshape((-1,9)).astype(n.float32))
cudaworker.cukrns.slice('linear',2,gP,gfV,gpts,gRs,None,1,stream=stream)
cudaworker.cukrns.shifted_squared_error(N_R, N_S, N_T, gphis.reshape((N_R*N_S)), gP, gQ, gts, gpts, 2.0*n.pi/N, None, broadcast_Q=True, stream=stream)
if stream is not None:
stream.synchronize()
gret += gphis
ret = gret.get()
return ret.reshape((N_R))
def demo(N=128, rad=0.5):
TtoF = sincint.gentrunctofull(N=N, rad=rad)
xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True)
print('shape of TtoF:', TtoF.shape)
print('slice shape:', trunc_xy.shape[0])
trunc_slice = np.arange(trunc_xy.shape[0])
sliced_image = TtoF.dot(trunc_slice).reshape(N, N)
trunc_xy_idx = np.int_(trunc_xy + int(N/2))
# Compare speed for getting slices in this way
new_trunc_slice = sliced_image[trunc_xy_idx[:, 0], trunc_xy_idx[:, 1]]
print('error:', sum(trunc_slice - new_trunc_slice))
pol_trunc_xy = correlation.cart2pol(trunc_xy)
# inside of rad
# sort trunc_xy coordinates
sorted_idx = np.lexsort((pol_trunc_xy[:, 1], pol_trunc_xy[:, 0])) # lexsort; first, sort rho; second, sort theta
sorted_pol_trunc_xy = pol_trunc_xy[sorted_idx]
# reconstuct sorted coordinates into original state
reco_pol_trunc_xy = sorted_pol_trunc_xy[sorted_idx.argsort()]
print('error for reconstructed coordinates:', sum(correlation.pol2cart(reco_pol_trunc_xy) - trunc_xy))
reco_trunc_slice = trunc_slice[sorted_idx.argsort()]
bingo_sliced_image = TtoF.dot(reco_trunc_slice).reshape(N, N)
# outside of rad
xy_outside = xy[~truncmask]
sliced_image_outside_rad = np.zeros((N, N))
sliced_image_outside_rad[~truncmask.reshape(N, N)] = np.arange(xy_outside.shape[0])
pol_xy_outside = correlation.cart2pol(xy_outside)
outside_sorted_idx = np.lexsort((pol_xy_outside[:, 1], pol_xy_outside[:, 0])) # lexsort; first, sort rho; second, sort theta
sorted_pol_xy_outside = pol_xy_outside[outside_sorted_idx]
reco_pol_xy_outside = np.arange(xy_outside.shape[0])[outside_sorted_idx.argsort()]
bingo_sliced_image_outside_rad = np.zeros((N, N))
bingo_sliced_image_outside_rad[~truncmask.reshape(N, N)] = reco_pol_xy_outside
fig, axes = plt.subplots(2, 2)
ax = axes.flatten()
ax[0].imshow(sliced_image)
ax[1].imshow(bingo_sliced_image)
ax[2].imshow(sliced_image_outside_rad)
ax[3].imshow(bingo_sliced_image_outside_rad)
plt.show()
def compare_interpolation(N=128, rad=1):
_, trunc_xy, _ = geometry.gencoords(N, 2, rad, True)
pol_trunc_xy = correlation.cart2pol(trunc_xy)
sorted_idx = np.lexsort((pol_trunc_xy[:, 1], pol_trunc_xy[:, 0])) # lexsort; first, sort rho; second, sort theta
sorted_pol_trunc_xy = pol_trunc_xy[sorted_idx]
interpolation = ['none', 'nearest', 'nearest_decimal_1', 'nearest_half']
fig, ax = plt.subplots(nrows=len(interpolation), sharex=True)
# fig, ax = plt.subplots()
def round_to(n, precision):
# correction = 0.5 if n >= 0 else -0.5
correction = np.ones_like(n) * 0.5
correction[n < 0] = -0.5
return np.int_(n / precision + correction) * precision
def round_half(n):
return round_to(n, 0.5)
def get_ip_func(ip_method):
if 'none' == ip_method.lower():
return lambda x: x
elif 'nearest' == ip_method.lower():
return np.round
elif 'nearest_decimal_1' == ip_method.lower():
return lambda x: np.round(x, 1)
elif 'nearest_half' == ip_method.lower():
return round_half
else:
raise ValueError('please input correct interpolation method.')
for i, ip in enumerate(interpolation):
ip_func = get_ip_func(ip)
ip_pol_xy = ip_func(sorted_pol_trunc_xy[:, 0])
unique_value, unique_index, unique_inverse, unique_counts = np.unique(ip_pol_xy,
return_index=True, return_inverse=True, return_counts=True)
ax[i].plot(unique_value, unique_counts, label='interpolation: {}'.format(ip))
ax[i].legend(frameon=False)
ax[i].set_ylabel('counts')
ax[-1].set_xlabel('radius')
plt.show()
def compute_fsc(VF1, VF2, maxrad, width=1.0, thresholds=[0.143, 0.5]):
assert VF1.shape == VF2.shape
N = VF1.shape[0]
r = np.sqrt(np.sum(gencoords(N, 3).reshape((N, N, N, 3))**2, axis=3))
prev_rad = -np.inf
fsc = []
rads = []
resInd = len(thresholds) * [None]
for i, rad in enumerate(np.arange(1.5, maxrad * N / 2.0, width)):
cxyz = np.logical_and(r >= prev_rad, r < rad)
cF1 = VF1[cxyz]
cF2 = VF2[cxyz]
if len(cF1) == 0:
break
cCorr = np.vdot(cF1, cF2) / np.sqrt(
np.vdot(cF1, cF1) * np.vdot(cF2, cF2))
for j, thr in enumerate(thresholds):
if cCorr < thr and resInd[j] is None:
resInd[j] = i
fsc.append(cCorr.real)
rads.append(rad / (N / 2.0))
prev_rad = rad
fsc = np.array(fsc)
rads = np.array(rads)
resolutions = []
for rI, thr in zip(resInd, thresholds):
if rI is None:
resolutions.append(rads[-1])
elif rI == 0:
resolutions.append(np.inf)
else:
x = (thr - fsc[rI]) / (fsc[rI - 1] - fsc[rI])
resolutions.append(x * rads[rI - 1] + (1 - x) * rads[rI])
return rads, fsc, thresholds, resolutions
def estimate_noise_variance(self, esttype='robust', zerosub=False, rad=1.0):
N = self.get_num_pixels()
Cs = np.sum(geometry.gencoords(N, 2).reshape((N**2, 2))**2,
axis=1).reshape((N, N)) > (rad * N / 2.0 - 1.5)**2
vals = []
for img in self:
cvals = img[Cs]
vals.append(cvals)
if esttype == 'robust':
if zerosub:
var = (
1.4826 * np.median(np.abs(np.asarray(vals) - np.median(vals))))**2
else:
var = (1.4826 * np.median(np.abs(vals)))**2
elif esttype == 'mle':
var = np.mean(np.asarray(vals)**2, dtype=np.float64)
if zerosub:
var -= np.mean(vals, dtype=np.float64)**2
return var
def estimate_noise_variance(self, esttype='robust', zerosub=False, rad=1.0):
N = self.get_num_pixels()
Cs = np.sum(geometry.gencoords(N, 2).reshape((N**2, 2))**2,
axis=1).reshape((N, N)) > (rad * N / 2.0 - 1.5)**2
vals = []
for img in self:
cvals = img[Cs]
vals.append(cvals)
if esttype == 'robust':
if zerosub:
var = (
1.4826 * np.median(np.abs(np.asarray(vals) - np.median(vals))))**2
else:
var = (1.4826 * np.median(np.abs(vals)))**2
elif esttype == 'mle':
var = np.mean(np.asarray(vals)**2, dtype=np.float64)
if zerosub:
var -= np.mean(vals, dtype=np.float64)**2
return var
def compute_fsc(VF1, VF2, maxrad, width=1.0, thresholds=[0.143, 0.5]):
assert VF1.shape == VF2.shape
N = VF1.shape[0]
r = np.sqrt(np.sum(gencoords(N, 3).reshape((N, N, N, 3))**2, axis=3))
prev_rad = -np.inf
fsc = []
rads = []
resInd = len(thresholds) * [None]
for i, rad in enumerate(np.arange(1.5, maxrad * N / 2.0, width)):
cxyz = np.logical_and(r >= prev_rad, r < rad)
cF1 = VF1[cxyz]
cF2 = VF2[cxyz]
if len(cF1) == 0:
break
cCorr = np.vdot(cF1, cF2) / np.sqrt(np.vdot(cF1, cF1) * np.vdot(cF2, cF2))
for j, thr in enumerate(thresholds):
if cCorr < thr and resInd[j] is None:
resInd[j] = i
fsc.append(cCorr.real)
rads.append(rad / (N / 2.0))
prev_rad = rad
fsc = np.array(fsc)
rads = np.array(rads)
resolutions = []
for rI, thr in zip(resInd, thresholds):
if rI is None:
resolutions.append(rads[-1])
elif rI == 0:
resolutions.append(np.inf)
else:
x = (thr - fsc[rI]) / (fsc[rI - 1] - fsc[rI])
resolutions.append(x * rads[rI - 1] + (1 - x) * rads[rI])
return rads, fsc, thresholds, resolutions
def merge_slices(slices, Rs, N, rad, beamstop_rad=None, res=None):
center = int(N / 2)
if beamstop_rad is None:
coords = gencoords(N, 2, rad)
else:
coords = gencoords_centermask(N, 2, rad, beamstop_rad)
assert slices.shape[1] == coords.shape[0]
if res is None:
res = np.zeros((N, ) * 3, dtype=np.float32)
else:
assert res.shape == (N, ) * 3
assert res.dtype == slices.dtype
res[:] = 0.0
model_weight = np.zeros((N, ) * 3)
for i, R in enumerate(Rs):
curr_slices = slices[i, :]
for j, xy in enumerate(coords):
voxel_intensity = curr_slices[j]
rot_coord = R.dot(xy.T).reshape(1, -1)[0] + center
rot_coord = np.int_(np.round(rot_coord))
in_x = rot_coord[0] >= 0 and rot_coord[0] < N
in_y = rot_coord[1] >= 0 and rot_coord[1] < N
in_z = rot_coord[2] >= 0 and rot_coord[2] < N
if in_x and in_y and in_z:
index_coord = tuple(rot_coord)
model_voxel_intensity = res[index_coord]
model_weight[index_coord] += 1
voxel_weight = model_weight[index_coord]
delta_intensity = voxel_intensity - model_voxel_intensity
model_voxel_intensity += delta_intensity / voxel_weight
res[index_coord] = model_voxel_intensity
return res
def getslices_interp(V, Rs, rad, beamstop_rad=None, res=None):
ndim = V.ndim
assert ndim > 1
num_slices = len(Rs)
# if ndim == 2:
# assert Rs.shape[1] == 2
# elif ndim == 3:
# assert Rs.shape[1] == 3
# Rs.shape[2] == 2
N = V.shape[0]
center = int(N / 2)
if beamstop_rad is None:
coords = gencoords(N, 2, rad)
else:
coords = gencoords_centermask(N, 2, rad, beamstop_rad)
N_T = coords.shape[0]
grid = (np.arange(N), ) * ndim
slicing_func = RegularGridInterpolator(grid,
V,
bounds_error=False,
fill_value=0.0)
if res is None:
res = np.zeros((num_slices, N_T), dtype=V.dtype)
else:
assert res.shape[0] == Rs.shape[0]
assert res.dtype == V.dtype
res[:] = 0
for i, R in enumerate(Rs):
rotated_coords = R.dot(coords.T).T + center
res[i] = slicing_func(rotated_coords)
# res[i] = interpn(grid, V, rotated_coords)
# res[i] = spinterp.map_coordinates(V, rotated_coords.T)
return res
def __init__(self, model, dataset_params, ctf_params,
interp_params={'kern': 'lanczos', 'kernsize': 4.0, 'zeropad': 0, 'dopremult': True},
load_cache=True):
self.dataset_params = dataset_params
if model is not None:
# assert False
assert isinstance(model, np.ndarray), "Unexpected data type for input model"
self.num_pixels = model.shape[0]
N = self.num_pixels
self.num_images = dataset_params['num_images']
assert self.num_images > 1, "it's better to make num_images larger than 1."
self.pixel_size = float(dataset_params['pixel_size'])
euler_angles = dataset_params['euler_angles']
self.is_sym = get_symmetryop(dataset_params.get('symmetry', None))
if euler_angles is None and self.is_sym is None:
pt = np.random.randn(self.num_images, 3)
pt /= np.linalg.norm(pt, axis=1, keepdims=True)
euler_angles = geometry.genEA(pt)
euler_angles[:, 2] = 2 * np.pi * np.random.rand(self.num_images)
elif euler_angles is None and self.is_sym is not None:
euler_angles = np.zeros((self.num_images, 3))
for i, ea in enumerate(euler_angles):
while True:
pt = np.random.randn(3)
pt /= np.linalg.norm(pt)
if self.is_sym.in_asymunit(pt.reshape(-1, 3)):
break
ea[0:2] = geometry.genEA(pt)[0][0:2]
ea[2] = 2 * np.pi * np.random.rand()
self.euler_angles = euler_angles.reshape((-1, 3))
if ctf_params is not None:
self.use_ctf = True
ctf_map = ctf.compute_full_ctf(None, N, ctf_params['psize'],
ctf_params['akv'], ctf_params['cs'], ctf_params['wgh'],
ctf_params['df1'], ctf_params['df2'], ctf_params['angast'],
ctf_params['dscale'], ctf_params.get('bfactor', 500))
self.ctf_params = copy(ctf_params)
if 'bfactor' in self.ctf_params.keys():
self.ctf_params.pop('bfactor')
else:
self.use_ctf = False
ctf_map = np.ones((N**2,), dtype=density.real_t)
kernel = 'lanczos'
ksize = 6
rad = 0.95
# premult = cryoops.compute_premultiplier(N, kernel, ksize)
TtoF = sincint.gentrunctofull(N=N, rad=rad)
base_coords = geometry.gencoords(N, 2, rad)
# premulter = premult.reshape((1, 1, -1)) \
# * premult.reshape((1, -1, 1)) \
# * premult.reshape((-1, 1, 1))
# fM = density.real_to_fspace(premulter * model)
fM = model
# if load_cache:
# try:
print("Generating Dataset ... :")
tic = time.time()
imgdata = np.empty((self.num_images, N, N), dtype=density.real_t)
for i, ea in zip(range(self.num_images), self.euler_angles):
R = geometry.rotmat3D_EA(*ea)[:, 0:2]
slop = cryoops.compute_projection_matrix(
[R], N, kernel, ksize, rad, 'rots')
# D = slop.dot(fM.reshape((-1,)))
rotated_coords = R.dot(base_coords.T).T + int(N/2)
D = interpn((np.arange(N),) * 3, fM, rotated_coords)
np.maximum(D, 0.0, out=D)
intensity = ctf_map.reshape((N, N)) * TtoF.dot(D).reshape((N, N))
np.maximum(1e-8, intensity, out=intensity)
intensity = np.float_( np.random.poisson(intensity) )
imgdata[i] = np.require(intensity, dtype=density.real_t)
self.imgdata = imgdata
print(" cost {} seconds.".format(time.time()-tic))
self.set_transform(interp_params)
# self.prep_processing()
else:
euler_angles = []
with open(self.dataset_params['gtpath']) as par:
par.readline()
# 'C PHI THETA PSI SHX SHY FILM DF1 DF2 ANGAST'
while True:
try:
line = par.readline().split()
euler_angles.append([float(line[1]), float(line[2]), float(line[3])])
except Exception:
break
self.euler_angles = np.deg2rad(np.asarray(euler_angles))
num_images = self.dataset_params.get('num_images', 200)
imgdata = mrc.readMRCimgs(self.dataset_params['inpath'], 0, num_images)
self.imgdata = np.transpose(imgdata, axes=(2, 0, 1))
self.num_images = self.imgdata.shape[0]
self.num_pixels = self.imgdata.shape[1]
N = self.num_pixels
self.pixel_size = self.dataset_params['resolution']
self.is_sym = self.dataset_params.get('symmetry', None)
self.use_ctf = False
ctf_map = np.ones((N**2,), dtype=density.real_t)
self.set_transform(interp_params)
def compute_full_ctf(rots, N, psize, akv, csf, wgh, dfmid1, dfmid2, angastf, dscale, bfactor):
freqs = geometry.gencoords(N, 2) / (N * psize)
return compute_ctf(freqs, rots, akv, csf, wgh, dfmid1, dfmid2, angastf, dscale, bfactor)
def genphantomdata(N_D, phantompath):
mscope_params = {
'akv': 200,
'wgh': 0.07,
'cs': 2.0,
'psize': 3.0,
'bfactor': 500.0
}
M = mrc.readMRC(phantompath)
N = M.shape[0]
rad = 0.95
M_totalmass = 1000000
# M_totalmass = 1500000
kernel = 'lanczos'
ksize = 6
tic = time.time()
N_D = int(N_D)
N = int(N)
rad = float(rad)
psize = mscope_params['psize']
bfactor = mscope_params['bfactor']
M_totalmass = float(M_totalmass)
ctfparfile = 'particle/examplectfs.par'
srcctf_stack = CTFStack(ctfparfile, mscope_params)
genctf_stack = GeneratedCTFStack(
mscope_params, parfields=['PHI', 'THETA', 'PSI', 'SHX', 'SHY'])
TtoF = sincint.gentrunctofull(N=N, rad=rad)
Cmap = np.sort(
np.random.random_integers(0,
srcctf_stack.get_num_ctfs() - 1, N_D))
cryoem.window(M, 'circle')
M[M < 0] = 0
if M_totalmass is not None:
M *= M_totalmass / M.sum()
# oversampling
oversampling_factor = 3
psize = psize * oversampling_factor
V = density.real_to_fspace_with_oversampling(M, oversampling_factor)
fM = V.real**2 + V.imag**2
# mrc.writeMRC('particle/EMD6044_fM_totalmass_{}_oversampling_{}.mrc'.format(str(int(M_totalmass)).zfill(5), oversampling_factor), fM, psz=psize)
print("Generating data...")
sys.stdout.flush()
imgdata = np.empty((N_D, N, N), dtype=density.real_t)
pardata = {'R': []}
prevctfI = None
coords = geometry.gencoords(N, 2, rad)
slicing_func = RegularGridInterpolator((np.arange(N), ) * 3,
fM,
bounds_error=False,
fill_value=0.0)
for i, srcctfI in enumerate(Cmap):
ellapse_time = time.time() - tic
remain_time = float(N_D - i) * ellapse_time / max(i, 1)
print("\r%.2f Percent.. (Elapsed: %s, Remaining: %s)" %
(i / float(N_D) * 100.0, format_timedelta(ellapse_time),
format_timedelta(remain_time)),
end='')
sys.stdout.flush()
# Get the CTF for this image
cCTF = srcctf_stack.get_ctf(srcctfI)
if prevctfI != srcctfI:
genctfI = genctf_stack.add_ctf(cCTF)
C = cCTF.dense_ctf(N, psize, bfactor).reshape((N, N))
prevctfI = srcctfI
# Randomly generate the viewing direction/shift
pt = np.random.randn(3)
pt /= np.linalg.norm(pt)
psi = 2 * np.pi * np.random.rand()
EA = geometry.genEA(pt)[0]
EA[2] = psi
# Rotate coordinates and get slice image by interpolation
R = geometry.rotmat3D_EA(*EA)[:, 0:2]
rotated_coords = R.dot(coords.T).T + int(N / 2)
slice_data = slicing_func(rotated_coords)
intensity = TtoF.dot(slice_data)
np.maximum(intensity, 0.0, out=intensity)
# Add poisson noise
img = np.float_(np.random.poisson(intensity.reshape(N, N)))
np.maximum(1e-8, img, out=img)
imgdata[i] = np.require(img, dtype=density.real_t)
genctf_stack.add_img(genctfI,
PHI=EA[0] * 180.0 / np.pi,
THETA=EA[1] * 180.0 / np.pi,
PSI=EA[2] * 180.0 / np.pi,
SHX=0.0,
SHY=0.0)
pardata['R'].append(R)
print("\n\rDone in ", time.time() - tic, " seconds.")
return imgdata, genctf_stack, pardata, mscope_params
import sincint
M = mrc.readMRC('./particle/EMD-6044.mrc')
# M = mrc.readMRC('./particle/1AON.mrc')
# M = M / np.sum(M)
M = M[:124, :124, :124]
mrc.writeMRC('./particle/EMD-6044-cropped.mrc', M, psz=3.0)
N = M.shape[0]
print(M.shape)
rad = 1
kernel = 'lanczos'
ksize = 4
xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True)
# premult = cryoops.compute_premultiplier(N, kernel='lanczos', kernsize=6)
premult = cryoops.compute_premultiplier(N, kernel, ksize)
TtoF = sincint.gentrunctofull(N=N, rad=rad)
fM = density.real_to_fspace(M)
prefM = density.real_to_fspace(premult.reshape(
(1, 1, -1)) * premult.reshape((1, -1, 1)) * premult.reshape((-1, 1, 1)) * M)
EAs_grid = healpix.gen_EAs_grid(nside=2, psi_step=360)
Rs = [geometry.rotmat3D_EA(*EA)[:, 0:2] for EA in EAs_grid]
slice_ops = cryoops.compute_projection_matrix(Rs, N, kern='lanczos', kernsize=ksize, rad=rad, projdirtype='rots')
slices_sampled = cryoem.getslices(fM, slice_ops).reshape((EAs_grid.shape[0], trunc_xy.shape[0]))
premult_slices_sampled = cryoem.getslices(prefM, slice_ops).reshape((EAs_grid.shape[0], trunc_xy.shape[0]))
cs = 2.0
df1, df2, angast = 44722, 49349, 45.0 * (np.pi / 180.0)
dscale = 1.0
v1 = compute_ctf(fcoords, rots, akv, cs, wgh, df1,
df2, angast, dscale).reshape((-1,))
v2 = compute_ctf(rotfcoords, None, akv, cs, wgh, df1,
df2, angast, dscale).reshape((-1,))
# This being small confirms that using the rots parameter is equivalent to
# rotating the coordinates
print(np.abs(v1 - v2).max())
N = 512
psz = 5.6
rad = 0.25
fcoords = geometry.gencoords(N, 2, rad) / (N * psz)
ctf1_rot = compute_full_ctf(
rots, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None)
ctf2_full = compute_full_ctf(
None, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None)
P_rot = coops.compute_inplanerot_matrix(rots, N, 'lanczos', 10, rad)
ctf2_rot = P_rot.dot(ctf2_full).reshape((-1,))
P_null = coops.compute_inplanerot_matrix(np.array([0]), N, 'linear', 2, rad)
ctf1_rot = P_null.dot(ctf1_rot).reshape((-1,))
roterr = ctf1_rot - ctf2_rot
relerr = np.abs(roterr) / np.maximum(np.abs(ctf1_rot), np.abs(ctf2_rot))
# This being small confirms that compute_inplane_rotmatrix and rots use
def set_data(self,cparams,minibatch):
self.params = cparams
self.minibatch = minibatch
factoredRI = cparams.get('likelihood_factored_slicing',True)
max_freq = cparams['max_frequency']
psize = cparams['pixel_size']
rad_cutoff = cparams.get('rad_cutoff', 1.0)
rad = min(rad_cutoff,max_freq*2.0*psize)
self.xy, self.trunc_xy, self.truncmask = gencoords(self.N, 2, rad, True)
self.trunc_freq = np.require(self.trunc_xy / (self.N*psize), dtype=np.float32)
self.N_T = self.trunc_xy.shape[0]
interp_change = self.rad != rad or self.factoredRI != factoredRI
if interp_change:
print("Iteration {0}: freq = {3}, rad = {1}, N_T = {2}".format(cparams['iteration'], rad, self.N_T, max_freq))
self.rad = rad
self.factoredRI = factoredRI
# Setup the quadrature schemes
if not factoredRI:
self.set_proj_quad(rad)
else:
self.set_slice_quad(rad)
self.set_inplane_quad(rad)
# Check shift quadrature
self.set_shift_quad(rad)
# Setup inlier model
self.inlier_sigma2 = cparams['sigma']**2
base_sigma2 = self.cryodata.noise_var
if isinstance(self.inlier_sigma2,np.ndarray):
self.inlier_sigma2 = self.inlier_sigma2.reshape(self.truncmask.shape)
self.inlier_sigma2_trunc = self.inlier_sigma2[self.truncmask != 0]
self.inlier_const = (self.N_T/2.0)*np.log(2.0*np.pi) + 0.5*np.sum(np.log(self.inlier_sigma2_trunc))
else:
self.inlier_sigma2_trunc = self.inlier_sigma2
self.inlier_const = (self.N_T/2.0)*np.log(2.0*np.pi*self.inlier_sigma2)
# Compute the likelihood for the image content outside of rad
_,_,fspace_truncmask = gencoords(self.fspace_stack.get_num_pixels(), 2, rad*self.fspace_stack.get_num_pixels()/self.N, True)
self.imgpower = np.empty((self.minibatch['N_M'],),dtype=density.real_t)
self.imgpower_trunc = np.empty((self.minibatch['N_M'],),dtype=density.real_t)
for idx,Idx in enumerate(self.minibatch['img_idxs']):
Img = self.fspace_stack.get_image(Idx)
self.imgpower[idx] = np.sum(Img.real**2) + np.sum(Img.imag**2)
Img_trunc = Img[fspace_truncmask.reshape(Img.shape) == 0]
self.imgpower_trunc[idx] = np.sum(Img_trunc.real**2) + np.sum(Img_trunc.imag**2)
like_trunc = 0.5*self.imgpower_trunc/base_sigma2
self.inlier_like_trunc = like_trunc
self.inlier_const += ((self.N**2 - self.N_T)/2.0)*np.log(2.0*np.pi*base_sigma2)
# Setup the envelope function
envelope = self.params.get('exp_envelope',None)
if envelope is not None:
envelope = envelope.reshape((-1,))
envelope = envelope[self.truncmask != 0]
envelope = np.require(envelope,dtype=np.float32)
else:
bfactor = self.params.get('learn_like_envelope_bfactor',500.0)
if bfactor is not None:
freqs = np.sqrt(np.sum(self.trunc_xy**2,axis=1))/(psize*self.N)
envelope = ctf.envelope_function(freqs,bfactor)
self.envelope = envelope
import sincint
M = mrc.readMRC('./particle/EMD-6044.mrc')
# M = mrc.readMRC('./particle/1AON.mrc')
# M = M / np.sum(M)
M = M[:124, :124, :124]
mrc.writeMRC('./particle/EMD-6044-cropped.mrc', M, psz=3.0)
N = M.shape[0]
print(M.shape)
rad = 1
kernel = 'lanczos'
ksize = 4
xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True)
# premult = cryoops.compute_premultiplier(N, kernel='lanczos', kernsize=6)
premult = cryoops.compute_premultiplier(N, kernel, ksize)
TtoF = sincint.gentrunctofull(N=N, rad=rad)
fM = density.real_to_fspace(M)
prefM = density.real_to_fspace(
premult.reshape((1, 1, -1)) * premult.reshape(
(1, -1, 1)) * premult.reshape((-1, 1, 1)) * M)
EAs_grid = healpix.gen_EAs_grid(nside=2, psi_step=360)
Rs = [geometry.rotmat3D_EA(*EA)[:, 0:2] for EA in EAs_grid]
slice_ops = cryoops.compute_projection_matrix(Rs,
N,
kern='lanczos',
kernsize=ksize,
cs = 2.0
df1, df2, angast = 44722, 49349, 45.0 * (np.pi / 180.0)
dscale = 1.0
v1 = compute_ctf(fcoords, rots, akv, cs, wgh, df1,
df2, angast, dscale).reshape((-1,))
v2 = compute_ctf(rotfcoords, None, akv, cs, wgh, df1,
df2, angast, dscale).reshape((-1,))
# This being small confirms that using the rots parameter is equivalent to
# rotating the coordinates
print(np.abs(v1 - v2).max())
N = 512
psz = 5.6
rad = 0.25
fcoords = geometry.gencoords(N, 2, rad) / (N * psz)
ctf1_rot = compute_full_ctf(
rots, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None)
ctf2_full = compute_full_ctf(
None, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None)
P_rot = coops.compute_inplanerot_matrix(rots, N, 'lanczos', 10, rad)
ctf2_rot = P_rot.dot(ctf2_full).reshape((-1,))
P_null = coops.compute_inplanerot_matrix(np.array([0]), N, 'linear', 2, rad)
ctf1_rot = P_null.dot(ctf1_rot).reshape((-1,))
roterr = ctf1_rot - ctf2_rot
relerr = np.abs(roterr) / np.maximum(np.abs(ctf1_rot), np.abs(ctf2_rot))
# This being small confirms that compute_inplane_rotmatrix and rots use
def compute_full_ctf(rots, N, psize, akv, csf, wgh, dfmid1, dfmid2, angastf, dscale, bfactor):
freqs = geometry.gencoords(N, 2) / (N * psize)
return compute_ctf(freqs, rots, akv, csf, wgh, dfmid1, dfmid2, angastf, dscale, bfactor)
def ctf_vis():
import cryoops as coops
fcoords = np.random.randn(10, 2)
rots = np.array([np.pi / 3.0])
R = np.array([[np.cos(rots), -np.sin(rots)],
[np.sin(rots), np.cos(rots)]]).reshape((2, 2))
rotfcoords = np.dot(fcoords, R.T)
akv = 200
wgh = 0.07
cs = 2.0
df1, df2, angast = 44722, 49349, 45.0 * (np.pi / 180.0)
dscale = 1.0
v1 = ctf.compute_ctf(fcoords, rots, akv, cs, wgh, df1,
df2, angast, dscale).reshape((-1,))
v2 = ctf.compute_ctf(rotfcoords, None, akv, cs, wgh, df1,
df2, angast, dscale).reshape((-1,))
# This being small confirms that using the rots parameter is equivalent to
# rotating the coordinates
print(np.abs(v1 - v2).max())
# N = 512
N = 128
psz = 5.6
rad = 0.25
fcoords = geometry.gencoords(N, 2, rad) / (N * psz)
ctf1_rot = ctf.compute_full_ctf(
rots, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None)
ctf2_full = ctf.compute_full_ctf(
None, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None)
print(ctf1_rot.shape)
fig, ax = plt.subplots(1, 2)
ax[0].imshow(ctf1_rot.reshape(N, N))
ax[1].imshow(ctf2_full.reshape(N, N))
ax[0].set_title('inplane-rotation: 0')
ax[1].set_title('inplane-rotation: 60 degree')
fig, ax = plt.subplots()
x = np.arange(int(-N/2.0), int(N/2.0))
ax.plot(x, ctf1_rot.reshape(N, N)[int(N/2), :], label='0 degree')
ax.plot(x, ctf2_full.reshape(N, N)[int(N/2), :], label='60 degree')
ax.set_xlabel('pixel')
ax.set_ylabel('CTF value')
ax.legend()
plt.show()
P_rot = coops.compute_inplanerot_matrix(rots, N, 'lanczos', 10, rad)
ctf2_rot = P_rot.dot(ctf2_full).reshape((-1,))
P_null = coops.compute_inplanerot_matrix(np.array([0]), N, 'linear', 2, rad)
ctf1_rot = P_null.dot(ctf1_rot).reshape((-1,))
print(fcoords.shape)
print(P_rot.shape)
print(ctf2_rot.shape)
roterr = ctf1_rot - ctf2_rot
relerr = np.abs(roterr) / np.maximum(np.abs(ctf1_rot), np.abs(ctf2_rot))
# This being small confirms that compute_inplane_rotmatrix and rots use
# the same rotation convention
print(relerr.max(), relerr.mean())
