def plot_projs_with_slider(mrcs_files, log_scale=True, show_ac_image=False):
for mrcs in mrcs_files:
image_stack = mrc.readMRCimgs(mrcs, 0)
size = image_stack.shape
N = size[0]
mask = gen_dense_beamstop_mask(N, 2, 0.003, psize=9)
print('image size: {0}x{1}, number of images: {2}'.format(*size))
# plot projections
fig = plt.figure(figsize=(8, 8))
gs = GridSpec(
2,
2,
width_ratios=[1, 0.075],
height_ratios=[1, 0.075],
)
# original
ax = fig.add_subplot(gs[0, 0])
curr_img = image_stack[:, :, 0] * mask
if show_ac_image:
curr_ac_img = correlation.calc_full_ac(curr_img, 0.95) * mask
curr_img = curr_ac_img
if log_scale:
curr_img = log(curr_img)
im = ax.imshow(curr_img, origin='lower')
ticks = [0, int(N / 4.0), int(N / 2.0), int(N / 4.0 * 3), int(N - 1)]
ax.set_xticks(ticks)
ax.set_yticks(ticks)
ax.set_title('Slice Viewer (log scale: {}) for {}'.format(
log_scale, os.path.basename(mrcs)))
ax_divider = make_axes_locatable(ax)
cax = ax_divider.append_axes("right", size="7%", pad="2%")
cbar = fig.colorbar(im, cax=cax) # colorbar
# slider
ax_slider = fig.add_subplot(gs[1, 0])
idx_slider = Slider(ax_slider,
'index:',
0,
size[2] - 1,
valinit=0,
valfmt='%d')
def update(val):
idx = int(idx_slider.val)
curr_img = image_stack[:, :, idx] * mask
if show_ac_image:
curr_ac_img = correlation.calc_full_ac(curr_img, 0.95) * mask
curr_img = curr_ac_img
if log_scale:
curr_img = log(curr_img)
im.set_data(curr_img)
cbar.set_clim(vmin=curr_img.min(), vmax=curr_img.max())
cbar.draw_all()
fig.canvas.draw_idle()
idx_slider.on_changed(update)
plt.show()
def gen_mrcs_from_EAs(EAs, phantompath, dstpath):
ang_star = os.path.join(os.path.dirname(dstpath), 'EAs.star')
star.easy_writeSTAR_relion(ang_star, EAs=EAs)
projs_star, projs_mrcs = relion.project(phantompath, dstpath, ang=ang_star)
model_projs = mrc.readMRCimgs(projs_mrcs, idx=0)
return model_projs
def gen_mrcs_from_EAs(EAs, phantompath, dstpath):
ang_star = os.path.join(os.path.dirname(dstpath), 'EAs.star')
star.easy_writeSTAR_relion(ang_star, EAs=EAs)
projs_star, projs_mrcs = relion.project(phantompath, dstpath, ang=ang_star)
model_projs = mrc.readMRCimgs(projs_mrcs, idx=0)
return model_projs
def reconstruct(projs_path, nside, iteration_loops, **kwargs):
try:
WD = kwargs['WD']
except KeyError as error:
raise KeyError('lack working directory')
exp_samples = mrc.readMRCimgs(projs_path, idx=0)
input_shape = exp_samples.shape
print('size of input images: {0}x{1}, number of input images: {2}'.format(
*input_shape))
print('generating random phantom density')
N = input_shape[0]
M = cryoem.generate_phantom_density(N, 0.95 * N / 2.0, 5 * N / 128.0, 30)
output_mrc = os.path.join(WD, 'initial_random_model.mrc')
mrc.writeMRC(output_mrc, M)
print('Start projection matching')
total_iteration = iteration_loops
for it in range(total_iteration):
print('Iteration {0}'.format(it))
dir_suffix = 'it' + str(it).zfill(3)
iter_dir = os.path.join(WD, dir_suffix)
os.makedirs(iter_dir, exist_ok=True)
tic = time.time()
input_model = output_mrc
ori = projection_matching(input_model,
exp_samples,
nside=nside,
dir_suffix=dir_suffix,
**kwargs)
ori_star = os.path.join(WD, dir_suffix, 'orientations.star')
star.easy_writeSTAR_xmipp(ori_star, EAs=ori, imgs_path=projs_path)
output_mrc = os.path.join(WD, dir_suffix, dir_suffix + '_reco.mrc')
xmipp.reconstruct_fourier(ori_star, output_mrc, thr=4)
toc = time.time() - tic
print('Iteration {0} finished, spread time: {1:.4f} seconds'.format(
it, toc))
print('Total iteration is {0}, remian {1:.2f} minutes'.format(
total_iteration, (total_iteration - it) * toc / 60))
print(
'---------------------------------------------------------------------------------'
)
def plot_projs(mrcs_files, log_scale=True, plot_randomly=True):
for mrcs in mrcs_files:
image_stack = mrc.readMRCimgs(mrcs, 0)
size = image_stack.shape
N = size[0]
mask = gen_dense_beamstop_mask(N, 2, 0.003, psize=9)
print('image size: {0}x{1}, number of images: {2}'.format(*size))
print('Select indices randomly:', plot_randomly)
fig, axes = plt.subplots(3, 3, figsize=(12.9, 9.6))
for i, ax in enumerate(axes.flat):
row, col = unravel_index(i, (3, 3))
if plot_randomly:
num = randint(0, size[2])
else:
num = i
print('index:', num)
if log_scale:
img = log(maximum(image_stack[:, :, num], 1e-6)) * mask
else:
img = image_stack[:, :, num] * mask
im = ax.imshow(img, origin='lower') # cmap='Greys'
ticks = [
0,
int(N / 4.0),
int(N / 2.0),
int(N * 3.0 / 4.0),
int(N - 1)
]
if row == 2:
ax.set_xticks([])
else:
ax.set_xticks(ticks)
if col == 0:
ax.set_yticks([])
else:
ax.set_yticks(ticks)
fig.subplots_adjust(right=0.8)
cbarar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7])
fig.colorbar(im, cax=cbarar_ax)
fig.suptitle('{} before normalization'.format(mrcs))
# fig.tight_layout()
plt.show()
def reconstruct(projs_path, nside, iteration_loops, **kwargs):
try:
WD = kwargs['WD']
except KeyError as error:
raise KeyError('lack working directory')
exp_samples = mrc.readMRCimgs(projs_path, idx=0)
input_shape = exp_samples.shape
print('size of input images: {0}x{1}, number of input images: {2}'.format(
*input_shape))
print('generating random phantom density')
N = input_shape[0]
M = cryoem.generate_phantom_density(N, 0.95 * N / 2.0, 5 * N / 128.0, 30)
output_mrc = os.path.join(WD, 'initial_random_model.mrc')
mrc.writeMRC(output_mrc, M)
print('Start projection matching')
total_iteration = iteration_loops
for it in range(total_iteration):
print('Iteration {0}'.format(it))
dir_suffix = 'it' + str(it).zfill(3)
iter_dir = os.path.join(WD, dir_suffix)
os.makedirs(iter_dir, exist_ok=True)
tic = time.time()
input_model = output_mrc
ori = projection_matching(
input_model, exp_samples, nside=nside, dir_suffix=dir_suffix, **kwargs)
ori_star = os.path.join(WD, dir_suffix, 'orientations.star')
star.easy_writeSTAR_xmipp(ori_star, EAs=ori, imgs_path=projs_path)
output_mrc = os.path.join(WD, dir_suffix, dir_suffix + '_reco.mrc')
xmipp.reconstruct_fourier(ori_star, output_mrc, thr=4)
toc = time.time() - tic
print('Iteration {0} finished, spread time: {1:.4f} seconds'.format(it, toc))
print('Total iteration is {0}, remian {1:.2f} minutes'.format(
total_iteration, (total_iteration-it)*toc/60))
print('---------------------------------------------------------------------------------')
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input_imgs', type=str)
parser.add_argument('-r', '--random', type=bool, default=False)
args = parser.parse_args()
input_imgs = os.path.abspath(args.input_imgs)
plot_randomly = args.random
print(plot_randomly)
image_stack = mrc.readMRCimgs(input_imgs, 0)
size = image_stack.shape
print('image size: {0}x{1}, number of images: {2}'.format(*size))
plt.figure(1)
for i in range(9):
plt.subplot(331 + i)
if plot_randomly:
num = randint(0, size[2])
else:
num = i
print(num)
plt.imshow(image_stack[:, :, num])
plt.show()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input_imgs', type=str)
parser.add_argument('-r', '--random', type=bool, default=False)
args = parser.parse_args()
input_imgs = os.path.abspath(args.input_imgs)
plot_randomly = args.random
print(plot_randomly)
image_stack = mrc.readMRCimgs(input_imgs, 0)
size = image_stack.shape
print('image size: {0}x{1}, number of images: {2}'.format(*size))
plt.figure(1)
for i in range(9):
plt.subplot(331+i)
if plot_randomly:
num = randint(0, size[2])
else:
num = i
print(num)
plt.imshow(image_stack[:, :, num])
plt.show()
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)
N = M.shape[0]
kernel = 'lanczos'
ksize = 6
premult = cryoops.compute_premultiplier(zp_N, kernel, ksize)
V = density.real_to_fspace(premult.reshape((1, 1, -1)) * premult.reshape((1, -1, 1)) * premult.reshape((-1, 1, 1)) * ZP_M)
# V = density.real_to_fspace(ZP_M)
ZP_fM = V.real ** 2 + V.imag ** 2
fM = ZP_fM[zp_M_slicer]
# mask_3d_outlier = geometry.gen_dense_beamstop_mask(N, 3, 0.015, psize=2.8)
# fM *= mask_3d_outlier
# 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))