def get_mean(self, coeffs, ctf_fb=None, ctf_idx=None): """ Calculate the mean vector from the expansion coefficients with CTF information. :param coeffs: A coefficient vector (or an array of coefficient vectors) to be averaged. :param ctf_fb: The CFT functions in the FB expansion. :param ctf_idx: An array of the CFT function indices for all 2D images. If ctf_fb or ctf_idx is None, the identity filter will be applied. :return: The mean value vector for all images. """ if coeffs.size == 0: raise RuntimeError('The coefficients need to be calculated!') if (ctf_fb is None) or (ctf_idx is None): ctf_idx = np.zeros(coeffs.shape[1], dtype=int) ctf_fb = [BlkDiagMatrix.eye_like(RadialCTFFilter().fb_mat(self.basis))] b = np.zeros(self.basis.count, dtype=coeffs.dtype) A = BlkDiagMatrix.zeros_like(ctf_fb[0]) for k in np.unique(ctf_idx[:]).T: coeff_k = coeffs[:, ctf_idx == k] weight = np.size(coeff_k, 1)/np.size(coeffs, 1) mean_coeff_k = self._get_mean(coeff_k) ctf_fb_k = ctf_fb[k] ctf_fb_k_t = ctf_fb_k.T b += weight * ctf_fb_k_t.apply(mean_coeff_k) A += weight * (ctf_fb_k_t @ ctf_fb_k) mean_coeff = A.solve(b) return mean_coeff
def testBlkDiagMatrixEye(self): result = [ np.array([[1., 0., 0., 0.], [0., 1., 0., 0.], [0., 0., 1., 0.], [0., 0., 0., 1.]]), np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]), np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]), np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]), np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]), np.array([[1., 0.], [0., 1.]]), np.array([[1., 0.], [0., 1.]]), np.array([[1., 0.], [0., 1.]]), np.array([[1., 0.], [0., 1.]]), np.array([[1., 0.], [0., 1.]]), np.array([[1., 0.], [0., 1.]]), np.array([[1.]]), np.array([[1.]]), np.array([[1.]]), np.array([[1.]]), np.array([[1.]]), np.array([[1.]]) ] blk_eye = BlkDiagMatrix.eye(self.blk_partition) self.allallfunc(blk_eye, result) blk_eye = BlkDiagMatrix.eye_like(self.blk_a) self.allallfunc(blk_eye, result)
def testBlkDiagMatrixZeros(self): result = [ np.array([[0., 0., 0., 0.], [0., 0., 0., 0.], [0., 0., 0., 0.], [0., 0., 0., 0.]]), np.array([[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]), np.array([[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]), np.array([[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]), np.array([[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]), np.array([[0., 0.], [0., 0.]]), np.array([[0., 0.], [0., 0.]]), np.array([[0., 0.], [0., 0.]]), np.array([[0., 0.], [0., 0.]]), np.array([[0., 0.], [0., 0.]]), np.array([[0., 0.], [0., 0.]]), np.array([[0.]]), np.array([[0.]]), np.array([[0.]]), np.array([[0.]]), np.array([[0.]]), np.array([[0.]]) ] blk_zeros = BlkDiagMatrix.zeros(self.blk_partition) self.allallfunc(blk_zeros, result) blk_zeros = BlkDiagMatrix.zeros_like(self.blk_a) self.allallfunc(blk_zeros, result)
def _calc_op(self): src = self.src basis = self.basis ctf_fb = self.ctf_fb ctf_idx = self.ctf_idx A_mean = BlkDiagMatrix.zeros_like(ctf_fb[0]) A_covar = [None for _ in ctf_fb] M_covar = BlkDiagMatrix.zeros_like(ctf_fb[0]) for k in np.unique(ctf_idx): weight = np.count_nonzero(ctf_idx == k) / src.n ctf_fb_k = ctf_fb[k] ctf_fb_k_t = ctf_fb_k.T ctf_fb_k_sq = ctf_fb_k_t @ ctf_fb_k A_mean_k = weight * ctf_fb_k_sq A_mean += A_mean_k A_covar_k = np.sqrt(weight) * ctf_fb_k_sq A_covar[k] = A_covar_k M_covar += A_covar_k self.A_mean = A_mean self.A_covar = A_covar self.M_covar = M_covar
def testBlkDiagMatrixCompat(self): """ Check incompatible matrix raises exception. """ # Create a differently shaped matrix x = BlkDiagMatrix.from_list(self.blk_a[1:-1]) # code should raise with pytest.raises(RuntimeError): _ = x + self.blk_a
def testBlkDiagMatrixInPlace(self): """ Tests sequence of in place optimized arithmetic (add, sub, mul) """ result_1 = [x + x + 10. for x in self.blk_a] result_2 = [np.ones(x.shape) * 5. for x in self.blk_a] # make a block diagonal object to mutate blk_c = self.blk_a.copy() # store the python object ids of each array in blk_c # we want to ensure the actual object refs are _not_ changing # for in place operations. id0 = [id(x) for x in blk_c] blk_c += self.blk_a self.allallid(blk_c, id0) blk_c += 10. self.allallid(blk_c, id0) blk_a5 = BlkDiagMatrix.ones(self.blk_partition) id1 = [id(x) for x in blk_a5] blk_a5 *= 5. self.allallid(blk_a5, id1) blk_c -= blk_a5 blk_c -= blk_a5 self.allallid(blk_c, id0) blk_c -= self.blk_a self.allallid(blk_c, id0) self.allallfunc(blk_c, self.blk_a)
def _solve_covar(self, A_covar, b_covar, M, covar_est_opt): ctf_fb = self.ctf_fb def precond_fun(S, x): p = np.size(S, 0) ensure(np.size(x) == p*p, 'The sizes of S and x are not consistent.') x = m_reshape(x, (p, p)) y = S @ x @ S y = m_reshape(y, (p**2,)) return y def apply(A, x): p = np.size(A[0], 0) x = m_reshape(x, (p, p)) y = np.zeros_like(x) for k in range(0, len(A)): y = y + A[k] @ x @ A[k].T y = m_reshape(y, (p**2,)) return y cg_opt = covar_est_opt covar_coeff = BlkDiagMatrix.zeros_like(ctf_fb[0]) for ell in range(0, len(b_covar)): A_ell = [] for k in range(0, len(A_covar)): A_ell.append(A_covar[k][ell]) p = np.size(A_ell[0], 0) b_ell = m_reshape(b_covar[ell], (p ** 2,)) S = inv(M[ell]) cg_opt['preconditioner'] = lambda x: precond_fun(S, x) covar_coeff_ell, _, _ = conj_grad(lambda x: apply(A_ell, x), b_ell, cg_opt) covar_coeff[ell] = m_reshape(covar_coeff_ell, (p, p)) return covar_coeff
def get_cwf_coeffs(self, coeffs, ctf_fb, ctf_idx, mean_coeff, covar_coeff, noise_var=1): """ Estimate the expansion coefficients using the Covariance Wiener Filtering (CWF) method. :param coeffs: A coefficient vector (or an array of coefficient vectors) to be calculated. :param ctf_fb: The CFT functions in the FB expansion. :param ctf_idx: An array of the CFT function indices for all 2D images. If ctf_fb or ctf_idx is None, the identity filter will be applied. :param mean_coeff: The mean value vector from all images. :param covar_coeff: The block diagonal covariance matrix of the clean coefficients represented by a cell array. :param noise_var: The estimated variance of noise. The value should be zero for `coeffs` from clean images of simulation data. :return: The estimated coefficients of the unfiltered images in certain math basis. These are obtained using a Wiener filter with the specified covariance for the clean images and white noise of variance `noise_var` for the noise. """ if mean_coeff is None: mean_coeff = self.get_mean() if covar_coeff is None: covar_coeff = self.get_covar(noise_var=noise_var, mean_coeff=mean_coeff) if (ctf_fb is None) or (ctf_idx is None): ctf_idx = np.zeros(coeffs.shape[1], dtype=int) ctf_fb = [BlkDiagMatrix.eye_like(covar_coeff)] noise_covar_coeff = noise_var * BlkDiagMatrix.eye_like(covar_coeff) coeffs_est = np.zeros_like(coeffs) for k in np.unique(ctf_idx[:]): coeff_k = coeffs[:, ctf_idx == k] ctf_fb_k = ctf_fb[k] ctf_fb_k_t = ctf_fb_k.T sig_covar_coeff = ctf_fb_k @ covar_coeff @ ctf_fb_k_t sig_noise_covar_coeff = sig_covar_coeff + noise_covar_coeff mean_coeff_k = ctf_fb_k.apply(mean_coeff) coeff_est_k = coeff_k - mean_coeff_k[:, np.newaxis] coeff_est_k = sig_noise_covar_coeff.solve(coeff_est_k) coeff_est_k = (covar_coeff @ ctf_fb_k_t).apply(coeff_est_k) coeff_est_k = coeff_est_k + mean_coeff[:, np.newaxis] coeffs_est[:, ctf_idx == k] = coeff_est_k return coeffs_est
def _get_covar(self, coeffs, mean_coeff=None, do_refl=True): """ Calculate the covariance matrix from the expansion coefficients without CTF information. :param coeffs: A coefficient vector (or an array of coefficient vectors) calculated from 2D images. :param mean_coeff: The mean vector calculated from the `coeffs`. :param do_refl: If true, enforce invariance to reflection (default false). :return: The covariance matrix of coefficients for all images. """ if coeffs.size == 0: raise RuntimeError('The coefficients need to be calculated first!') if mean_coeff is None: mean_coeff = self._get_mean(coeffs) # Initialize a totally empty BlkDiagMatrix, build incrementally. covar_coeff = BlkDiagMatrix.empty(0, dtype=coeffs.dtype) ell = 0 mask = self.basis._indices["ells"] == ell coeff_ell = coeffs[mask, ...] - mean_coeff[mask, np.newaxis] covar_ell = np.array(coeff_ell @ coeff_ell.T/np.size(coeffs, 1)) covar_coeff.append(covar_ell) for ell in range(1, self.basis.ell_max+1): mask = self.basis._indices["ells"] == ell mask_pos = [mask[i] and (self.basis._indices['sgns'][i] == +1) for i in range(len(mask))] mask_neg = [mask[i] and (self.basis._indices['sgns'][i] == -1) for i in range(len(mask))] covar_ell_diag = np.array(coeffs[mask_pos, :] @ coeffs[mask_pos, :].T + coeffs[mask_neg, :] @ coeffs[mask_neg, :].T) / (2 * np.size(coeffs, 1)) if do_refl: covar_coeff.append(covar_ell_diag) covar_coeff.append(covar_ell_diag) else: covar_ell_off = np.array((coeffs[mask_pos, :] @ coeffs[mask_neg, :].T / np.size(coeffs, 1) - coeffs[mask_neg, :] @ coeffs[mask_pos, :].T)/(2 * np.size(coeffs, 1))) hsize = np.size(covar_ell_diag, 0) covar_coeff_blk = np.zeros((2 * hsize, 2 * hsize)) fsize = np.size(covar_coeff_blk, 0) covar_coeff_blk[0:hsize, 0:hsize] = covar_ell_diag[0:hsize, 0:hsize] covar_coeff_blk[hsize:fsize, hsize:fsize] = covar_ell_diag[0:hsize, 0:hsize] covar_coeff_blk[0:hsize, hsize:fsize] = covar_ell_off[0:hsize, 0:hsize] covar_coeff_blk[hsize:fsize, 0:hsize] = covar_ell_off.T[0:hsize, 0:hsize] covar_coeff.append(covar_coeff_blk) return covar_coeff
def _build(self): src = self.src if self.basis is None: from aspire.basis.ffb_2d import FFBBasis2D self.basis = FFBBasis2D((src.L, src.L)) if src.filters is None: logger.info(f'CTF filters are not included in Cov2D denoising') # set all CTF filters to an identity filter self.ctf_idx = np.zeros(src.n, dtype=int) self.ctf_fb = [BlkDiagMatrix.eye_like(RadialCTFFilter().fb_mat(self.basis))] else: logger.info(f'Represent CTF filters in FB basis') unique_filters = list(set(src.filters)) self.ctf_idx = np.array([unique_filters.index(f) for f in src.filters]) self.ctf_fb = [f.fb_mat(self.basis) for f in unique_filters]
def _calc_rhs(self): src = self.src basis = self.basis ctf_fb = self.ctf_fb ctf_idx = self.ctf_idx zero_coeff = np.zeros((basis.count,)) b_mean = [np.zeros(basis.count) for _ in ctf_fb] b_covar = BlkDiagMatrix.zeros_like(ctf_fb[0]) for start in range(0, src.n, self.batch_size): batch = np.arange(start, min(start + self.batch_size, src.n)) im = src.images(batch[0], len(batch)) coeff = basis.evaluate_t(im.data) for k in np.unique(ctf_idx[batch]): coeff_k = coeff[:, ctf_idx[batch] == k] weight = np.size(coeff_k, 1) / src.n mean_coeff_k = self._get_mean(coeff_k) ctf_fb_k = ctf_fb[k] ctf_fb_k_t = ctf_fb_k.T b_mean_k = weight * ctf_fb_k_t.apply(mean_coeff_k) b_mean[k] += b_mean_k covar_coeff_k = self._get_covar(coeff_k, zero_coeff) b_covar_k = ctf_fb_k_t @ covar_coeff_k b_covar_k = b_covar_k @ ctf_fb_k b_covar_k *= weight b_covar += b_covar_k self.b_mean = b_mean self.b_covar = b_covar
def testBlkDiagMatrixOnes(self): result = [ np.array([[1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.]]), np.array([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]), np.array([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]), np.array([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]), np.array([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]), np.array([[1., 1.], [1., 1.]]), np.array([[1., 1.], [1., 1.]]), np.array([[1., 1.], [1., 1.]]), np.array([[1., 1.], [1., 1.]]), np.array([[1., 1.], [1., 1.]]), np.array([[1., 1.], [1., 1.]]), np.array([[1.]]), np.array([[1.]]), np.array([[1.]]), np.array([[1.]]), np.array([[1.]]), np.array([[1.]]) ] blk_ones = BlkDiagMatrix.ones(self.blk_partition) self.allallfunc(blk_ones, result)
def setUp(self): self.blk_a = BlkDiagMatrix.from_list([ np.array([[-0.30656809, -0.34287864, -0.00854488, 0.5275285], [-0.34287864, -0.19752432, 0.17833916, -0.22052178], [-0.00854488, 0.17833916, -0.4125285, -0.30338836], [0.5275285, -0.22052178, -0.30338836, 0.12254553]]), np.array([[-0.08041961, -0.29729055, -0.3960436], [-0.29729055, -0.56196307, 0.0607334], [-0.3960436, 0.0607334, -0.09843568]]), np.array([[-0.08041961, -0.29729055, -0.3960436], [-0.29729055, -0.56196307, 0.0607334], [-0.3960436, 0.0607334, -0.09843568]]), np.array([[0.04041207, -0.05335218, -0.55965714], [-0.05335218, -0.62303735, -0.26739795], [-0.55965714, -0.26739795, 0.07961834]]), np.array([[0.04041207, -0.05335218, -0.55965714], [-0.05335218, -0.62303735, -0.26739795], [-0.55965714, -0.26739795, 0.07961834]]), np.array([[0.05209133, 0.17877923], [0.17877923, -0.43246999]]), np.array([[0.05209133, 0.17877923], [0.17877923, -0.43246999]]), np.array([[0.00296652, 0.32398467], [0.32398467, -0.12518209]]), np.array([[0.00296652, 0.32398467], [0.32398467, -0.12518209]]), np.array([[-0.06749134, 0.37452234], [0.37452234, 0.1789093]]), np.array([[-0.06749134, 0.37452234], [0.37452234, 0.1789093]]), np.array([[-0.13551364]]), np.array([[-0.13551364]]), np.array([[-0.18975111]]), np.array([[-0.18975111]]), np.array([[-0.22661312]]), np.array([[-0.22661312]]) ]) self.blk_b = BlkDiagMatrix.from_list([ np.array([[-0.27457111, -0.33770709, -0.09067737, 0.52007584], [-0.33770709, -0.24677034, 0.20639731, -0.2078888], [-0.09067737, 0.20639731, -0.36507922, -0.37742705], [0.52007584, -0.2078888, -0.37742705, 0.30641696]]), np.array([[-0.05219921, -0.23403317, -0.46057571], [-0.23403317, -0.58669018, 0.02892691], [-0.46057571, 0.02892691, -0.01291359]]), np.array([[-0.05219921, -0.23403317, -0.46057571], [-0.23403317, -0.58669018, 0.02892691], [-0.46057571, 0.02892691, -0.01291359]]), np.array([[0.04180323, 0.03254542, -0.56899236], [0.03254542, -0.5865448, -0.33459247], [-0.56899236, -0.33459247, 0.23441378]]), np.array([[0.04180323, 0.03254542, -0.56899236], [0.03254542, -0.5865448, -0.33459247], [-0.56899236, -0.33459247, 0.23441378]]), np.array([[0.02685116, 0.25494821], [0.25494821, -0.33546046]]), np.array([[0.02685116, 0.25494821], [0.25494821, -0.33546046]]), np.array([[-0.04041477, 0.36934457], [0.36934457, 0.0278976]]), np.array([[-0.04041477, 0.36934457], [0.36934457, 0.0278976]]), np.array([[-0.11861509, 0.37524325], [0.37524325, 0.38311618]]), np.array([[-0.11861509, 0.37524325], [0.37524325, 0.38311618]]), np.array([[-0.18418421]]), np.array([[-0.18418421]]), np.array([[-0.22672867]]), np.array([[-0.22672867]]), np.array([[-0.24393745]]), np.array([[-0.24393745]]) ]) self.blk_partition = [(4, 4), (3, 3), (3, 3), (3, 3), (3, 3), (2, 2), (2, 2), (2, 2), (2, 2), (2, 2), (2, 2), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1)]
def testBlkDiagMatrixSolve(self): result = np.array([[ -2108.56909528, -1599.04367287, 12983.7980645, 4907.54823872, 11596.60358098 ], [ -13261.71782311, 19064.68118668, -7058.51163424, -23505.65277468, 10718.98836096 ], [ 509.34615374, -6346.73823956, 12742.41838169, 26050.66900347, -11082.31307888 ], [ -3307.59390503, -2868.45197933, -9056.27254017, -13121.74603406, 10435.30033164 ], [ 5194.61155267, 4738.83799474, 1275.99271119, 5695.36574982, 317.72274105 ], [ -7192.22592633, -7706.51971094, -2706.50848048, -5226.00224303, -10609.21894342 ], [ 2106.79729152, 720.663365, 477.29928858, 2865.35110109, 2523.12141081 ], [ 3293.89400066, 8296.32959773, 6274.69897306, -6876.89810121, -5862.81117304 ], [ -5037.67048311, 4648.20122872, -6478.90991549, 5611.52398572, 5231.26195727 ], [ 319.85059444, -13526.13244267, 5392.6728323, -5960.86618271, -9989.22959649 ], [ -6810.8349795, -5587.85203171, -4557.56983224, -21895.30335831, -14766.02390764 ], [ -1487.82495778, -17267.75186366, 2518.13873842, 16122.34779131, -10824.59882957 ], [ 1696.60328706, 7524.33425367, -4412.35449613, -8159.19336023, 3283.61842699 ], [ 16533.21857068, 1048.77054126, 270.18096236, 3130.66029662, 3802.78276484 ], [ 3987.29419815, 14592.92355119, -10535.19114199, -11751.50898025, 10280.91948022 ], [ 8221.61307636, 4963.62327863, 15482.47214868, 9215.02309094, 2026.10702896 ], [ 10605.58662119, 4939.08507375, -6298.18788412, 2199.78159731, 6348.7733857 ], [ 8939.66366506, 5979.70412996, 3984.91381065, 5566.93943721, -1421.24116192 ], [ -21606.51783477, 4822.39145187, 6685.88851131, 15766.49583437, -12707.40988432 ], [ -4230.44592724, 5420.21451676, -235.85955642, 18100.20123655, -4768.83642194 ], [ 9450.93104008, 8900.03420762, 5202.48555337, -201.77556792, 5466.57970732 ], [ 11583.03440042, 75.37078542, 19981.80654578, -358.04535766, -6139.4067836 ], [ 12333.72389213, 1438.96251241, -8437.87301414, 17219.37822955, 11248.37259986 ], [ 2193.30045686, 7010.7452714, 7716.40327545, -7755.61096007, 4158.2657834 ], [ -16260.99699839, -24157.30011252, 8957.72261107, 7178.33063461, 8681.06123862 ], [ 9062.40740653, 14709.69068774, 5430.20220187, -6558.47236417, -7115.41480134 ], [ 7468.93152378, -11727.69088497, 13044.41696129, -1429.80060688, 615.43370603 ], [ 5660.19210206, 758.33179619, 457.57835518, 8210.45409543, -6029.85682174 ], [ -11618.39974191, -4792.87662117, 7891.4438672, -13769.59940549, 21193.75984068 ], [ -10062.76974695, -12184.86042594, 3139.30111484, 5233.94758207, 5060.12404909 ], [ 20338.94250519, -5108.5476594, -7916.199259, -13155.55120958, -2631.98451114 ], [ -8575.21386665, -3075.44542941, 6705.35386167, 1366.04418311, -3451.27892975 ], [ 11471.10856558, -18990.29443441, -5960.30374281, 10147.12732058, 2454.46580692 ], [ 23451.62021758, 9401.91573045, -5248.02166459, -7872.80360867, 4596.97842653 ]]) sn_matrix = [ np.array([[ 1.29245300e-08, -6.49163277e-10, 2.45070400e-09, 1.53410891e-10 ], [ -6.49163277e-10, 1.20780198e-08, 3.75068433e-09, -2.53453097e-09 ], [ 2.45070400e-09, 3.75068433e-09, 1.07092249e-08, 1.63163421e-09 ], [ 1.53410891e-10, -2.53453097e-09, 1.63163421e-09, 1.57226938e-08 ]]), np.array([[1.29101793e-08, 8.87468230e-09, 3.67875235e-09], [8.87468230e-09, 5.05913176e-08, 2.61651494e-08], [3.67875235e-09, 2.61651494e-08, 3.12014975e-08]]), np.array([[1.29101793e-08, 8.87468230e-09, 3.67875235e-09], [8.87468230e-09, 5.05913176e-08, 2.61651494e-08], [3.67875235e-09, 2.61651494e-08, 3.12014975e-08]]), np.array([[1.09084572e-08, 6.16238273e-12, -1.67985227e-09], [6.16238273e-12, 9.71880814e-09, 6.56599364e-09], [-1.67985227e-09, 6.56599364e-09, 1.98674901e-08]]), np.array([[1.09084572e-08, 6.16238273e-12, -1.67985227e-09], [6.16238273e-12, 9.71880814e-09, 6.56599364e-09], [-1.67985227e-09, 6.56599364e-09, 1.98674901e-08]]), np.array([[1.44374266e-08, -4.43138611e-09], [-4.43138611e-09, 1.13592375e-08]]), np.array([[1.44374266e-08, -4.43138611e-09], [-4.43138611e-09, 1.13592375e-08]]), np.array([[1.01359586e-08, 1.48988673e-09], [1.48988673e-09, 8.98955124e-09]]), np.array([[1.01359586e-08, 1.48988673e-09], [1.48988673e-09, 8.98955124e-09]]), np.array([[1.12133339e-08, 1.87797581e-09], [1.87797581e-09, 9.22078294e-09]]), np.array([[1.12133339e-08, 1.87797581e-09], [1.87797581e-09, 9.22078294e-09]]), np.array([[9.17068059e-09]]), np.array([[9.17068059e-09]]), np.array([[8.62262996e-09]]), np.array([[8.62262996e-09]]), np.array([[8.80222582e-09]]), np.array([[8.80222582e-09]]) ] coeff = np.array( [[ -1.79024086e-05, -4.90370074e-05, 2.02230179e-04, 1.40516221e-04, 1.17363697e-04 ], [ -1.48512889e-04, 2.14767206e-04, -2.29352560e-05, -1.56122232e-04, 5.39212146e-05 ], [ -5.48500768e-05, -5.06209331e-06, 1.27030097e-04, 1.81437248e-04, -3.30330057e-05 ], [ -1.78844625e-05, -1.04020683e-04, -1.01716162e-04, -1.03475356e-04, 1.20600188e-04 ], [ 1.09850319e-05, -4.56252384e-06, -5.79024234e-06, 3.76900005e-05, -8.07696512e-05 ], [ -2.62638993e-04, -3.28971040e-04, -1.13113193e-04, -1.38873438e-04, -4.67896828e-04 ], [ -1.03340746e-04, -1.61723452e-04, -5.12296850e-05, -2.63840441e-05, -1.97697809e-04 ], [ -1.00631180e-06, 9.85991203e-05, 4.33475295e-05, -6.09100454e-05, -6.60120574e-05 ], [ -2.17261186e-04, -4.51273621e-05, -1.30990539e-04, 6.68971521e-05, -4.87438362e-05 ], [ -1.09714163e-04, -2.69894566e-04, 2.18208856e-05, -6.44599929e-05, -1.96370002e-04 ], [ -7.71549133e-05, -7.37010252e-05, -4.22884340e-05, -2.25038388e-04, -1.66657239e-04 ], [ -3.36196989e-06, -1.18451671e-04, -4.52626978e-06, 1.02981866e-04, -8.37329754e-05 ], [ 3.53793964e-05, 4.54964533e-05, -6.34722823e-05, -1.94625852e-05, 1.89677483e-05 ], [ 1.66565383e-04, 3.19224191e-06, -2.31259304e-05, 1.85983791e-05, 3.81422875e-05 ], [ 9.28366905e-05, 1.74423406e-04, -7.30022809e-07, -5.36855858e-05, 1.13245124e-04 ], [0.00016175, 0.00019267, 0.00023797, 0.00010066, 0.00010137], [ 1.13502277e-04, 4.48093004e-05, -1.08588317e-04, 7.08992725e-06, 9.79580181e-05 ], [ 5.45503135e-05, 4.60378864e-05, 7.31752847e-05, 5.34881056e-05, -4.42780821e-05 ], [ -2.93195776e-04, 4.56038593e-05, 9.75722094e-05, 1.47418646e-04, -1.62329742e-04 ], [ 4.76921830e-05, 4.01996255e-05, -3.23069382e-05, 1.35737054e-04, 2.14109414e-06 ], [ 1.13051655e-04, 9.03226722e-05, 8.25028066e-05, -2.57863583e-06, 4.62620049e-05 ], [ 1.18207098e-04, 1.39375924e-05, 1.87378588e-04, -3.51928983e-06, -4.70459273e-05 ], [ 1.28281884e-04, 2.50304808e-05, -7.40293647e-05, 1.62979923e-04, 1.20208384e-04 ], [ 3.80926384e-05, 6.51673450e-05, 5.67955276e-05, -4.40645390e-05, 5.41397444e-05 ], [ -1.65321007e-04, -2.43259429e-04, 1.10643723e-04, 6.81763658e-05, 8.39810614e-05 ], [ 5.30247326e-05, 9.02680397e-05, 6.68931022e-05, -4.69935188e-05, -4.93068724e-05 ], [ 9.43813269e-05, -1.30082385e-04, 1.47130724e-04, -6.13797435e-07, -4.42286161e-06 ], [ 6.62178755e-05, -1.50319069e-05, 2.87163302e-05, 7.30216841e-05, -5.44442313e-05 ], [ -1.06548633e-04, -4.39539406e-05, 7.23699111e-05, -1.26276598e-04, 1.94361202e-04 ], [ -9.22824472e-05, -1.11743463e-04, 2.87895278e-05, 4.79988615e-05, 4.64047814e-05 ], [ 1.75375175e-04, -4.40491161e-05, -6.82584569e-05, -1.13435450e-04, -2.26946285e-05 ], [ -7.39408960e-05, -2.65184279e-05, 5.78177851e-05, 1.17788935e-05, -2.97591011e-05 ], [ 1.00971288e-04, -1.67156860e-04, -5.24639395e-05, 8.93173061e-05, 2.16047623e-05 ], [ 2.06426457e-04, 8.27577854e-05, -4.61942718e-05, -6.92981952e-05, 4.04636422e-05 ]]) A = BlkDiagMatrix.from_list(sn_matrix) coeff_est = A.solve(coeff) self.allallfunc(result, coeff_est)
def get_covar(self, coeffs, ctf_fb=None, ctf_idx=None, mean_coeff=None, do_refl=True, noise_var=1, covar_est_opt=None): """ Calculate the covariance matrix from the expansion coefficients and CTF information. :param coeffs: A coefficient vector (or an array of coefficient vectors) to be calculated. :param ctf_fb: The CFT functions in the FB expansion. :param ctf_idx: An array of the CFT function indices for all 2D images. If ctf_fb or ctf_idx is None, the identity filter will be applied. :param mean_coeff: The mean value vector from all images. :param noise_var: The estimated variance of noise. The value should be zero for `coeffs` from clean images of simulation data. :param covar_est_opt: The optimization parameter list for obtaining the Cov2D matrix. :return: The basis coefficients of the covariance matrix in the form of cell array representing a block diagonal matrix. These block diagonal matrices are implemented as BlkDiagMatrix instances. The covariance is calculated from the images represented by the coeffs array, along with all possible rotations and reflections. As a result, the computed covariance matrix is invariant to both reflection and rotation. The effect of the filters in ctf_fb are accounted for and inverted to yield a covariance estimate of the unfiltered images. """ if coeffs.size == 0: raise RuntimeError('The coefficients need to be calculated!') if (ctf_fb is None) or (ctf_idx is None): ctf_idx = np.zeros(coeffs.shape[1], dtype=int) ctf_fb = [BlkDiagMatrix.eye_like(RadialCTFFilter().fb_mat(self.basis))] def identity(x): return x default_est_opt = {'shrinker': 'None', 'verbose': 0, 'max_iter': 250, 'iter_callback': [], 'store_iterates': False, 'rel_tolerance': 1e-12, 'precision': 'float64', 'preconditioner': identity} covar_est_opt = fill_struct(covar_est_opt, default_est_opt) if mean_coeff is None: mean_coeff = self.get_mean(coeffs, ctf_fb, ctf_idx) b_coeff = BlkDiagMatrix.zeros_like(ctf_fb[0]) b_noise = BlkDiagMatrix.zeros_like(ctf_fb[0]) A = [] for k in range(0, len(ctf_fb)): A.append(BlkDiagMatrix.zeros_like(ctf_fb[0])) M = BlkDiagMatrix.zeros_like(ctf_fb[0]) for k in np.unique(ctf_idx[:]): coeff_k = coeffs[:, ctf_idx == k] weight = np.size(coeff_k, 1)/np.size(coeffs, 1) ctf_fb_k = ctf_fb[k] ctf_fb_k_t = ctf_fb_k.T mean_coeff_k = ctf_fb_k.apply(mean_coeff) covar_coeff_k = self._get_covar(coeff_k, mean_coeff_k) b_coeff += weight * (ctf_fb_k_t @ covar_coeff_k @ ctf_fb_k) ctf_fb_k_sq = ctf_fb_k_t @ ctf_fb_k b_noise += weight * ctf_fb_k_sq A[k] = np.sqrt(weight) * ctf_fb_k_sq M += A[k] if covar_est_opt['shrinker'] == 'None': b = b_coeff - noise_var * b_noise else: b = self.shrink_covar_backward(b_coeff, b_noise, np.size(coeffs, 1), noise_var, covar_est_opt['shrinker']) cg_opt = covar_est_opt covar_coeff = BlkDiagMatrix.zeros_like(ctf_fb[0]) def precond_fun(S, x): p = np.size(S, 0) ensure(np.size(x) == p*p, 'The sizes of S and x are not consistent.') x = m_reshape(x, (p, p)) y = S @ x @ S y = m_reshape(y, (p ** 2,)) return y def apply(A, x): p = np.size(A[0], 0) x = m_reshape(x, (p, p)) y = np.zeros_like(x) for k in range(0, len(A)): y = y + A[k] @ x @ A[k].T y = m_reshape(y, (p ** 2,)) return y for ell in range(0, len(b)): A_ell = [] for k in range(0, len(A)): A_ell.append(A[k][ell]) p = np.size(A_ell[0], 0) b_ell = m_reshape(b[ell], (p ** 2,)) S = inv(M[ell]) cg_opt["preconditioner"] = lambda x: precond_fun(S, x) covar_coeff_ell, _, _ = conj_grad(lambda x: apply(A_ell, x), b_ell, cg_opt) covar_coeff[ell] = m_reshape(covar_coeff_ell, (p, p)) return covar_coeff