def downsample(self, ds_res): """ Downsample Image to a specific resolution. This method returns a new Image. :param ds_res: int - new resolution, should be <= the current resolution of this Image :return: The downsampled Image object. """ grid = grid_2d(self.res) grid_ds = grid_2d(ds_res) im_ds = np.zeros((self.n_images, ds_res, ds_res), dtype=self.dtype) # x, y values corresponding to 'grid'. This is what scipy interpolator needs to function. res_by_2 = self.res / 2 x = y = np.ceil(np.arange(-res_by_2, res_by_2)) / res_by_2 mask = (np.abs(grid["x"]) < ds_res / self.res) & (np.abs(grid["y"]) < ds_res / self.res) im_shifted = fft.centered_ifft2( fft.centered_fft2(xp.asarray(self.data)) * xp.asarray(mask)) im = np.real(xp.asnumpy(im_shifted)) for s in range(im_ds.shape[0]): interpolator = RegularGridInterpolator((x, y), im[s], bounds_error=False, fill_value=0) im_ds[s] = interpolator(np.dstack([grid_ds["x"], grid_ds["y"]])) return Image(im_ds)
def im_downsample(im, L_ds): """ Blur and downsample image :param im: Set of images to be downsampled in the form of an array L-by-L-by-K, where K is the number of images. :param L_ds: The desired resolution of the downsampled images. Must be smaller than L. :return: An array of the form L_ds-by-L_ds-by-K consisting of the blurred and downsampled images. """ N = im.shape[0] grid = grid_2d(N) grid_ds = grid_2d(L_ds) im_ds = np.zeros((L_ds, L_ds, im.shape[2])).astype(im.dtype) # x, y values corresponding to 'grid'. This is what scipy interpolator needs to function. x = y = np.ceil(np.arange(-N/2, N/2)) / (N/2) mask = (np.abs(grid['x']) < L_ds/N) & (np.abs(grid['y']) < L_ds/N) im = np.real(centered_ifft2(centered_fft2(im) * np.expand_dims(mask, 2))) for s in range(im_ds.shape[-1]): interpolator = RegularGridInterpolator( (x, y), im[:, :, s], bounds_error=False, fill_value=0 ) im_ds[:, :, s] = interpolator(np.dstack([grid_ds['x'], grid_ds['y']])) return im_ds
def downsample(self, ds_res): """ Downsample Image to a specific resolution. This method returns a new Image. :param ds_res: int - new resolution, should be <= the current resolution of this Image :return: The downsampled Image object. """ grid = grid_2d(self.res) grid_ds = grid_2d(ds_res) im_ds = np.zeros((ds_res, ds_res, self.n_images)).astype(self.dtype) # x, y values corresponding to 'grid'. This is what scipy interpolator needs to function. res_by_2 = self.res / 2 x = y = np.ceil(np.arange(-res_by_2, res_by_2)) / res_by_2 mask = (np.abs(grid['x']) < ds_res / self.res) & (np.abs(grid['y']) < ds_res / self.res) im = np.real( centered_ifft2(centered_fft2(self.data) * np.expand_dims(mask, 2))) for s in range(im_ds.shape[-1]): interpolator = RegularGridInterpolator((x, y), im[:, :, s], bounds_error=False, fill_value=0) im_ds[:, :, s] = interpolator(np.dstack([grid_ds['x'], grid_ds['y']])) return Image(im_ds)
def pca(self, signal, pixel_size, g_min, g_max): """ :param signal: Estimated power spectrum. :param pixel_size: Pixel size in \u212b (Angstrom). :param g_min: Inverse of minimun resolution for PSD. :param g_max: Inverse of maximum resolution for PSD. :return: ratio. """ # RCOPT signal = signal.asnumpy()[0].T N = signal.shape[0] center = N // 2 grid = grid_2d(N, normalized=True, dtype=self.dtype) rb = np.sqrt((grid["x"] / 2)**2 + (grid["y"] / 2)**2).T r_ctf = rb * (10 / pixel_size) grid = grid_2d(N, normalized=False, dtype=self.dtype) X = grid["x"].T Y = grid["y"].T signal -= np.min(signal) rad_sq_min = N * pixel_size / g_min rad_sq_max = N * pixel_size / g_max min_limit = r_ctf[center, (center + np.floor(rad_sq_min)).astype(int)] signal[r_ctf < min_limit] = 0 max_limit = r_ctf[center, (center + np.ceil(rad_sq_max)).astype(int)] signal = np.where(r_ctf > max_limit, 0, signal) moment_02 = Y**2 * signal moment_02 = np.sum(moment_02, axis=(0, 1)) moment_11 = Y * X * signal moment_11 = np.sum(moment_11, axis=(0, 1)) moment_20 = X**2 * signal moment_20 = np.sum(moment_20, axis=(0, 1)) moment_mat = np.zeros((2, 2)) moment_mat[0, 0] = moment_20 moment_mat[1, 1] = moment_02 moment_mat[0, 1] = moment_11 moment_mat[1, 0] = moment_11 moment_evals = npla.eigvalsh(moment_mat) ratio = moment_evals[0] / moment_evals[1] return ratio
def circ(size, x0=0, y0=0, radius=1, peak=1, dtype=np.float64): """ Returns a 2d `circ` function in a square 2d numpy array. where for r = sqrt(x**2 + y**2) circ(x,y) = peak : 0 <= r <= radius 0 : otherwise Default is a centered circle of spread=peak=1. :param size: The height and width of returned array (pixels) :param x0: x cordinate of center (pixels) :param y0: y cordinate of center (pixels) :param radius: radius of circle :param peak: peak height at center :param dtype: dtype of returned array :return: Numpy array (2D) """ # Construct centered mesh g = grid_2d(size, shifted=True, normalized=False, dtype=dtype) vals = ((g["x"] - x0)**2 + (g["y"] - y0)**2) < radius * radius return (peak * vals).astype(dtype)
def gaussian_2d(size, x0=0, y0=0, sigma_x=1, sigma_y=1, peak=1, dtype=np.float64): """ Returns a 2d Gaussian in a square 2d numpy array. Default is a centered disc of spread=peak=1. :param size: The height and width of returned array (pixels) :param x0: x cordinate of center (pixels) :param y0: y cordinate of center (pixels) :param sigma_x: spread in x direction :param sigma_y: spread in y direction :param peak: peak height at center :param dtype: dtype of returned array :return: Numpy array (2D) """ # Construct centered mesh g = grid_2d(size, shifted=True, normalized=False, dtype=dtype) p = (g["x"] - x0)**2 / (2 * sigma_x**2) + (g["y"] - y0)**2 / (2 * sigma_y**2) return (peak * np.exp(-p)).astype(dtype, copy=False)
def estimate_noise_psd(self): """ :return: The estimated noise variance of the images in the Source used to create this estimator. TODO: How's this initial estimate of variance different from the 'estimate' method? """ # Run estimate using saved parameters g2d = grid_2d(self.L) mask = g2d['r'] >= self.bgRadius mean_est = 0 noise_psd_est = np.zeros((self.L, self.L)).astype(self.src.dtype) for i in range(0, self.n, self.batchSize): images = self.src.images(i, self.batchSize).asnumpy() images_masked = (images * np.expand_dims(mask, 2)) _denominator = self.n * np.sum(mask) mean_est += np.sum(images_masked) / _denominator im_masked_f = centered_fft2(images_masked) noise_psd_est += np.sum(np.abs(im_masked_f**2), axis=2) / _denominator mid = self.L // 2 noise_psd_est[mid, mid] -= mean_est**2 return noise_psd_est
def background_subtract_2d(self, signal, background_p1, max_col): """ Subtract background from estimated power spectrum :param signal: Estimated power spectrum :param background_p1: 1-D background estimation :param max_col: Internal variable, returned as the second parameter from opt1d. :return: 2-tuple of NumPy arrays (Estimated PSD without noise and estimated noise). """ signal = signal.asnumpy() N = signal.shape[1] grid = grid_2d(N, normalized=False, dtype=self.dtype) radii = np.sqrt((grid["x"] / 2)**2 + (grid["y"] / 2)**2).T background = np.zeros(signal.shape, dtype=self.dtype) for r in range(max_col + 2, background_p1.shape[1]): background[:, (r < radii) & (radii <= r + 1)] = background_p1[max_col, r] mask = radii <= max_col + 2 background[:, mask] = signal[:, mask] signal = signal - background signal = np.maximum(0, signal) return Image(signal), Image(background)
def unique_coords_nd(N, ndim, shifted=False, normalized=True): """ Generate unique polar coordinates from 2D or 3D rectangular coordinates. :param N: length size of a square or cube. :param ndim: number of dimension, 2 or 3. :param shifted: shifted half pixel or not for odd N. :param normalized: normalize the grid or not. :return: The unique polar coordinates in 2D or 3D """ ensure(ndim in (2, 3), 'Only two- or three-dimensional basis functions are supported.') ensure(N > 0, 'Number of grid points should be greater than 0.') if ndim == 2: grid = grid_2d(N, shifted=shifted, normalized=normalized) mask = grid['r'] <= 1 # Minor differences in r/theta/phi values are unimportant for the purpose # of this function, so round off before proceeding # TODO: numpy boolean indexing will return a 1d array (like MATLAB) # However, it always searches in row-major order, unlike MATLAB (column-major), # with no options to change the search order. The results we'll be getting back are thus not comparable. # We transpose the appropriate ndarrays before applying the mask to obtain the same behavior as MATLAB. r = grid['r'].T[mask].round(5) phi = grid['phi'].T[mask].round(5) r_unique, r_idx = np.unique(r, return_inverse=True) ang_unique, ang_idx = np.unique(phi, return_inverse=True) else: grid = grid_3d(N, shifted=shifted, normalized=normalized) mask = grid['r'] <= 1 # In Numpy, elements in the indexed array are always iterated and returned in row-major (C-style) order. # To emulate a behavior where iteration happens in Fortran order, we swap axes 0 and 2 of both the array # being indexed (r/theta/phi), as well as the mask itself. # TODO: This is only for the purpose of getting the same behavior as MATLAB while porting the code, and is # likely not needed in the final version. # Minor differences in r/theta/phi values are unimportant for the purpose of this function, # so we round off before proceeding. mask_ = np.swapaxes(mask, 0, 2) r = np.swapaxes(grid['r'], 0, 2)[mask_].round(5) theta = np.swapaxes(grid['theta'], 0, 2)[mask_].round(5) phi = np.swapaxes(grid['phi'], 0, 2)[mask_].round(5) r_unique, r_idx = np.unique(r, return_inverse=True) ang_unique, ang_idx = np.unique(np.vstack([theta, phi]), axis=1, return_inverse=True) return { 'r_unique': r_unique, 'ang_unique': ang_unique, 'r_idx': r_idx, 'ang_idx': ang_idx, 'mask': mask }
def rotated_grids(L, rot_matrices): """ Generate rotated Fourier grids in 3D from rotation matrices :param L: The resolution of the desired grids. :param rot_matrices: An array of size k-by-3-by-3 containing K rotation matrices :return: A set of rotated Fourier grids in three dimensions as specified by the rotation matrices. Frequencies are in the range [-pi, pi]. """ grid2d = grid_2d(L, dtype=rot_matrices.dtype) num_pts = L**2 num_rots = rot_matrices.shape[0] pts = np.pi * np.vstack([ grid2d["x"].flatten(), grid2d["y"].flatten(), np.zeros(num_pts, dtype=rot_matrices.dtype), ]) pts_rot = np.zeros((3, num_pts, num_rots), dtype=rot_matrices.dtype) for i in range(num_rots): pts_rot[:, :, i] = rot_matrices[i, :, :] @ pts pts_rot = pts_rot.reshape((3, L, L, num_rots)) # Note we return grids as (Z, Y, X) return pts_rot[::-1]
def evaluate_grid(self, L, dtype=np.float32, *args, **kwargs): grid2d = grid_2d(L, dtype=dtype) omega = np.pi * np.vstack( (grid2d["x"].flatten("F"), grid2d["y"].flatten("F"))) h = self.evaluate(omega, *args, **kwargs) h = m_reshape(h, grid2d["x"].shape) return h
def evaluate_grid(self, L, *args, **kwargs): grid2d = grid_2d(L) omega = np.pi * np.vstack( (grid2d['x'].flatten('F'), grid2d['y'].flatten('F'))) h = self.evaluate(omega, *args, **kwargs) h = m_reshape(h, grid2d['x'].shape) return h
def evaluate_grid(self, L, *args, **kwargs): # Todo: remove redundancy wrt a single Filter's evaluate_grid grid2d = grid_2d(L) omega = np.pi * np.vstack( (grid2d['x'].flatten('F'), grid2d['y'].flatten('F'))) h = self.evaluate(omega, *args, **kwargs) h = m_reshape(h, grid2d['x'].shape + (len(self.filters), )) return h
def testNormBackground(self): bg_radius = 1.0 grid = grid_2d(self.L) mask = grid["r"] > bg_radius self.sim.normalize_background() imgs_nb = self.sim.images(start=0, num=self.n).asnumpy() new_mean = np.mean(imgs_nb[:, mask]) new_variance = np.var(imgs_nb[:, mask]) # new mean of noise should be close to zero and variance should be close to 1 self.assertTrue(new_mean < 1e-7 and abs(new_variance - 1) < 1e-7)
def eval_filter_grid(self, L, power=1): grid2d = grid_2d(L, dtype=self.dtype) omega = np.pi * np.vstack((grid2d["x"].flatten(), grid2d["y"].flatten())) h = np.empty((omega.shape[-1], len(self.filter_indices)), dtype=self.dtype) for i, filt in enumerate(self.unique_filters): idx_k = np.where(self.filter_indices == i)[0] if len(idx_k) > 0: filter_values = filt.evaluate(omega) if power != 1: filter_values **= power h[:, idx_k] = np.column_stack((filter_values,) * len(idx_k)) h = np.reshape(h, grid2d["x"].shape + (len(self.filter_indices),)) return h
def invert_contrast(self, batch_size=512): """ invert the global contrast of images Check if all images in a stack should be globally phase flipped so that the molecule corresponds to brighter pixels and the background corresponds to darker pixels. This is done by comparing the mean in a small circle around the origin (supposed to correspond to the molecule) with the mean of the noise, and making sure that the mean of the molecule is larger. From the implementation level, we modify the `ImageSource` in-place by appending a `Multiple` filter to the generation pipeline. :param batch_size: Batch size of images to query. :return: On return, the `ImageSource` object has been modified in place. """ logger.info("Apply contrast inversion on source object") L = self.L grid = grid_2d(L, shifted=True) # Get mask indices of signal and noise samples assuming molecule signal_mask = grid["r"] < 0.5 noise_mask = grid["r"] > 0.8 # Calculate mean values in batch_size signal_mean = 0.0 noise_mean = 0.0 for i in range(0, self.n, batch_size): images = self.images(i, batch_size).asnumpy() signal = images * signal_mask noise = images * noise_mask signal_mean += np.sum(signal) noise_mean += np.sum(noise) signal_denominator = self.n * np.sum(signal_mask) noise_denominator = self.n * np.sum(noise_mask) signal_mean /= signal_denominator noise_mean /= noise_denominator if signal_mean < noise_mean: logger.info("Need to invert contrast") scale_factor = -1.0 else: logger.info("No need to invert contrast") scale_factor = 1.0 logger.info("Adding Scaling Xform to end of generation pipeline") self.generation_pipeline.add_xform(Multiply(scale_factor))
def eval_filter_grid(self, L, power=1): grid2d = grid_2d(L) omega = np.pi * np.vstack( (grid2d['x'].flatten(), grid2d['y'].flatten())) h = np.empty((omega.shape[-1], len(self.filters))) for f in set(self.filters): idx_k = np.where(self.filters == f)[0] if len(idx_k) > 0: filter_values = f.evaluate(omega) if power != 1: filter_values **= power h[:, idx_k] = np.column_stack((filter_values, ) * len(idx_k)) h = np.reshape(h, grid2d['x'].shape + (len(self.filters), )) return h
def normalize_bg(imgs, bg_radius=1.0, do_ramp=True): """ Normalize backgrounds and apply to a stack of images :param imgs: A stack of images in N-by-L-by-L array :param bg_radius: Radius cutoff to be considered as background (in image size) :param do_ramp: When it is `True`, fit a ramping background to the data and subtract. Namely perform normalization based on values from each image. Otherwise, a constant background level from all images is used. :return: The modified images """ L = imgs.shape[-1] grid = grid_2d(L) mask = grid["r"] > bg_radius if do_ramp: # Create matrices and reshape the background mask # for fitting a ramping background ramp_mask = np.vstack(( grid["x"][mask].flatten(), grid["y"][mask].flatten(), np.ones(grid["y"][mask].flatten().size), )).T ramp_all = np.vstack( (grid["x"].flatten(), grid["y"].flatten(), np.ones(L * L))).T mask_reshape = mask.reshape((L * L)) imgs = imgs.reshape((-1, L * L)) # Fit a ramping background and apply to images coeff = lstsq(ramp_mask, imgs[:, mask_reshape].T)[0] # RCOPT imgs = imgs - (ramp_all @ coeff).T # RCOPT imgs = imgs.reshape((-1, L, L)) # Apply mask images and calculate mean and std values of background imgs_masked = imgs * mask denominator = np.sum(mask) first_moment = np.sum(imgs_masked, axis=(1, 2)) / denominator second_moment = np.sum(imgs_masked**2, axis=(1, 2)) / denominator mean = first_moment.reshape(-1, 1, 1) variance = second_moment.reshape(-1, 1, 1) - mean**2 std = np.sqrt(variance) return (imgs - mean) / std
def __init__( self, pixel_size, cs, amplitude_contrast, voltage, psd_size, num_tapers, dtype=np.float32, ): """ Instantiate a CtfEstimator instance. :param pixel_size: Size of the pixel in \u212b (Angstrom). :param cs: Spherical aberration in mm. :param amplitude_contrast: Amplitude contrast. :param voltage: Voltage of electron microscope. :param psd_size: Block size (in pixels) for PSD estimation. :param num_tapers: Number of tapers to apply in PSD estimation. :returns: CtfEstimator instance. """ self.pixel_size = pixel_size self.cs = cs self.amplitude_contrast = amplitude_contrast self.voltage = voltage self.psd_size = psd_size self.num_tapers = num_tapers self.lmbd = voltage_to_wavelength(voltage) / 10.0 # (Angstrom) self.dtype = np.dtype(dtype) grid = grid_2d(psd_size, normalized=True, dtype=self.dtype) # Note this mesh for x,y is transposed, and range is -half to half. rb = np.sqrt((grid["x"] / 2)**2 + (grid["y"] / 2)**2).T self.r_ctf = rb * (10 / pixel_size) # units: inverse nm # Note this mesh for theta is transposed. self.theta = grid["phi"].T self.defocus1 = 0 self.defocus2 = 0 self.angle = 0 # Radians self.h = 0
def _estimate_noise_variance(self): """ Any additional arguments/keyword-arguments are passed on to the Source's 'images' method :return: The estimated noise variance of the images in the Source used to create this estimator. TODO: How's this initial estimate of variance different from the 'estimate' method? """ # Run estimate using saved parameters g2d = grid_2d(self.L) mask = g2d['r'] >= self.bgRadius first_moment = 0 second_moment = 0 for i in range(0, self.n, self.batchSize): images = self.src.images(start=i, num=self.batchSize) images_masked = (images * np.expand_dims(mask, 2)) _denominator = self.n * np.sum(mask) first_moment += np.sum(images_masked) / _denominator second_moment += np.sum(np.abs(images_masked**2)) / _denominator return second_moment - first_moment**2
def evaluate_grid(self, L, dtype=np.float32, *args, **kwargs): """ Generates a two dimensional grid with prescribed dtype, yielding the values (omega) which are then evaluated by the filter's evaluate method. Passes arbritrary args and kwargs down to self.evaluate method. :param L: Number of grid points (L by L). :param dtype: dtype of grid, defaults np.float32. :return: Filter values at omega's points. """ grid2d = grid_2d(L, dtype=dtype) omega = np.pi * np.vstack( (grid2d["x"].flatten("F"), grid2d["y"].flatten("F"))) h = self.evaluate(omega, *args, **kwargs) h = m_reshape(h, grid2d["x"].shape) return h
def rotated_grids(L, rot_matrices): """ Generate rotated Fourier grids in 3D from rotation matrices :param L: The resolution of the desired grids. :param rot_matrices: An array of size 3-by-3-by-K containing K rotation matrices :return: A set of rotated Fourier grids in three dimensions as specified by the rotation matrices. Frequencies are in the range [-pi, pi]. """ # TODO: Flattening and reshaping at end may not be necessary! grid2d = grid_2d(L) num_pts = L**2 num_rots = rot_matrices.shape[-1] pts = np.pi * np.vstack([ grid2d['x'].flatten('F'), grid2d['y'].flatten('F'), np.zeros(num_pts) ]) pts_rot = np.zeros((3, num_pts, num_rots)) for i in range(num_rots): pts_rot[:, :, i] = rot_matrices[:, :, i] @ pts pts_rot = m_reshape(pts_rot, (3, L, L, num_rots)) return pts_rot
def inverse_r(size, x0=0, y0=0, peak=1, dtype=np.float64): """ Returns a 2d inverse radius function in a square 2d numpy array. Where inverse_r(x,y): 1/sqrt(1 + x**2 + y**2) Default is a centered circle of peak=1. :param size: The height and width of returned array (pixels) :param x0: x cordinate of center (pixels) :param y0: y cordinate of center (pixels) :param peak: peak height at center :param dtype: dtype of returned array :return: Numpy array (2D) """ # Construct centered mesh g = grid_2d(size, shifted=True, normalized=False, dtype=dtype) # Compute the denominator vals = np.sqrt(1 + (g["x"] - x0)**2 + (g["y"] - y0)**2) return (peak / vals).astype(dtype)
def downsample(insamples, szout): """ Blur and downsample 1D to 3D objects such as, curves, images or volumes :param insamples: Set of objects to be downsampled in the form of an array, the last dimension is the number of objects. :param szout: The desired resolution of for output objects. :return: An array consists of the blurred and downsampled objects. """ ensure( insamples.ndim - 1 == szout.ndim, 'The number of downsampling dimensions is not the same as that of objects.' ) L_in = insamples.shape[0] L_out = szout.shape[0] ndata = insamples.shape(-1) outdims = szout outdims.push_back(ndata) outsamples = np.zeros((outdims)).astype(insamples.dtype) if insamples.ndim == 2: # one dimension object grid_in = grid_1d(L_in) grid_out = grid_1d(L_out) # x values corresponding to 'grid'. This is what scipy interpolator needs to function. x = np.ceil(np.arange(-L_in / 2, L_in / 2)) / (L_in / 2) mask = (np.abs(grid_in['x']) < L_out / L_in) insamples_fft = np.real( centered_ifft1(centered_fft1(insamples) * np.expand_dims(mask, 1))) for idata in range(ndata): interpolator = RegularGridInterpolator((x, ), insamples_fft[:, idata], bounds_error=False, fill_value=0) outsamples[:, :, idata] = interpolator(np.dstack([grid_out['x']])) elif insamples.ndim == 3: grid_in = grid_2d(L_in) grid_out = grid_2d(L_out) # x, y values corresponding to 'grid'. This is what scipy interpolator needs to function. x = y = np.ceil(np.arange(-L_in / 2, L_in / 2)) / (L_in / 2) mask = (np.abs(grid_in['x']) < L_out / L_in) & (np.abs(grid_in['y']) < L_out / L_in) insamples_fft = np.real( centered_ifft2(centered_fft2(insamples) * np.expand_dims(mask, 2))) for idata in range(ndata): interpolator = RegularGridInterpolator((x, y), insamples_fft[:, :, idata], bounds_error=False, fill_value=0) outsamples[:, :, idata] = interpolator( np.dstack([grid_out['x'], grid_out['y']])) elif insamples.ndim == 4: grid_in = grid_3d(L_in) grid_out = grid_3d(L_out) # x, y, z values corresponding to 'grid'. This is what scipy interpolator needs to function. x = y = z = np.ceil(np.arange(-L_in / 2, L_in / 2)) / (L_in / 2) mask = (np.abs(grid_in['x']) < L_out / L_in) & (np.abs( grid_in['y']) < L_out / L_in) & (np.abs(grid_in['z']) < L_out / L_in) insamples_fft = np.real( centered_ifft3(centered_fft3(insamples) * np.expand_dims(mask, 3))) for idata in range(ndata): interpolator = RegularGridInterpolator((x, y, z), insamples_fft[:, :, :, idata], bounds_error=False, fill_value=0) outsamples[:, :, :, idata] = interpolator( np.dstack([grid_out['x'], grid_out['y'], grid_out['z']])) return outsamples
def vol2img(volume, rots, L=None, dtype=None): """ Generate 2D images from the input volume and rotation angles The function handles odd and even-sized arrays correctly. The center of an odd array is taken to be at (n+1)/2, and an even array is n/2+1. :param volume: A 3D volume objects. :param rots: A n-by-3-by-3 array of rotation angles. :param L: The output size of 2D images. :return: An array consists of 2D images. """ if L is None: L = np.size(volume, 0) if dtype is None: dtype = volume.dtype lv = np.size(volume, 0) if L > lv + 1: # For compatibility with gen_projections, allow one pixel aliasing. # More precisely, it should be N>nv, however, by using nv+1 the # results match those of gen_projections. if np.mod(L - lv, 2) == 1: raise RuntimeError( 'Upsampling from odd to even sizes or vice versa is ' 'currently not supported') dL = np.floor((L - lv) / 2) fv = centered_fft3(volume) padded_volume = np.zeros((L, L, L), dtype=dtype) padded_volume[dL + 1:dL + lv + 1, dL + 1:dL + lv + 1, dL + 1:dL + lv + 1] = fv volume = centered_ifft3(padded_volume) ensure( np.norm(np.imag(volume[:])) / np.norm(volume[:]) < 1.0e-5, "The image part of volume is related large (>1.0e-5).") # The new volume size lv = L grid2d = grid_2d(lv, shifted=True, normalized=False) num_pts = lv**2 num_rots = rots.shape[0] pts = np.pi * np.vstack([ grid2d['x'].flatten('F'), grid2d['y'].flatten('F'), np.zeros(num_pts) ]) pts_rot = np.zeros((3, num_pts, num_rots)) for i in range(num_rots): pts_rot[:, :, i] = rots[i, :, :].T @ pts pts_rot = m_reshape(pts_rot, (3, lv**2 * num_rots)) pts_rot = -2 * pts_rot / lv im_f = Plan(volume.shape, -pts_rot).transform(volume) im_f = m_reshape(im_f, (lv, lv, -1)) if lv % 2 == 0: pts_rot = m_reshape(pts_rot, (3, lv, lv, num_rots)) im_f = im_f * np.exp(1j * np.sum(pts_rot, 0) / 2) im_f = im_f * np.expand_dims( np.exp(2 * np.pi * 1j * (grid2d['x'] + grid2d['y'] - 1) / (2 * lv)), 2) im = centered_ifft2(im_f) if lv % 2 == 0: im = im * m_reshape( np.exp(2 * np.pi * 1j * (grid2d['x'] + grid2d['y']) / (2 * lv)), (lv, lv, 1)) return np.real(im)
def adaptive_support(img_src, energy_threshold=0.99): """ Determine size of the compact support in both real and Fourier Space. Returns c_limit (support radius in Fourier space), and R_limit (support radius in real space). Fourier c_limit is scaled in range [0, 0.5]. R_limit is in pixels [0, Image.res/2]. :param img_src: Input `Source` of images. :param energy_threshold: [0, 1] threshold limit :return: (c_limit, R_limit) """ if not isinstance(img_src, ImageSource): raise RuntimeError( "adaptive_support expects `Source` instance or subclass.") # Sanity Check Threshold is in range if energy_threshold <= 0 or energy_threshold > 1: raise ValueError( f"Given energy_threshold {energy_threshold} outside sane range [0,1]" ) L = img_src.L N = L // 2 r = grid_2d(L, shifted=False, normalized=False, dtype=img_src.dtype)["r"] # Estimate noise noise_est = WhiteNoiseEstimator(img_src) noise_var = noise_est.estimate() # Transform to Fourier space img = img_src.images(0, img_src.n).asnumpy() imgf = fft.centered_fft2(img) # Compute the Variance and Power Spectrum # Mean along image stack. variance_map = np.mean(np.abs(img)**2, axis=0) pspec = np.mean(np.abs(imgf)**2, axis=0) # Compute the Radial Variance and Radial Power Spectrum radial_var = np.zeros(N) radial_pspec = np.zeros(N) for i in range(N): mask = (r >= i) & (r < i + 1) # Mean along radial track defined by mask radial_var[i] = np.mean(variance_map[mask]) radial_pspec[i] = np.mean(pspec[mask]) # Subtract the noise variance radial_pspec -= noise_var radial_var -= noise_var # Lower bound variance and power by 0 np.clip(radial_pspec, 0, a_max=None, out=radial_pspec) np.clip(radial_var, 0, a_max=None, out=radial_var) # Construct range of Fourier limits. We need a half-sample correction # since each ring is centered between two integer radii. Same for spatial # domain (R). c = (np.arange(N) + 0.5) / (2 * N) R = np.arange(N) + 0.5 # Calculate cumulative energy cum_pspec = np.cumsum(radial_pspec * c) cum_var = np.cumsum(radial_var * R) # Normalize energies [0,1] # Multiply threshold to avoid unstable division c_energy_threshold = energy_threshold * cum_pspec[-1] R_energy_threshold = energy_threshold * cum_var[-1] # First note legacy code *=L for Fourier limit, # but then only uses divided by L... so just removed here. # This makes it consistent with Nyquist, ie [0, .5] # Second note, we attempt to find the cutoff, # but when a search fails returns the last (-1) element, # essentially the maximal radius. # Third note, to increase accuracy, we take a weighted average of the two # points around the cutoff. This mostly affects c since R is rounded. ind = np.argmax(cum_pspec > c_energy_threshold) if ind > 0: c_limit = (cum_pspec[ind - 1] * c[ind - 1] + cum_pspec[ind] * c[ind]) / (cum_pspec[ind - 1] + cum_pspec[ind]) else: c_limit = c[-1] ind = np.argmax(cum_var > R_energy_threshold) if ind > 0: R_limit = round( (cum_var[ind - 1] * R[ind - 1] + cum_var[ind] * R[ind]) / (cum_var[ind - 1] + cum_var[ind])) else: R_limit = R[-1] return c_limit, R_limit
def testGrid2d(self): grid2d = grid_2d(8) self.assertTrue(np.allclose(grid2d['x'], np.load(os.path.join(DATA_DIR, 'grid2d_8_x.npy')))) self.assertTrue(np.allclose(grid2d['y'], np.load(os.path.join(DATA_DIR, 'grid2d_8_y.npy')))) self.assertTrue(np.allclose(grid2d['r'], np.load(os.path.join(DATA_DIR, 'grid2d_8_r.npy')))) self.assertTrue(np.allclose(grid2d['phi'], np.load(os.path.join(DATA_DIR, 'grid2d_8_phi.npy'))))
def estimate_ctf( data_folder, pixel_size, cs, amplitude_contrast, voltage, num_tapers, psd_size, g_min, g_max, output_dir, dtype=np.float32, ): """ Given paramaters estimates CTF from experimental data and returns CTF as a mrc file. """ dtype = np.dtype(dtype) assert dtype in (np.float32, np.float64) dir_content = os.scandir(data_folder) mrc_files = [ f.name for f in dir_content if os.path.splitext(f)[1] == ".mrc" ] mrcs_files = [ f.name for f in dir_content if os.path.splitext(f)[1] == ".mrcs" ] file_names = mrc_files + mrcs_files amp = amplitude_contrast amplitude_contrast = np.arctan(amplitude_contrast / np.sqrt(1 - amplitude_contrast**2)) lmbd = voltage_to_wavelength(voltage) / 10 # (Angstrom) ctf_object = CtfEstimator(pixel_size, cs, amplitude_contrast, voltage, psd_size, num_tapers, dtype=dtype) # Note for repro debugging, suggest use of doubles, # closer to original code. ffbbasis = FFBBasis2D((psd_size, psd_size), 2, dtype=dtype) results = [] for name in file_names: with mrcfile.open(os.path.join(data_folder, name), mode="r", permissive=True) as mrc: micrograph = mrc.data # Try to match dtype used in Basis instance micrograph = micrograph.astype(dtype, copy=False) micrograph_blocks = ctf_object.preprocess_micrograph( micrograph, psd_size) tapers_1d = ctf_object.tapers(psd_size, num_tapers / 2, num_tapers) signal_observed = ctf_object.estimate_psd(micrograph_blocks, tapers_1d) amplitude_spectrum, _ = ctf_object.elliptical_average( ffbbasis, signal_observed, True) # absolute differenceL 10^-14. Relative error: 10^-7 # Optionally changing to: linprog_method='simplex', # will more deterministically repro results in exchange for speed. signal_1d, background_1d = ctf_object.background_subtract_1d( amplitude_spectrum, linprog_method="interior-point") avg_defocus, low_freq_skip = ctf_object.opt1d( signal_1d, pixel_size, cs, lmbd, # (Angstrom) amplitude_contrast, signal_observed.shape[-1], ) low_freq_skip = 12 signal, background_2d = ctf_object.background_subtract_2d( signal_observed, background_1d, low_freq_skip) ratio = ctf_object.pca(signal_observed, pixel_size, g_min, g_max) signal, additional_background = ctf_object.elliptical_average( ffbbasis, signal.sqrt(), False) background_2d = background_2d + additional_background initial_df1 = (avg_defocus * 2) / (1 + ratio) initial_df2 = (avg_defocus * 2) - initial_df1 grid = grid_2d(psd_size, normalized=True, dtype=dtype) rb = np.sqrt((grid["x"] / 2)**2 + (grid["y"] / 2)**2).T r_ctf = rb * (10 / pixel_size) theta = grid["phi"].T angle = -5 / 12 * np.pi # Radians (-75 degrees) cc_array = np.zeros((6, 4)) for a in range(0, 6): df1, df2, angle_ast, p = ctf_object.gd( signal, initial_df1, initial_df2, angle + a * np.pi / 6.0, # Radians, + a*30degrees r_ctf, theta, pixel_size, g_min, g_max, amplitude_contrast, lmbd, # (Angstrom) cs, ) cc_array[a, 0] = df1 cc_array[a, 1] = df2 cc_array[a, 2] = angle_ast # Radians cc_array[a, 3] = p ml = np.argmax(cc_array[:, 3], -1) result = ( cc_array[ml, 0], cc_array[ml, 1], cc_array[ml, 2], # Radians cs, voltage, pixel_size, amp, name, ) ctf_object.write_star(*result, output_dir) results.append(result) ctf_object.set_df1(cc_array[ml, 0]) ctf_object.set_df2(cc_array[ml, 1]) ctf_object.set_angle(cc_array[ml, 2]) # Radians ctf_object.generate_ctf() with mrcfile.new(output_dir + "/" + os.path.splitext(name)[0] + "_noise.mrc", overwrite=True) as mrc: mrc.set_data(background_2d[0].astype(np.float32)) mrc.voxel_size = pixel_size mrc.close() df = (cc_array[ml, 0] + cc_array[ml, 1]) * np.ones( theta.shape, theta.dtype) + (cc_array[ml, 0] - cc_array[ ml, 1]) * np.cos(2 * theta - 2 * cc_array[ml, 2] * np.ones(theta.shape, theta.dtype)) ctf_im = -np.sin(np.pi * lmbd * r_ctf**2 / 2 * (df - lmbd**2 * r_ctf**2 * cs * 1e6) + amplitude_contrast) ctf_signal = np.zeros(ctf_im.shape, ctf_im.dtype) ctf_signal[:ctf_im.shape[0] // 2, :] = ctf_im[:ctf_im.shape[0] // 2, :] ctf_signal[ctf_im.shape[0] // 2 + 1:, :] = signal[:, :, ctf_im.shape[0] // 2 + 1] with mrcfile.new(output_dir + "/" + os.path.splitext(name)[0] + ".ctf", overwrite=True) as mrc: mrc.set_data(np.float32(ctf_signal)) mrc.voxel_size = pixel_size mrc.close() return results