def gaussian_blobs(K, alpha):
Q = np.zeros(shape=(3, 3, K)).astype(self.dtype)
D = np.zeros(shape=(3, 3, K)).astype(self.dtype)
mu = np.zeros(shape=(3, K)).astype(self.dtype)
for k in range(K):
V = randn(3, 3).astype(self.dtype) / np.sqrt(3)
Q[:, :, k] = qr(V)[0]
D[:, :, k] = alpha**2 / 16 * np.diag(np.sum(abs(V)**2, axis=0))
mu[:, k] = 0.5 * randn(3) / np.sqrt(3)
return Q, D, mu
def qrand(nrot, seed=0):
"""
Generate a set of quaternions from random normal distribution.
Each quaternions is a four-elements column vector. Returns a matrix of
size 4xn.
The 3-sphere S^3 in R^4 is a double cover of the rotation group SO(3), SO(3) = RP^3.
We identify unit norm quaternions a^2+b^2+c^2+d^2=1 with group elements.
The antipodal points (-a,-b,-c,-d) and (a,b,c,d) are identified as the same group elements,
so we take a>=0.
:param nrot: The number of quaternions for rotations.
:param seed: The random seed.
:return: An array consists of 4 dimensions quaternions
"""
q = randn(4, nrot, seed=seed)
l2_norm = np.sqrt(q[0, :]**2 + q[1, :]**2 + q[2, :]**2 + q[3, :]**2)
for i in range(4):
q[i, :] = q[i, :] / l2_norm
for k in range(nrot):
if q[0, k] < 0:
q[:, k] = -q[:, k]
return q
def _noise_images(self,
start=0,
num=None,
noise_seed=0,
noise_filter=None):
# Generate noisy images in interval [start, start+num-1] (a total of 'num' images)
end = self.n
if num is not None:
end = min(start + num, self.n)
all_idx = np.arange(start, end)
if noise_filter is None:
noise_filter = ScalarFilter(value=1, power=0.5)
im = np.zeros((self.L, self.L, len(all_idx)), dtype=self.dtype)
for idx in all_idx:
random_seed = noise_seed + 191 * (idx + 1)
im_s = randn(2 * self.L, 2 * self.L, seed=random_seed)
im_s = im_filter(im_s, noise_filter)
im_s = im_s[:self.L, :self.L]
im[:, :, idx - start] = im_s
return im
def __init__(self,
L=8,
n=1024,
vols=None,
states=None,
filters=None,
offsets=None,
amplitudes=None,
dtype='single',
C=2,
angles=None,
seed=0,
memory=None,
noise_filter=None):
"""
A Cryo-EM simulation
Other than the base class attributes, it has:
:param C: The number of distinct volumes
:param angles: A 3-by-n array of rotation angles
"""
super().__init__(L=L, n=n, dtype=dtype, memory=memory)
if offsets is None:
offsets = L / 16 * randn(2, n, seed=seed).T
if amplitudes is None:
min_, max_ = 2. / 3, 3. / 2
amplitudes = min_ + rand(n, seed=seed) * (max_ - min_)
states = states or randi(C, n, seed=seed)
angles = angles or uniform_random_angles(n, seed=seed)
self.states = states
if filters is not None:
self.filters = np.take(filters,
randi(len(filters), n, seed=seed) - 1)
else:
self.filters = None
self.offsets = offsets
self.amplitudes = amplitudes
self.angles = angles
self.C = C
if vols is None:
self.vols = self._gaussian_blob_vols(L=self.L, C=self.C, seed=seed)
else:
self.vols = vols
self.seed = seed
self.noise_adder = None
if noise_filter is not None and not isinstance(noise_filter,
ZeroFilter):
logger.info(f'Appending a NoiseAdder to generation pipeline')
self.noise_adder = NoiseAdder(seed=self.seed,
noise_filter=noise_filter)
def _forward(self, im, indices):
im = im.copy()
for i, idx in enumerate(indices):
# Note: The following random seed behavior is directly taken from MATLAB Cov3D code.
random_seed = self.seed + 191 * (idx + 1)
im_s = randn(2 * im.res, 2 * im.res, seed=random_seed)
im_s = Image(im_s).filter(self.noise_filter)[:, :, 0]
im[:, :, i] += im_s[:im.res, :im.res]
return im
def setUp(self):
L = 8
n = 32
C = 1
SNR = 1
pixel_size = 5
voltage = 200
defocus_min = 1.5e4
defocus_max = 2.5e4
defocus_ct = 7
Cs = 2.0
alpha = 0.1
filters = [
RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
for d in np.linspace(defocus_min, defocus_max, defocus_ct)
]
# Since FFBBasis2D doesn't yet implement dtype, we'll set this to double to match its built in types.
sim = Simulation(n=n, C=C, filters=filters, dtype='double')
vols = np.load(os.path.join(DATA_DIR, 'clean70SRibosome_vol.npy'))
vols = vols[..., np.newaxis]
vols = downsample(vols, (L * np.ones(3, dtype=int)))
sim.vols = vols
self.basis = FFBBasis2D((L, L))
# use new methods to generate random rotations and clean images
sim.rots = qrand_rots(n, seed=0)
self.imgs_clean = vol2img(vols[..., 0], sim.rots)
self.h_idx = np.array([filters.index(f) for f in sim.filters])
self.filters = filters
self.h_ctf_fb = [filt.fb_mat(self.basis) for filt in self.filters]
self.imgs_ctf_clean = sim.eval_filters(self.imgs_clean)
sim.cache(self.imgs_ctf_clean)
power_clean = anorm(self.imgs_ctf_clean)**2 / np.size(
self.imgs_ctf_clean)
self.noise_var = power_clean / SNR
self.imgs_ctf_noise = self.imgs_ctf_clean + np.sqrt(
self.noise_var) * randn(L, L, n, seed=0)
self.cov2d = RotCov2D(self.basis)
self.coeff_clean = self.basis.evaluate_t(self.imgs_clean)
self.coeff = self.basis.evaluate_t(self.imgs_ctf_noise)
def __init__(self,
L=8,
n=1024,
states=None,
filters=None,
offsets=None,
amplitudes=None,
dtype='single',
C=2,
rots=None):
"""
A Cryo-EM simulation
Other than the base class attributes, it has:
:param C: The no. of distinct volumes
:param rots: A 3-by-3-by-n array of rotation matrices corresponding to viewing directions
"""
offsets = offsets or L / 16 * randn(2, n, seed=0)
if amplitudes is None:
min_, max_ = 2. / 3, 3. / 2
amplitudes = min_ + rand(n, seed=0) * (max_ - min_)
states = states or randi(C, n, seed=0)
rots = rots or self._uniform_random_rotations(n, seed=0)
super().__init__(L=L,
n=n,
states=states,
filters=filters,
offsets=offsets,
amplitudes=amplitudes,
rots=rots,
dtype=dtype)
self.C = C
self.vols = self._gaussian_blob_vols(L=self.L, C=self.C, seed=0)
# Evaluate CTF in the 8X8 FB basis
h_ctf_fb = [filt.fb_mat(ffbbasis) for filt in filters]
# Apply the CTF to the clean images.
logger.info('Apply CTF filters to clean images.')
imgs_ctf_clean = Image(sim.eval_filters(imgs_clean))
sim.cache(imgs_ctf_clean)
# imgs_ctf_clean is an Image object. Convert to numpy array for subsequent statements
imgs_ctf_clean = imgs_ctf_clean.asnumpy()
# Apply the noise at the desired singal-noise ratio to the filtered clean images
logger.info('Apply noise filters to clean images.')
power_clean = anorm(imgs_ctf_clean)**2 / np.size(imgs_ctf_clean)
noise_var = power_clean / sn_ratio
imgs_noise = imgs_ctf_clean + np.sqrt(noise_var) * randn(
img_size, img_size, num_imgs, seed=0)
# Expand the images, both clean and noisy, in the Fourier-Bessel basis. This
# can be done exactly (that is, up to numerical precision) using the
# `basis.expand` function, but for our purposes, an approximation will do.
# Since the basis is close to orthonormal, we may approximate the exact
# expansion by applying the adjoint of the evaluation mapping using
# `basis.evaluate_t`.
logger.info('Get coefficients of clean and noisy images in FFB basis.')
coeff_clean = ffbbasis.evaluate_t(imgs_clean)
coeff_noise = ffbbasis.evaluate_t(imgs_noise)
# Create the Cov2D object and calculate mean and covariance for clean images without CTF.
# Given the clean Fourier-Bessel coefficients, we can estimate the symmetric
# mean and covariance. Note that these are not the same as the sample mean and
# covariance, since these functions use the rotational and reflectional