def cov3d(starfile, data_folder, pixel_size, max_rows, max_resolution, cg_tol):
"""Estimate mean volume and covariance from a starfile."""
source = RelionSource(starfile,
data_folder=data_folder,
pixel_size=pixel_size,
max_rows=max_rows)
source.downsample(max_resolution)
source.cache()
source.whiten()
basis = FBBasis3D((max_resolution, max_resolution, max_resolution))
mean_estimator = MeanEstimator(source, basis, batch_size=8192)
mean_est = mean_estimator.estimate()
noise_estimator = WhiteNoiseEstimator(source, batchSize=500)
# Estimate the noise variance. This is needed for the covariance estimation step below.
noise_variance = noise_estimator.estimate()
logger.info(f"Noise Variance = {noise_variance}")
# Passing in a mean_kernel argument to the following constructor speeds up some calculations
covar_estimator = CovarianceEstimator(source,
basis,
mean_kernel=mean_estimator.kernel)
covar_estimator.estimate(mean_est, noise_variance, tol=cg_tol)
def setUpClass(cls):
cls.dtype = np.float32
cls.sim = Simulation(
n=1024,
unique_filters=[
RadialCTFFilter(defocus=d)
for d in np.linspace(1.5e4, 2.5e4, 7)
],
dtype=cls.dtype,
)
basis = FBBasis3D((8, 8, 8), dtype=cls.dtype)
cls.noise_variance = 0.0030762743633643615
cls.mean_estimator = MeanEstimator(cls.sim, basis)
cls.mean_est = Volume(
np.load(os.path.join(DATA_DIR,
"mean_8_8_8.npy")).astype(cls.dtype))
# Passing in a mean_kernel argument to the following constructor speeds up some calculations
cls.covar_estimator = CovarianceEstimator(
cls.sim,
basis,
mean_kernel=cls.mean_estimator.kernel,
preconditioner="none")
cls.covar_estimator_with_preconditioner = CovarianceEstimator(
cls.sim,
basis,
mean_kernel=cls.mean_estimator.kernel,
preconditioner="circulant",
)
def testEstimateResolutionError(self):
"""
Test mismatched resolutions yields a relevant error message.
"""
with raises(ValueError, match=r".*resolution.*"):
# This basis is intentionally the wrong resolution.
incorrect_basis = FBBasis3D((2 * self.resolution, ) * 3,
dtype=self.dtype)
_ = MeanEstimator(self.sim, incorrect_basis, preconditioner="none")
def setUp(self):
self.dtype = np.float32
sim = Simulation(
n=1024,
unique_filters=[
RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7)
],
dtype=self.dtype,
)
basis = FBBasis3D((8, 8, 8), dtype=self.dtype)
self.estimator = MeanEstimator(sim, basis, preconditioner="none")
self.estimator_with_preconditioner = MeanEstimator(
sim, basis, preconditioner="circulant"
)
def setUp(self):
self.dtype = np.float32
self.resolution = 8
self.n = 1024
# Generate a stack of images
self.sim = sim = Simulation(
n=self.n,
L=self.resolution,
unique_filters=[IdentityFilter()],
seed=0,
dtype=self.dtype,
# We'll use random angles
offsets=np.zeros((self.n, 2)), # No offsets
amplitudes=np.ones((self.n)), # Constant amplitudes
)
# Expose images as numpy array.
self.ims_np = sim.images(0, sim.n).asnumpy()
self.im = Image(self.ims_np)
# Vol estimation requires a 3D basis
self.basis = FBBasis3D((self.resolution,) * 3, dtype=self.dtype)
# Estimate noise in the ImageSource instance noise_estimator = AnisotropicNoiseEstimator(source) # Apply whitening to ImageSource source.whiten(noise_estimator.filter) # Display subset of the images images = source.images(0, 10) images.show() # %% # Estimate Mean Volume # -------------------- # We'll create a 3D Fourier Bessel Basis corresponding to volume resolution L. basis = FBBasis3D((L, L, L)) # Estimate mean Volume mean_estimator = MeanEstimator(source, basis, batch_size=8192) mean_est = mean_estimator.estimate() # %% # Visualize Volume # ---------------- # MeanEstimator.estimate() returns a Volume Instance, # which is wrapper on an ndarray representing a stack of one or more 3d volumes. # We will visualize the data via orthogonal projections along the three axes. vol = mean_est[0] # Visualize volume L = vol.shape[0]
# Specify parameters
num_vols = 2 # number of volumes
img_size = 8 # image size in square
num_imgs = 1024 # number of images
num_eigs = 16 # number of eigen-vectors to keep
# Create a simulation object with specified filters
sim = Simulation(
L=img_size,
n=num_imgs,
C=num_vols,
unique_filters=[RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7)],
)
# Specify the normal FB basis method for expending the 2D images
basis = FBBasis3D((img_size, img_size, img_size))
# Estimate the noise variance. This is needed for the covariance estimation step below.
noise_estimator = WhiteNoiseEstimator(sim, batchSize=500)
noise_variance = noise_estimator.estimate()
logger.info(f"Noise Variance = {noise_variance}")
# %%
# Estimate Mean Volume and Covariance
# -----------------------------------
#
# Estimate the mean. This uses conjugate gradient on the normal equations for
# the least-squares estimator of the mean volume. The mean volume is represented internally
# using the basis object, but the output is in the form of an
# L-by-L-by-L array.
def setUp(self):
self.dtype = np.float32
self.basis = FBBasis3D((8, 8, 8), dtype=self.dtype)