def precompute_projections(self, fM):
if self.using_precomp_slicing:
# precompute all possible slices
getslices(fM.reshape((-1, )), self.slice_ops, res=self.slices)
self.precomp_slices = self.slices.reshape((-1, self.N_T))
else:
# do on-the-fly slicing
self.precomp_slices = None
def precompute_projections(self, fM):
if self.using_precomp_slicing:
# precompute all possible slices
getslices(fM.reshape((-1,)), self.slice_ops, res=self.slices)
self.precomp_slices = self.slices.reshape((-1, self.N_T))
else:
# do on-the-fly slicing
self.precomp_slices = None
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 prep_operators(self,fM,idx, slicing = True, res=None):
Idx = self.minibatch['img_idxs'][idx]
CIdx = self.minibatch['ctf_idxs'][idx]
cCTF = self.cryodata.ctfstack.get_ctf(CIdx)
Img = self.fspace_stack.get_image(Idx)
factoredRI = self.factoredRI
if not factoredRI:
W_R = self.proj_quad['W_R']
W_I = self.proj_quad['W_I']
else:
W_R = self.slice_quad['W']
W_I = self.inplane_quad['W']
W_S = self.shift_quad['W']
envelope = self.envelope
tic = time.time()
samples_R, sampleweights_R = self.sampler_R.sample(Idx)
if samples_R is None:
N_R_sampled = self.N_R
W_R_sampled = W_R
else:
N_R_sampled = len(samples_R)
W_R_sampled = np.require(W_R[samples_R] * sampleweights_R, dtype = W_R.dtype)
sampleinfo_R = N_R_sampled, samples_R, sampleweights_R
if res is not None:
res['kern_timing']['prep_sample_R'][idx] = time.time() - tic
tic = time.time()
samples_I, sampleweights_I = self.sampler_I.sample(Idx)
if samples_I is None:
N_I_sampled = self.N_I
W_I_sampled = W_I
else:
N_I_sampled = len(samples_I)
W_I_sampled = np.require(W_I[samples_I] * sampleweights_I, dtype = W_I.dtype)
sampleinfo_I = N_I_sampled, samples_I, sampleweights_I
if res is not None:
res['kern_timing']['prep_sample_I'][idx] = time.time() - tic
tic = time.time()
samples_S, sampleweights_S = self.sampler_S.sample(Idx)
if samples_S is None:
N_S_sampled = self.N_S
S_sampled = self.shift_ops
W_S_sampled = W_S
else:
N_S_sampled = len(samples_S)
S_sampled = self.shift_ops[samples_S]
W_S_sampled = np.require( W_S[samples_S] * sampleweights_S , dtype = W_S.dtype)
sampleinfo_S = N_S_sampled, samples_S, sampleweights_S
if res is not None:
res['kern_timing']['prep_sample_S'][idx] = time.time() - tic
if slicing:
if not factoredRI:
if samples_R is None and samples_I is None:
full_samples = None
else:
full_samples = []
it_R = xrange(self.N_R) if samples_R is None else samples_R
it_I = xrange(self.N_I) if samples_I is None else samples_I
for r in it_R:
full_samples += [(r,i) for i in it_I]
full_samples = np.array(full_samples)
samples_R = self.N_I*full_samples[:,0] + full_samples[:,1]
samples_I = np.array(0)
W_R_sampled = (W_R_sampled.reshape((N_R_sampled,1)) * W_I_sampled.reshape((1,N_I_sampled))).reshape((N_R_sampled*N_I_sampled,))
W_I_sampled = np.array([1.0], dtype = W_I.dtype)
N_R_sampled = N_R_sampled*N_I_sampled
N_I_sampled = 1
if self.using_precomp_slicing:
slice_ops = self.slice_ops
if samples_R is None:
slices_sampled = self.precomp_slices
else:
slices_sampled = self.precomp_slices[samples_R]
else:
slice_ops = self.quad_domain_RI.compute_operator(self.interp_params,samples_R)
slices_sampled = getslices(fM.reshape((-1,)), slice_ops).reshape((N_R_sampled,self.N_T))
rotd_sampled = Img[self.truncmask.reshape(Img.shape)].reshape((N_I_sampled,self.N_T))
rotc_sampled = cCTF.compute(self.trunc_freq).reshape((1,self.N_T))
else:
tic = time.time()
if self.using_precomp_slicing:
slice_ops = self.slice_ops
if samples_R is None:
slices_sampled = self.precomp_slices
else:
slices_sampled = self.precomp_slices[samples_R]
else:
slice_ops = self.quad_domain_R.compute_operator(self.slice_interp,samples_R)
slices_sampled = getslices(fM.reshape((-1,)), slice_ops).reshape((N_R_sampled,self.N_T))
if res is not None:
res['kern_timing']['prep_slice'][idx] = time.time() - tic
tic = time.time()
if samples_I is None:
rotc_sampled = cCTF.compute(self.trunc_freq,self.quad_domain_I.theta).T
else:
rotc_sampled = cCTF.compute(self.trunc_freq,self.quad_domain_I.theta[samples_I]).T
if res is not None:
res['kern_timing']['prep_rot_ctf'][idx] = time.time() - tic
# Generate the rotated versions of the current image
if self.using_precomp_inplane:
if samples_I is None:
rotd_sampled = getslices(Img,self.inplane_ops).reshape((N_I_sampled,self.N_T))
else:
rotd_sampled = getslices(Img,self.inplane_ops).reshape((self.N_I,self.N_T))[samples_I]
else:
inplane_ops = self.quad_domain_I.compute_operator(self.inplane_interp,samples_I)
rotd_sampled = getslices(Img,inplane_ops).reshape((N_I_sampled,self.N_T))
if res is not None:
res['kern_timing']['prep_rot_img'][idx] = time.time() - tic
else:
slice_ops = None
slices_sampled = None
rotc_sampled = None
rotd_sampled = None
return slice_ops, envelope, \
W_R_sampled, sampleinfo_R, slices_sampled, samples_R, \
W_I_sampled, sampleinfo_I, rotd_sampled, rotc_sampled, \
W_S_sampled, sampleinfo_S, S_sampled
# fM = mrc.readMRC('particle/1AON_fM_totalmass_5000.mrc') * mask_3d_outlier
imgdata = mrc.readMRCimgs('data/1AON_xfel_5000_totalmass_05000/imgdata.mrc', 420, 1)
curr_img = imgdata[:, :, 0]
zp_img = np.zeros((zp_N,)*2, dtype=density.real_t)
zp_img[zp_M_slicer[0:2]] = curr_img
# curr_img = zp_img
slice_interp = {'projdirs': dirs, 'N': N, 'kern': 'lanczos', 'kernsize': 4, 'rad': rad, 'sym': None} # 'zeropad': 0, 'dopremult': True
inplane_interp = {'thetas': theta, 'N': N, 'kern': 'lanczos', 'kernsize': 4, 'rad': rad, 'onlyRs': False} # 'zeropad': 1, 'dopremult': True
inplane_interp['N_src'] = N # zp_N
slice_ops = cryoops.compute_projection_matrix(**slice_interp)
inplane_ops = cryoops.compute_inplanerot_matrix(**inplane_interp)
slices_sampled = cryoem.getslices(fM, slice_ops).reshape((N_R, N_T))
rotd_sampled = cryoem.getslices(curr_img, inplane_ops).reshape((N_I, N_T))
slices_sampled += 1.0
print(slices_sampled[0])
print("slices_sampled max():", slices_sampled.max(axis=1))
print("slices_sampled min():", slices_sampled.min(axis=1))
print("slices_sampled < 1.0:", (slices_sampled<1.0).sum(axis=1))
print((rotd_sampled < 1.0).sum(axis=1))
# fig, ax = plt.subplots()
# im = ax.imshow( ( curr_img ) )
# fig.colorbar(im, ax=ax)
