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 trunc_to_full(trunc_samples, N, rad, beamstop_rad=None):
"""convert truncated samples to samples in full shape"""
num_samples = trunc_samples.shape[0]
TtoF = sincint.gentrunctofull(N=N, rad=rad, beamstop_rad=beamstop_rad)
full_samples = np.zeros((num_samples, N, N), dtype=trunc_samples.dtype)
for i, sample in enumerate(trunc_samples):
full_samples[i] = TtoF.dot(sample).reshape((N,N))
if num_samples == 1:
full_samples = full_samples.reshape((N, N))
return full_samples
def correlation_trunc_example(M):
N = M.shape[0]
rad = 1
proj = M.sum(axis=0)
FtoT = sincint.genfulltotrunc(N=N, rad=rad)
TtoF = sincint.gentrunctofull(N=N, rad=rad)
trunc = FtoT.dot(proj.flatten())
corr_trunc = correlation.calc_angular_correlation(trunc, N, 1)
corr_proj = TtoF.dot(corr_trunc).reshape(N, N)
fig, ax = plt.subplots(1,2)
ax[0].imshow(proj)
ax[1].imshow(corr_proj)
plt.show()
def premult_test(model, kernel='lanczos', kernsize=6):
if isinstance(model, str):
M = mrc.readMRC(model)
elif isinstance(model, np.ndarray):
M = model
shape = np.asarray(M.shape)
assert (shape - shape.mean()).sum() == 0
N = M.shape[0]
rad = 0.6
premult = cryoops.compute_premultiplier(N, kernel, kernsize)
TtoF = sincint.gentrunctofull(N=N, rad=rad)
premulter = premult.reshape((1, 1, -1)) \
* premult.reshape((1, -1, 1)) \
* premult.reshape((-1, 1, 1))
fM = density.real_to_fspace(M)
prefM = density.real_to_fspace(premulter * M)
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(fM.reshape((-1, )))
premult_trunc_slice = slop.dot(prefM.reshape((-1, )))
proj = density.fspace_to_real(TtoF.dot(trunc_slice).reshape(N, N))
premult_proj = density.fspace_to_real(
TtoF.dot(premult_trunc_slice).reshape(N, N))
fig, ax = plt.subplots(1, 3, figsize=(14.4, 4.8))
im_proj = ax[0].imshow(proj, origin='lower')
fig.colorbar(im_proj, ax=ax[0])
ax[0].set_title('no premulter')
im_pre = ax[1].imshow(premult_proj, origin='lower')
fig.colorbar(im_pre, ax=ax[1])
ax[1].set_title('with premulter')
im_diff = ax[2].imshow(proj - premult_proj, origin='lower')
fig.colorbar(im_diff, ax=ax[2])
ax[2].set_title('difference of two image')
fig.tight_layout()
plt.show()
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 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 project(model, euler_angles, rad=0.95, truncate=False):
if isinstance(model, str):
M = mrc.readMRC(model)
elif isinstance(model, np.ndarray):
M = model
N = M.shape[0]
kernel = 'lanczos'
ksize = 6
premult = cryoops.compute_premultiplier(N, kernel, ksize)
TtoF = sincint.gentrunctofull(N=N, rad=rad)
premulter = premult.reshape((1, 1, -1)) \
* premult.reshape((1, -1, 1)) \
* premult.reshape((-1, 1, 1))
# premulter = 1
fM = density.real_to_fspace(premulter * M)
euler_angles = euler_angles.reshape((-1, 3))
num_projs = euler_angles.shape[0]
if truncate:
projs = np.zeros((num_projs, TtoF.shape[1]), dtype=fM.dtype)
else:
projs = np.zeros((num_projs, N, N), dtype=M.dtype)
for i, ea in enumerate(euler_angles):
rot_matrix = geometry.rotmat3D_EA(*ea)[:, 0:2]
slop = cryoops.compute_projection_matrix([rot_matrix], N, kernel, ksize, rad, 'rots')
trunc_slice = slop.dot(fM.reshape((-1,)))
if truncate:
projs[i, :] = trunc_slice
else:
projs[i, :, :] = density.fspace_to_real(TtoF.dot(trunc_slice).reshape(N, N))
if num_projs == 1 and not truncate:
projs = projs.reshape((N, N))
return projs
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 genphantomdata(N_D, phantompath, ctfparfile):
mscope_params = {
'akv': 200,
'wgh': 0.07,
'cs': 2.0,
'psize': 2.8,
'bfactor': 500.0
}
N = 128
rad = 0.95
shift_sigma = 3.0
sigma_noise = 25.0
M_totalmass = 80000
kernel = 'lanczos'
ksize = 6
premult = cryoops.compute_premultiplier(N, kernel, ksize)
tic = time.time()
N_D = int(N_D)
N = int(N)
rad = float(rad)
psize = mscope_params['psize']
bfactor = mscope_params['bfactor']
shift_sigma = float(shift_sigma)
sigma_noise = float(sigma_noise)
M_totalmass = float(M_totalmass)
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 = n.sort(
n.random.random_integers(0,
srcctf_stack.get_num_ctfs() - 1, N_D))
M = mrc.readMRC(phantompath)
cryoem.window(M, 'circle')
M[M < 0] = 0
if M_totalmass is not None:
M *= M_totalmass / M.sum()
V = density.real_to_fspace(
premult.reshape((1, 1, -1)) * premult.reshape(
(1, -1, 1)) * premult.reshape((-1, 1, 1)) * M)
print "Generating data..."
sys.stdout.flush()
imgdata = n.empty((N_D, N, N), dtype=density.real_t)
pardata = {'R': [], 't': []}
prevctfI = None
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)),
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**2, ))
prevctfI = srcctfI
# Randomly generate the viewing direction/shift
pt = n.random.randn(3)
pt /= n.linalg.norm(pt)
psi = 2 * n.pi * n.random.rand()
EA = geom.genEA(pt)[0]
EA[2] = psi
shift = n.random.randn(2) * shift_sigma
R = geom.rotmat3D_EA(*EA)[:, 0:2]
slop = cryoops.compute_projection_matrix([R], N, kernel, ksize, rad,
'rots')
S = cryoops.compute_shift_phases(shift.reshape((1, 2)), N, rad)[0]
D = slop.dot(V.reshape((-1, )))
D *= S
imgdata[i] = density.fspace_to_real((C * TtoF.dot(D)).reshape(
(N, N))) + n.require(n.random.randn(N, N) * sigma_noise,
dtype=density.real_t)
genctf_stack.add_img(genctfI,
PHI=EA[0] * 180.0 / n.pi,
THETA=EA[1] * 180.0 / n.pi,
PSI=EA[2] * 180.0 / n.pi,
SHX=shift[0],
SHY=shift[1])
pardata['R'].append(R)
pardata['t'].append(shift)
print "\rDone in ", time.time() - tic, " seconds."
return imgdata, genctf_stack, pardata, mscope_params
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
# 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')
def genphantomdata(N_D, phantompath, ctfparfile):
# mscope_params = {'akv': 200, 'wgh': 0.07,
# 'cs': 2.0, 'psize': 2.8, 'bfactor': 500.0}
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
shift_sigma = 3.0
sigma_noise = 25.0
M_totalmass = 80000
kernel = 'lanczos'
ksize = 6
premult = cryoops.compute_premultiplier(N, kernel, ksize)
tic = time.time()
N_D = int(N_D)
N = int(N)
rad = float(rad)
psize = mscope_params['psize']
bfactor = mscope_params['bfactor']
shift_sigma = float(shift_sigma)
sigma_noise = float(sigma_noise)
M_totalmass = float(M_totalmass)
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()
V = density.real_to_fspace(
premult.reshape((1, 1, -1)) * premult.reshape((1, -1, 1)) * premult.reshape((-1, 1, 1)) * M)
print("Generating data...")
sys.stdout.flush()
imgdata = np.empty((N_D, N, N), dtype=density.real_t)
pardata = {'R': [], 't': []}
prevctfI = None
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)))
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**2,))
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
shift = np.random.randn(2) * shift_sigma
R = geometry.rotmat3D_EA(*EA)[:, 0:2]
slop = cryoops.compute_projection_matrix(
[R], N, kernel, ksize, rad, 'rots')
S = cryoops.compute_shift_phases(shift.reshape((1, 2)), N, rad)[0]
D = slop.dot(V.reshape((-1,)))
D *= S
imgdata[i] = density.fspace_to_real((C * TtoF.dot(D)).reshape((N, N))) + np.require(
np.random.randn(N, N) * sigma_noise, 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=shift[0], SHY=shift[1])
pardata['R'].append(R)
pardata['t'].append(shift)
print("\rDone in ", time.time() - tic, " seconds.")
return imgdata, genctf_stack, pardata, mscope_params
def plot_figures():
oversampling_factor = 6
psize = 3 * oversampling_factor
freq = 0.05
rad = 0.5 * 2.0 * psize
beamstop_freq = 0.005
beamstop_rad = beamstop_freq * 2.0 * psize
# M = mrc.readMRC('particle/EMD-6044-cropped.mrc')
# M[M < 0] = 0
# M_totalmass = 2000000
# if M_totalmass is not None:
# M *= M_totalmass / M.sum()
# fM = density.real_to_fspace_with_oversampling(M, oversampling_factor=oversampling_factor)
# fM = fM.real ** 2 + fM.imag ** 2
# mrc.writeMRC("particle/EMD-6044-cropped_fM_totalmass_%d_oversampling_%d.mrc" % (M_totalmass, oversampling_factor), fM, psz=psize)
# exit()
fM = mrc.readMRC(
'particle/EMD-6044-cropped_fM_totalmass_2000000_oversampling_6.mrc')
N = fM.shape[0]
TtoF = sincint.gentrunctofull(N=N, rad=rad, beamstop_rad=beamstop_rad)
theta = np.arange(0, 2 * np.pi, 2 * np.pi / 60)
degree_R, resolution_R = SK97Quadrature.compute_degree(N, 0.3, 1.0)
dirs, weights = SK97Quadrature.get_quad_points(degree_R, None)
Rs = np.vstack([
geometry.rotmat3D_dir(vec)[:, 0:2].reshape((1, 3, 2)) for vec in dirs
])
N_R = dirs.shape[0]
N_I = theta.shape[0]
N_T = TtoF.shape[1]
# generate slicing operators
dir_slice_interp = {
'projdirs': dirs,
'N': N,
'kern': 'lanczos',
'kernsize': 6,
'projdirtype': 'dirs',
'onlyRs': True,
'rad': rad,
'sym': None
} # 'zeropad': 0, 'dopremult': True
R_slice_interp = {
'projdirs': Rs,
'N': N,
'kern': 'lanczos',
'kernsize': 6,
'projdirtype': 'rots',
'onlyRs': True,
'rad': rad,
'sym': None
} # 'zeropad': 0, 'dopremult': True
inplane_interp = {
'thetas': theta,
'N': N,
'kern': 'lanczos',
'kernsize': 6,
'onlyRs': True,
'rad': rad
} # 'zeropad': 1, 'dopremult': True
inplane_interp['N_src'] = N # zp_N
slice_ops = cryoops.compute_projection_matrix(**dir_slice_interp)
inplane_ops = cryoops.compute_inplanerot_matrix(**inplane_interp)
# generate slices and inplane-rotated slices
slices_sampled = cryoem.getslices_interp(
fM, slice_ops, rad, beamstop_rad=beamstop_rad).reshape((N_R, N_T))
curr_img = TtoF.dot(slices_sampled[0]).reshape(N, N)
rotd_sampled = cryoem.getslices_interp(curr_img,
inplane_ops,
rad,
beamstop_rad=beamstop_rad).reshape(
(N_I, N_T))
## plot figures
fig, axes = plt.subplots(3, 4, figsize=(9.6, 7.2))
for i, ax in enumerate(axes.flatten()):
img = TtoF.dot(slices_sampled[i]).reshape(N, N)
ax.imshow(img, origin='lower')
fig.suptitle('slices_sampled')
fig, axes = plt.subplots(3, 4, figsize=(9.6, 7.2))
for i, ax in enumerate(axes.flatten()):
img = TtoF.dot(slices_sampled[i]).reshape(N, N)
ax.imshow(img, origin='lower')
fig.suptitle('rotd_sampled')
plt.show()
import pyximport
pyximport.install(setup_args={"include_dirs": np.get_include()},
reload_support=True)
import sincint
if __name__ == '__main__':
psize = 3.0 * 3
freq = 0.005
rad = freq * 2.0 * psize
beamstop_freq = 0.003
beamstop_rad = beamstop_freq * 2.0 * psize
fM = mrc.readMRC('../particle/EMD-6044-cropped-non-negative-256.mrc')
N = fM.shape[0]
TtoF = sincint.gentrunctofull(N=N, rad=rad, beamstop_rad=beamstop_rad)
theta = np.arange(0, 2 * np.pi, 2 * np.pi / 12)
degree_R, resolution_R = SK97Quadrature.compute_degree(N, rad, 1.0)
dirs, weights = SK97Quadrature.get_quad_points(degree_R, None)
Rs = np.vstack([
geometry.rotmat3D_dir(vec)[:, 0:2].reshape((1, 3, 2)) for vec in dirs
])
N_R = dirs.shape[0]
N_I = theta.shape[0]
N_T = TtoF.shape[1]
print(N_R, N_I, N_T)
# generate slicing operators
dir_slice_interp = {
# 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]))
S = cryoops.compute_shift_phases(np.asarray([100, -20]).reshape((1,2)), N, rad)[0]