def set_data(self, model, cparams):
assert isinstance(model, np.ndarray), "Unexpected input numpy array."
self.model = model
self.slices_sampled = None
self.rotd_sampled = None
self.rad = None
cparams['iteration'] = 0
max_freq = cparams['max_frequency']
rad_cutoff = cparams.get('rad_cutoff', 1.0)
fake_oversampling_factor = 1.0
rad = min(rad_cutoff, max_freq * 2 * self.psize)
self.beamstop_rad = cparams.get('beamstop_freq', 0.003) * 2 * self.psize
self.xy, self.trunc_xy, self.truncmask = geometry.gencoords_centermask(self.N, 2, rad, self.beamstop_rad, True)
self.trunc_freq = np.require(self.trunc_xy / (self.N * self.psize), dtype=np.float32)
self.N_T = self.trunc_xy.shape[0]
# set CTF and envelope
if self.cryodata.use_ctf:
radius_freqs = np.sqrt(np.sum(self.trunc_xy**2,axis=1))/(self.psize*self.N)
self.envelope = ctf.envelope_function(radius_freqs, cparams.get('learn_like_envelope_bfactor', None))
else:
self.envelope = np.ones(self.N_T, dtype=np.float32)
# print("Iteration {0}: freq = {3}, rad = {1:.4f}, beamstop_rad={4:.4f}, N_T = {2}".format(cparams['iteration'], rad, self.N_T, max_freq, self.beamstop_rad))
self.set_quad(rad)
# Setup inlier model
self.inlier_sigma2 = 1.0 # cparams['sigma']**2
base_sigma2 = 1.0 # self.cryodata.noise_var
self.inlier_sigma2_trunc = self.inlier_sigma2
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, n.ndarray):
sigma2 = sigma2.reshape((N, N))
if bfactor is not None:
coords = gencoords(N, 2).reshape((N**2, 2))
freqs = n.sqrt(n.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:
n.clip(exp_env, env_lb, env_ub, out=exp_env)
return exp_env
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 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 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
def get_env_func(N,resolution,bfactor = None):
freq_radius = np.linspace(0,N/2,N/2+1)/(N*resolution)
env = ctf.envelope_function(freq_radius,bfactor)
return freq_radius, env
def get_env_func(N,resolution,bfactor = None):
freq_radius = n.linspace(0,N/2,N/2+1)/(N*resolution)
env = ctf.envelope_function(freq_radius,bfactor)
return freq_radius, env
