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 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 premult_vis(N=128):
kernel_function = ['lanczos', 'sinc']
for kernel in kernel_function:
fig, ax = plt.subplots()
kernel_size = np.arange(3, 11)
for ksize in kernel_size:
premult = cryoops.compute_premultiplier(N, kernel, ksize)
ax.plot(range(premult.shape[0]),
premult,
label='kernel size: {}'.format(ksize))
ax.set_title('kernel function: {}'.format(kernel))
ax.set_xlabel('N')
ax.legend(frameon=False)
plt.show()
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 set_transform(self, interp_params, caching=True):
self.caching = caching
self.transformed = {}
self.interp_params = interp_params
zeropad = interp_params.get('zeropad', 0)
kernel = interp_params.get('kern', 'lanczos')
kernsize = interp_params.get('kernsize', 4)
self.zeropad = int(zeropad * (self.get_num_pixels() / 2))
Nzp = 2 * self.zeropad + self.num_pixels
self.zpimg = np.zeros((Nzp, Nzp), dtype=density.real_t)
if interp_params.get('dopremult', True):
premult = cryoops.compute_premultiplier(Nzp, kernel, kernsize)
reshape = ((-1, 1), (1, -1))
self.premult = np.prod([premult.reshape(rs) for rs in reshape])
else:
self.premult = None
if self.premult is not None:
assert self.premult.shape[0] == Nzp
assert self.premult.shape[1] == Nzp
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 set_inplane_quad(self,rad):
# Get (and generate if needed) the quadrature scheme for inplane rotation
params = self.params
tic = time.time()
degree_I = params.get('quad_degree_I','auto')
usFactor_I = params.get('quad_undersample_I',params.get('quad_undersample',1.0))
kern_I = params.get('interp_kernel_I',params.get('interp_kernel',None))
kernsize_I = params.get('interp_kernel_size_I',params.get('interp_kernel_size',None))
zeropad_I = params.get('interp_zeropad_I',params.get('interp_zeropad',0))
dopremult_I = params.get('interp_premult_I',params.get('interp_premult',True))
maxAngle = quadrature.compute_max_angle(self.N,rad,usFactor_I)
if degree_I == 'auto':
degree_I = np.uint32(np.ceil(2.0 * np.pi / maxAngle))
resolution_I = max(0.5*quadrature.compute_max_angle(self.N,rad), 2.0*np.pi / degree_I)
inplane_params = { 'degree':degree_I }
interp_params_I = { 'N':self.N, 'kern':kern_I, 'kernsize':kernsize_I, 'rad':rad, 'zeropad':zeropad_I, 'dopremult':dopremult_I }
domain_change_I = self.inplane_params != inplane_params
interp_change_I = self.inplane_interp != interp_params_I
transform_change = self.inplane_interp is None or \
self.inplane_interp['kern'] != interp_params_I['kern'] or \
self.inplane_interp['kernsize'] != interp_params_I['kernsize'] or \
self.inplane_interp['zeropad'] != interp_params_I['zeropad']
if domain_change_I:
inplane_quad = {}
inplane_quad['resolution'] = resolution_I
inplane_quad['thetas'] = np.linspace(0, 2.0*np.pi, degree_I, endpoint=False)
inplane_quad['thetas'] += inplane_quad['thetas'][1]/2.0
inplane_quad['W'] = np.require((2.0*np.pi/float(degree_I))*np.ones((degree_I,)),dtype=np.float32)
self.quad_domain_I = quadrature.FixedCircleDomain(inplane_quad['thetas'],
inplane_quad['resolution'])
self.N_I = len(self.quad_domain_I)
self.sampler_I.setup(params, self.N_D, self.N_D_Train, self.quad_domain_I)
self.inplane_quad = inplane_quad
self.inplane_params = inplane_params
if domain_change_I or interp_change_I:
print(" Inplane Ops: %d, " % self.N_I); sys.stdout.flush()
if self.N_I < self.otf_thresh_I:
self.using_precomp_inplane = True
print("generated in", end=''); sys.stdout.flush()
self.inplane_ops = self.quad_domain_I.compute_operator(interp_params_I)
print(" {0} secs.".format(time.time() - tic))
else:
self.using_precomp_inplane = False
print("generating OTF.")
self.inplane_ops = None
self.inplane_interp = interp_params_I
if transform_change:
if dopremult_I:
premult = cryoops.compute_premultiplier(self.N + 2*int(interp_params_I['zeropad']*(self.N/2)),
interp_params_I['kern'],interp_params_I['kernsize'])
premult = premult.reshape((-1,1)) * premult.reshape((1,-1))
else:
premult = None
self.fspace_stack.set_transform(premult,interp_params_I['zeropad'])
def set_slice_quad(self,rad):
# Get (and generate if needed) the quadrature scheme for slicing
params = self.params
tic = time.time()
N = self.N
degree_R = params.get('quad_degree_R','auto')
quad_scheme_R = params.get('quad_type_R','sk97')
sym = get_symmetryop(params.get('symmetry',None)) if params.get('perfect_symmetry',True) else None
usFactor_R = params.get('quad_undersample_R',params.get('quad_undersample',1.0))
kern_R = params.get('interp_kernel_R',params.get('interp_kernel',None))
kernsize_R = params.get('interp_kernel_size_R',params.get('interp_kernel_size',None))
zeropad_R = params.get('interp_zeropad_R',params.get('interp_zeropad',0))
dopremult_R = params.get('interp_premult_R',params.get('interp_premult',True))
quad_R = quadrature.quad_schemes[('dir',quad_scheme_R)]
if degree_R == 'auto':
degree_R,resolution_R = quad_R.compute_degree(N,rad,usFactor_R)
resolution_R = max(0.5*quadrature.compute_max_angle(self.N,rad), resolution_R)
slice_params = { 'quad_type':quad_scheme_R, 'degree':degree_R,
'sym':sym }
interp_params_R = { 'N':self.N, 'kern':kern_R, 'kernsize':kernsize_R, 'rad':rad, 'zeropad':zeropad_R, 'dopremult':dopremult_R }
domain_change_R = slice_params != self.slice_params
interp_change_R = self.slice_interp != interp_params_R
transform_change = self.slice_interp is None or \
self.slice_interp['kern'] != interp_params_R['kern'] or \
self.slice_interp['kernsize'] != interp_params_R['kernsize'] or \
self.slice_interp['zeropad'] != interp_params_R['zeropad']
if domain_change_R:
slice_quad = {}
slice_quad['resolution'] = resolution_R
slice_quad['degree'] = degree_R
slice_quad['symop'] = sym
slice_quad['dir'],slice_quad['W'] = quad_R.get_quad_points(degree_R,slice_quad['symop'])
slice_quad['W'] = np.require(slice_quad['W'], dtype=np.float32)
self.quad_domain_R = quadrature.FixedSphereDomain(slice_quad['dir'],
slice_quad['resolution'],\
sym=sym)
self.N_R = len(self.quad_domain_R)
self.sampler_R.setup(params, self.N_D, self.N_D_Train, self.quad_domain_R)
self.slice_quad = slice_quad
self.slice_params = slice_params
if domain_change_R or interp_change_R:
symorder = 1 if self.slice_quad['symop'] is None else self.slice_quad['symop'].get_order()
print(" Slice Ops: %d, " % self.N_R); sys.stdout.flush()
if self.N_R*symorder < self.otf_thresh_R:
self.using_precomp_slicing = True
print("generated in", end=''); sys.stdout.flush()
self.slice_ops = self.quad_domain_R.compute_operator(interp_params_R)
print(" {0} secs.".format(time.time() - tic))
Gsz = (self.N_R,self.N_T)
self.G = np.empty(Gsz, dtype=self.G_datatype)
self.slices = np.empty(np.prod(Gsz), dtype=np.complex64)
else:
self.using_precomp_slicing = False
print("generating OTF.")
self.slice_ops = None
self.G = np.empty((self.N,self.N,self.N),dtype=self.G_datatype)
self.slices = None
self.slice_interp = interp_params_R
if transform_change:
if dopremult_R:
premult = cryoops.compute_premultiplier(self.N + 2*int(interp_params_R['zeropad']*(self.N/2)),
interp_params_R['kern'],interp_params_R['kernsize'])
premult = premult.reshape((-1,1,1)) * premult.reshape((1,-1,1)) * premult.reshape((1,1,-1))
else:
premult = None
self.slice_premult = premult
self.slice_zeropad = interp_params_R['zeropad']
assert interp_params_R['zeropad'] == 0,'Zero padding for slicing not yet implemented'
def set_proj_quad(self,rad):
# Get (and generate if needed) the quadrature scheme for slicing
params = self.params
tic = time.time()
N = self.N
quad_scheme_R = params.get('quad_type_R','sk97')
quad_R = quadrature.quad_schemes[('dir',quad_scheme_R)]
degree_R = params.get('quad_degree_R','auto')
degree_I = params.get('quad_degree_I','auto')
usFactor_R = params.get('quad_undersample_R',params.get('quad_undersample',1.0))
usFactor_I = params.get('quad_undersample_I',params.get('quad_undersample',1.0))
kern_R = params.get('interp_kernel_R',params.get('interp_kernel',None))
kernsize_R = params.get('interp_kernel_size_R',params.get('interp_kernel_size',None))
zeropad_R = params.get('interp_zeropad_R',params.get('interp_zeropad',0))
dopremult_R = params.get('interp_premult_R',params.get('interp_premult',True))
sym = get_symmetryop(params.get('symmetry',None)) if params.get('perfect_symmetry',True) else None
maxAngle = quadrature.compute_max_angle(self.N,rad,usFactor_I)
if degree_I == 'auto':
degree_I = np.uint32(np.ceil(2.0 * np.pi / maxAngle))
if degree_R == 'auto':
degree_R,resolution_R = quad_R.compute_degree(N,rad,usFactor_R)
resolution_R = max(0.5*quadrature.compute_max_angle(self.N,rad), resolution_R)
resolution_I = max(0.5*quadrature.compute_max_angle(self.N,rad), 2.0*np.pi / degree_I)
slice_params = { 'quad_type':quad_scheme_R, 'degree':degree_R,
'sym':sym }
inplane_params = { 'degree':degree_I }
proj_params = { 'quad_type_R':quad_scheme_R, 'degree_R':degree_R,
'sym':sym, 'degree_I':degree_I }
interp_params_RI = { 'N':self.N, 'kern':kern_R, 'kernsize':kernsize_R, 'rad':rad, 'zeropad':zeropad_R, 'dopremult':dopremult_R }
interp_change_RI = self.proj_interp != interp_params_RI
transform_change = self.slice_interp is None or \
self.slice_interp['kern'] != interp_params_RI['kern'] or \
self.slice_interp['kernsize'] != interp_params_RI['kernsize'] or \
self.slice_interp['zeropad'] != interp_params_RI['zeropad']
domain_change_R = self.slice_params != slice_params
domain_change_I = self.inplane_params != inplane_params
domain_change_RI = self.proj_params != proj_params
if domain_change_RI:
proj_quad = {}
proj_quad['resolution'] = max(resolution_R,resolution_I)
proj_quad['degree_R'] = degree_R
proj_quad['degree_I'] = degree_I
proj_quad['symop'] = sym
proj_quad['dir'],proj_quad['W_R'] = quad_R.get_quad_points(degree_R,proj_quad['symop'])
proj_quad['W_R'] = np.require(proj_quad['W_R'], dtype=np.float32)
proj_quad['thetas'] = np.linspace(0, 2.0*np.pi, degree_I, endpoint=False)
proj_quad['thetas'] += proj_quad['thetas'][1]/2.0
proj_quad['W_I'] = np.require((2.0*np.pi/float(degree_I))*np.ones((degree_I,)),dtype=np.float32)
self.quad_domain_RI = quadrature.FixedSO3Domain( proj_quad['dir'],
-proj_quad['thetas'],
proj_quad['resolution'],
sym=sym)
self.N_RI = len(self.quad_domain_RI)
self.proj_quad = proj_quad
self.proj_params = proj_params
if domain_change_R:
self.quad_domain_R = quadrature.FixedSphereDomain(proj_quad['dir'],
proj_quad['resolution'],
sym=sym)
self.N_R = len(self.quad_domain_R)
self.sampler_R.setup(params, self.N_D, self.N_D_Train, self.quad_domain_R)
self.slice_params = slice_params
if domain_change_I:
self.quad_domain_I = quadrature.FixedCircleDomain(proj_quad['thetas'],
proj_quad['resolution'])
self.N_I = len(self.quad_domain_I)
self.sampler_I.setup(params, self.N_D, self.N_D_Train, self.quad_domain_I)
self.inplane_params = inplane_params
if domain_change_RI or interp_change_RI:
symorder = 1 if self.proj_quad['symop'] is None else self.proj_quad['symop'].get_order()
print(" Projection Ops: %d (%d slice, %d inplane), " % (self.N_RI, self.N_R, self.N_I)); sys.stdout.flush()
if self.N_RI*symorder < self.otf_thresh_RI:
self.using_precomp_slicing = True
print("generated in", end=''); sys.stdout.flush()
self.slice_ops = self.quad_domain_RI.compute_operator(interp_params_RI)
print(" {0} secs.".format(time.time() - tic))
Gsz = (self.N_RI,self.N_T)
self.G = np.empty(Gsz, dtype=self.G_datatype)
self.slices = np.empty(np.prod(Gsz), dtype=np.complex64)
else:
self.using_precomp_slicing = False
print("generating OTF.")
self.slice_ops = None
self.G = np.empty((N,N,N),dtype=self.G_datatype)
self.slices = None
self.using_precomp_inplane = False
self.inplane_ops = None
self.proj_interp = interp_params_RI
if transform_change:
if dopremult_R:
premult = cryoops.compute_premultiplier(self.N + 2*int(interp_params_RI['zeropad']*(self.N/2)),
interp_params_RI['kern'],interp_params_RI['kernsize'])
premult = premult.reshape((-1,1,1)) * premult.reshape((1,-1,1)) * premult.reshape((1,1,-1))
else:
premult = None
self.slice_premult = premult
self.slice_zeropad = interp_params_RI['zeropad']
assert interp_params_RI['zeropad'] == 0,'Zero padding for slicing not yet implemented'
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
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')
def sagd_init(data_dir, model_file, use_angular_correlation=False):
data_params = {
'dataset_name': "1AON",
'inpath': os.path.join(data_dir, 'imgdata.mrc'),
'ctfpath': os.path.join(data_dir, 'defocus.txt'),
'microscope_params': {
'akv': 200,
'wgh': 0.07,
'cs': 2.0
},
'resolution': 2.8,
'sigma': 'noise_std',
'sigma_out': 'data_std',
'minisize': 20,
'test_imgs': 20,
'partition': 0,
'num_partitions': 0,
'random_seed': 1,
# 'symmetry': 'C7'
}
# Setup dataset
print("Loading dataset %s" % data_dir)
cryodata, _ = dataset_loading_test(data_params)
# mleDC, _, mleDC_est_std = cryodata.get_dc_estimate()
# modelscale = (np.abs(mleDC) + 2*mleDC_est_std)/cryodata.N
modelscale = 1.0
if model_file is not None:
print("Loading density map %s" % model_file)
M = readMRC(model_file)
else:
print("Generating random initial density map ...")
M = cryoem.generate_phantom_density(cryodata.N, 0.95 * cryodata.N / 2.0, \
5 * cryodata.N / 128.0, 30, seed=0)
M *= modelscale / M.sum()
slice_interp = {
'kern': 'lanczos',
'kernsize': 4,
'zeropad': 0,
'dopremult': True
}
# fM = SimpleKernel.get_fft(M, slice_interp)
M_totalmass = 5000
M *= M_totalmass / M.sum()
N = M.shape[0]
kernel = 'lanczos'
ksize = 6
premult = cryoops.compute_premultiplier(N, kernel, ksize)
V = density.real_to_fspace(
premult.reshape((1, 1, -1)) * premult.reshape(
(1, -1, 1)) * premult.reshape((-1, 1, 1)) * M)
M = V.real**2 + V.imag**2
freqs_3d = geometry.gencoords_base(N, 3) / (N * data_params['resolution'])
freq_radius_3d = np.sqrt((freqs_3d**2).sum(axis=1))
mask_3d_outlier = np.require(np.float_(freq_radius_3d > 0.015).reshape(
(N, N, N)),
dtype=density.real_t)
fM = M * mask_3d_outlier
cparams = {
'use_angular_correlation': use_angular_correlation,
'likelihood': 'UnknownRSLikelihood()',
'kernel': 'multicpu',
'prior_name': "'Null'",
'sparsity_lambda': 0.9,
'prior': 'NullPrior()',
# 'prior_name': "'CAR'",
# 'prior': 'CARPrior()',
# 'car_type': 'gauss0.5',
# 'car_tau': 75.0,
'iteration': 0,
'pixel_size': cryodata.pixel_size,
'max_frequency': 0.02,
'num_batches': cryodata.N_batches,
'interp_kernel_R': 'lanczos',
'interp_kernel_size_R': 4,
'interp_zeropad_R': 0.0,
'interp_premult_R': True,
'interp_kernel_I': 'lanczos',
'interp_kernel_size_I': 8,
'interp_zeropad_I': 0.0, # 1.0,
'interp_premult_I': True,
'sigma': cryodata.noise_var,
'modelscale': modelscale,
# 'symmetry': 'C7'
}
is_params = {
# importance sampling
# Ignore the first 50 iterations entirely
'is_prior_prob': max(0.05,
2**(-0.005 * max(0, cparams['iteration'] - 50))),
'is_temperature': max(1.0,
2**(750.0 / max(1, cparams['iteration'] - 50))),
'is_ess_scale': 10,
'is_fisher_chirality_flip': cparams['iteration'] < 2500,
'is_on_R': True,
'is_global_prob_R': 0.9,
'is_on_I': True,
'is_global_prob_I': 1e-10,
'is_on_S': True,
'is_global_prob_S': 0.9,
'is_gaussian_sigmascale_S': 0.67,
}
cparams.update(is_params)
return cryodata, (M, fM), cparams
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
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]))
S = cryoops.compute_shift_phases(np.asarray([100, -20]).reshape((1,2)), N, rad)[0]
