def test_trivial_cholesky(self): # Known starting point of SI_nn = <psit_m|S+alpha*I|psit_n> I_nn = np.eye(*self.S0_nn.shape) alpha = 1e-3 # shift eigenvalues away from zero SI_nn = self.S0_nn + alpha * I_nn # Try Cholesky decomposition SI_nn = L_nn * L_nn^dag L_nn = np.linalg.cholesky(SI_nn) # |psit_n> -> C_nn |psit_n> , C_nn^(-1) = L_nn^dag # <psit_m|SI|psit_n> -> <psit_m|C_nn^dag SI C_nn|psit_n> = diag(W_n) C_nn = np.linalg.inv(L_nn.T.conj()) # Set up Hermitian overlap operator: S = lambda x: x dS = lambda a, P_ni: np.dot(P_ni, self.setups[a].dO_ii) nblocks = self.get_optimal_number_of_blocks(self.blocking) overlap = MatrixOperator(self.ksl, nblocks, self.async, True) self.psit_nG = overlap.matrix_multiply(C_nn.T.copy(), self.psit_nG, self.P_ani) D_nn = overlap.calculate_matrix_elements(self.psit_nG, \ self.P_ani, S, dS).T.copy() # transpose to get <psit_m|A|psit_n> tri2full(D_nn, 'U') # upper to lower... if self.bd.comm.rank == 0: self.gd.comm.broadcast(D_nn, 0) self.bd.comm.broadcast(D_nn, 0) if memstats: self.mem_test = record_memory() # D_nn = C_nn^dag * S_nn * C_nn = I_nn - alpha * C_nn^dag * C_nn D0_nn = I_nn - alpha * np.dot(C_nn.T.conj(), C_nn) self.check_and_plot(D_nn, D0_nn, 6, 'trivial,cholesky') #XXX precision
def test_multiply_randomized(self): # Known starting point of S_nn = <psit_m|S|psit_n> S_nn = self.S0_nn if self.dtype == complex: C_nn = np.random.uniform(size=self.nbands**2) * \ np.exp(1j*np.random.uniform(0,2*np.pi,size=self.nbands**2)) else: C_nn = np.random.normal(size=self.nbands**2) C_nn = C_nn.reshape((self.nbands,self.nbands)) / np.linalg.norm(C_nn,2) world.broadcast(C_nn, 0) # Set up Hermitian overlap operator: S = lambda x: x dS = lambda a, P_ni: np.dot(P_ni, self.setups[a].dO_ii) nblocks = self.get_optimal_number_of_blocks(self.blocking) overlap = MatrixOperator(self.ksl, nblocks, self.async, True) self.psit_nG = overlap.matrix_multiply(C_nn.T.copy(), self.psit_nG, self.P_ani) D_nn = overlap.calculate_matrix_elements(self.psit_nG, \ self.P_ani, S, dS).T.copy() # transpose to get <psit_m|A|psit_n> tri2full(D_nn, 'U') # upper to lower... if self.bd.comm.rank == 0: self.gd.comm.broadcast(D_nn, 0) self.bd.comm.broadcast(D_nn, 0) if memstats: self.mem_test = record_memory() # D_nn = C_nn^dag * S_nn * C_nn D0_nn = np.dot(C_nn.T.conj(), np.dot(S_nn, C_nn)) self.check_and_plot(D_nn, D0_nn, 9, 'multiply,randomized')
def test_trivial_diagonalize(self): # Known starting point of S_nn = <psit_m|S|psit_n> S_nn = self.S0_nn # Eigenvector decomposition S_nn = V_nn * W_nn * V_nn^dag # Utilize the fact that they are analytically known (cf. Maple) band_indices = np.arange(self.nbands) V_nn = np.eye(self.nbands).astype(self.dtype) if self.dtype == complex: V_nn[1:,1] = np.conj(self.gamma)**band_indices[1:] * band_indices[1:]**0.5 V_nn[1,2:] = -self.gamma**band_indices[1:-1] * band_indices[2:]**0.5 else: V_nn[2:,1] = band_indices[2:]**0.5 V_nn[1,2:] = -band_indices[2:]**0.5 W_n = np.zeros(self.nbands).astype(self.dtype) W_n[1] = (1. + self.Qtotal) * self.nbands * (self.nbands - 1) / 2. # Find the inverse basis Vinv_nn = np.linalg.inv(V_nn) # Test analytical eigenvectors for consistency against analytical S_nn D_nn = np.dot(Vinv_nn, np.dot(S_nn, V_nn)) self.assertAlmostEqual(np.abs(D_nn.diagonal()-W_n).max(), 0, 8) self.assertAlmostEqual(np.abs(np.tril(D_nn, -1)).max(), 0, 4) self.assertAlmostEqual(np.abs(np.triu(D_nn, 1)).max(), 0, 4) del Vinv_nn, D_nn # Perform Gram Schmidt orthonormalization for diagonalization # |psit_n> -> C_nn |psit_n>, using orthonormalized basis Q_nn # <psit_m|S|psit_n> -> <psit_m|C_nn^dag S C_nn|psit_n> = diag(W_n) # using S_nn = V_nn * W_nn * V_nn^(-1) = Q_nn * W_nn * Q_nn^dag C_nn = V_nn.copy() gram_schmidt(C_nn) self.assertAlmostEqual(np.abs(np.dot(C_nn.T.conj(), C_nn) \ - np.eye(self.nbands)).max(), 0, 6) # Set up Hermitian overlap operator: S = lambda x: x dS = lambda a, P_ni: np.dot(P_ni, self.setups[a].dO_ii) nblocks = self.get_optimal_number_of_blocks(self.blocking) overlap = MatrixOperator(self.ksl, nblocks, self.async, True) self.psit_nG = overlap.matrix_multiply(C_nn.T.copy(), self.psit_nG, self.P_ani) D_nn = overlap.calculate_matrix_elements(self.psit_nG, \ self.P_ani, S, dS).T.copy() # transpose to get <psit_m|A|psit_n> tri2full(D_nn, 'U') # upper to lower... if self.bd.comm.rank == 0: self.gd.comm.broadcast(D_nn, 0) self.bd.comm.broadcast(D_nn, 0) if memstats: self.mem_test = record_memory() # D_nn = C_nn^dag * S_nn * C_nn = W_n since Q_nn^dag = Q_nn^(-1) D0_nn = np.dot(C_nn.T.conj(), np.dot(S_nn, C_nn)) self.assertAlmostEqual(np.abs(D0_nn-np.diag(W_n)).max(), 0, 9) self.check_and_plot(D_nn, D0_nn, 9, 'trivial,diagonalize')
def test_multiply_nonhermitian(self): alpha = np.random.normal(size=1).astype(self.dtype) if self.dtype == complex: alpha += 1j*np.random.normal(size=1) world.broadcast(alpha, 0) # Known starting point of S_nn = <psit_m|S|psit_n> S_NN = alpha*self.S0_nn if self.dtype == complex: C_NN = np.random.uniform(size=self.nbands**2) * \ np.exp(1j*np.random.uniform(0,2*np.pi,size=self.nbands**2)) else: C_NN = np.random.normal(size=self.nbands**2) C_NN = C_NN.reshape((self.nbands,self.nbands)) / np.linalg.norm(C_NN,2) world.broadcast(C_NN, 0) # Set up Hermitian overlap operator: S = lambda x: alpha*x dS = lambda a, P_ni: np.dot(alpha*P_ni, self.setups[a].dO_ii) nblocks = self.get_optimal_number_of_blocks(self.blocking) overlap = MatrixOperator(self.ksl, nblocks, self.async, False) if 0: #XXX non-hermitian case so Nn2nn not just uplo='L' but rather 'G' blockcomm = self.ksl.nndescriptor.blacsgrid.comm self.ksl.Nn2nn = Redistributor(blockcomm, self.ksl.Nndescriptor, self.ksl.nndescriptor) if self.bd.comm.rank == 0 and self.gd.comm.rank == 0: assert C_NN.shape == (self.bd.nbands,) * 2 tmp_NN = C_NN.T.copy() # C -> Fortran indexing else: tmp_NN = self.ksl.nndescriptor.as_serial().empty(dtype=C_NN.dtype) C_nn = self.ksl.nndescriptor.distribute_from_master(tmp_NN) self.psit_nG = overlap.matrix_multiply(C_nn, self.psit_nG, self.P_ani) D_nn = overlap.calculate_matrix_elements(self.psit_nG, self.P_ani, S, dS) if memstats: self.mem_test = record_memory() D_NN = self.ksl.nndescriptor.collect_on_master(D_nn) if self.bd.comm.rank == 0 and self.gd.comm.rank == 0: assert D_NN.shape == (self.bd.nbands,) * 2 D_NN = D_NN.T.copy() # Fortran -> C indexing else: assert D_NN.nbytes == 0 D_NN = np.empty((self.bd.nbands,) * 2, dtype=D_NN.dtype) if self.bd.comm.rank == 0: self.gd.comm.broadcast(D_NN, 0) self.bd.comm.broadcast(D_NN, 0) # D_nn = C_nn^dag * S_nn * C_nn D0_NN = np.dot(C_NN.T.conj(), np.dot(S_NN, C_NN)) self.check_and_plot(D_NN, D0_NN, 9, 'multiply,nonhermitian')
def test_multiply_randomized(self): # Known starting point of S_nn = <psit_m|S|psit_n> S_NN = self.S0_nn if self.dtype == complex: C_NN = np.random.uniform(size=self.nbands**2) * \ np.exp(1j*np.random.uniform(0,2*np.pi,size=self.nbands**2)) else: C_NN = np.random.normal(size=self.nbands**2) C_NN = C_NN.reshape((self.nbands,self.nbands)) / np.linalg.norm(C_NN,2) world.broadcast(C_NN, 0) # Set up Hermitian overlap operator: S = lambda x: x dS = lambda a, P_ni: np.dot(P_ni, self.setups[a].dO_ii) nblocks = self.get_optimal_number_of_blocks(self.blocking) overlap = MatrixOperator(self.ksl, nblocks, self.async, True) if self.bd.comm.rank == 0 and self.gd.comm.rank == 0: assert C_NN.shape == (self.bd.nbands,) * 2 tmp_NN = C_NN.T.copy() # C -> Fortran indexing else: tmp_NN = self.ksl.nndescriptor.as_serial().empty(dtype=C_NN.dtype) C_nn = self.ksl.nndescriptor.distribute_from_master(tmp_NN) self.psit_nG = overlap.matrix_multiply(C_nn, self.psit_nG, self.P_ani) D_nn = overlap.calculate_matrix_elements(self.psit_nG, self.P_ani, S, dS) if memstats: self.mem_test = record_memory() D_NN = self.ksl.nndescriptor.collect_on_master(D_nn) if self.bd.comm.rank == 0 and self.gd.comm.rank == 0: assert D_NN.shape == (self.bd.nbands,) * 2 D_NN = D_NN.T.copy() # Fortran -> C indexing tri2full(D_NN, 'U') # upper to lower... else: assert D_NN.nbytes == 0 D_NN = np.empty((self.bd.nbands,) * 2, dtype=D_NN.dtype) if self.bd.comm.rank == 0: self.gd.comm.broadcast(D_NN, 0) self.bd.comm.broadcast(D_NN, 0) # D_nn = C_nn^dag * S_nn * C_nn D0_NN = np.dot(C_NN.T.conj(), np.dot(S_NN, C_NN)) self.check_and_plot(D_NN, D0_NN, 9, 'multiply,randomized')
def run(psit_mG): overlap = MatrixOperator(ksl, K) if 0: overlap.work1_xG = work1_xG overlap.work2_xG = work2_xG #S_nn = np.empty((N, N)) def S(x): return x dS_aii = {0: np.ones((2, 2)) * 0.123, 1: np.ones((3, 3)) * 0.321} def dS(a, P_ni): return np.dot(P_ni, dS_aii[a]) S_nn = overlap.calculate_matrix_elements(psit_mG, P_ani, S, dS) t1 = time() if world.rank == 0: print(S_nn.round(5)) inverse_cholesky(S_nn) C_nn = S_nn t2 = time() if world.rank == 0: print('Cholesky Time %f' % (t2 - t1)) # Distribute matrix: world.broadcast(C_nn, 0) psit_mG = overlap.matrix_multiply(C_nn, psit_mG, P_ani) if world.rank == 0: print('Made it past matrix multiply') # Check: S_nn = overlap.calculate_matrix_elements(psit_mG, P_ani, S, dS) assert not (P_ani[0] - psit_mG[:, :2, 0, 0]).round(10).any() assert not (P_ani[1] - psit_mG[:, -1, -1, -3:]).round(10).any() if world.rank == 0: for n in range(N): assert abs(S_nn[n, n] - 1.0) < 1e-10 assert not S_nn[n + 1:, n].round(10).any() return psit_mG
def run(psit_mG): overlap = MatrixOperator(ksl, J) def H(psit_xG): Htpsit_xG = np.empty_like(psit_xG) kin(psit_xG, Htpsit_xG) for psit_G, y_G in zip(psit_xG, Htpsit_xG): y_G += vt_G * psit_G return Htpsit_xG dH_aii = {0: np.ones((2, 2)) * 0.123, 1: np.ones((3, 3)) * 0.321} def dH(a, P_ni): return np.dot(P_ni, dH_aii[a]) H_nn = overlap.calculate_matrix_elements(psit_mG, P_ani, H, dH) t1 = time() if world.rank == 0: eps_n, H_nn = np.linalg.eigh(H_nn) H_nn = np.ascontiguousarray(H_nn.T) t2 = time() if world.rank == 0: print('Diagonalization Time %f' % (t2 - t1)) print(eps_n) # Distribute matrix: world.broadcast(H_nn, 0) psit_mG = overlap.matrix_multiply(H_nn, psit_mG, P_ani) if world.rank == 0: print('Made it past matrix multiply') # Check: assert not (P_ani[0] - psit_mG[:, :2, 0, 0]).round(10).any() assert not (P_ani[1] - psit_mG[:, -1, -1, -3:]).round(10).any() H_nn = overlap.calculate_matrix_elements(psit_mG, P_ani, H, dH) if world.rank == 0: for n in range(N): assert abs(H_nn[n, n] - eps_n[n]) < 2e-8 assert not H_nn[n + 1:, n].round(8).any() return psit_mG
def test_trivial_diagonalize(self): #XXX XXX XXX # Known starting point of S_nn = <psit_m|S|psit_n> S_nn = self.S0_nn # Eigenvector decomposition S_nn = V_nn * W_nn * V_nn^dag # Utilize the fact that they are analytically known (cf. Maple) W_n = np.zeros(self.nbands).astype(self.dtype) W_n[1] = (1. + self.Qtotal) * self.nbands * (self.nbands - 1) / 2. # Set up Hermitian overlap operator: S = lambda x: x dS = lambda a, P_ni: np.dot(P_ni, self.setups[a].dO_ii) nblocks = self.get_optimal_number_of_blocks(self.blocking) overlap = MatrixOperator(self.ksl, nblocks, self.async, True) S_nn = overlap.calculate_matrix_elements(self.psit_nG, self.P_ani, S, dS) eps_N = self.bd.empty(global_array=True) # XXX dtype? C_nn = self.ksl.nndescriptor.empty(dtype=S_nn.dtype) self.ksl.nndescriptor.diagonalize_dc(S_nn, C_nn, eps_N, 'L') self.assertAlmostEqual(np.abs(np.sort(eps_N)-np.sort(W_n)).max(), 0, 9) #eps_n = self.bd.empty() #self.bd.distribute(eps_N, eps_n) # XXX only blocked groups, right? # Rotate wavefunctions to diagonalize the overlap self.psit_nG = overlap.matrix_multiply(C_nn, self.psit_nG, self.P_ani) # Recaulculate the overlap matrix, which should now be diagonal D_nn = overlap.calculate_matrix_elements(self.psit_nG, self.P_ani, S, dS) D_NN = self.ksl.nndescriptor.collect_on_master(D_nn) if self.bd.comm.rank == 0 and self.gd.comm.rank == 0: assert D_NN.shape == (self.bd.nbands,) * 2 D_NN = D_NN.T.copy() # Fortran -> C indexing tri2full(D_NN, 'U') # upper to lower... else: assert D_NN.nbytes == 0 D_NN = np.empty((self.bd.nbands,) * 2, dtype=D_NN.dtype) if self.bd.comm.rank == 0: self.gd.comm.broadcast(D_NN, 0) self.bd.comm.broadcast(D_NN, 0) D0_NN = np.diag(eps_N) self.check_and_plot(D_NN, D0_NN, 9, 'trivial,diagonalize')
def run(psit_mG): overlap = MatrixOperator(ksl, K) if 0: overlap.work1_xG = work1_xG overlap.work2_xG = work2_xG #S_nn = np.empty((N, N)) def S(x): return x dS_aii = {0: np.ones((2, 2)) * 0.123, 1: np.ones((3, 3)) * 0.321} def dS(a, P_ni): return np.dot(P_ni, dS_aii[a]) S_nn = overlap.calculate_matrix_elements(psit_mG, P_ani, S, dS) t1 = time() if world.rank == 0: print S_nn.round(5) inverse_cholesky(S_nn) C_nn = S_nn t2 = time() if world.rank == 0: print 'Cholesky Time %f' % (t2-t1) # Distribute matrix: world.broadcast(C_nn, 0) psit_mG = overlap.matrix_multiply(C_nn, psit_mG, P_ani) if world.rank == 0: print 'Made it past matrix multiply' # Check: S_nn = overlap.calculate_matrix_elements(psit_mG, P_ani, S, dS) assert not(P_ani[0] - psit_mG[:, :2, 0, 0]).round(10).any() assert not(P_ani[1] - psit_mG[:, -1, -1, -3:]).round(10).any() if world.rank == 0: for n in range(N): assert abs(S_nn[n, n] - 1.0) < 1e-10 assert not S_nn[n + 1:, n].round(10).any() return psit_mG
def run(psit_mG): overlap = MatrixOperator(ksl, J) def H(psit_xG): Htpsit_xG = np.empty_like(psit_xG) kin(psit_xG, Htpsit_xG) for psit_G, y_G in zip(psit_xG, Htpsit_xG): y_G += vt_G * psit_G return Htpsit_xG dH_aii = {0: np.ones((2, 2)) * 0.123, 1: np.ones((3, 3)) * 0.321} def dH(a, P_ni): return np.dot(P_ni, dH_aii[a]) H_nn = overlap.calculate_matrix_elements(psit_mG, P_ani, H, dH) t1 = time() if world.rank == 0: eps_n, H_nn = np.linalg.eigh(H_nn) H_nn = np.ascontiguousarray(H_nn.T) t2 = time() if world.rank == 0: print('Diagonalization Time %f' % (t2-t1)) print(eps_n) # Distribute matrix: world.broadcast(H_nn, 0) psit_mG = overlap.matrix_multiply(H_nn, psit_mG, P_ani) if world.rank == 0: print('Made it past matrix multiply') # Check: assert not(P_ani[0] - psit_mG[:, :2, 0, 0]).round(10).any() assert not(P_ani[1] - psit_mG[:, -1, -1, -3:]).round(10).any() H_nn = overlap.calculate_matrix_elements(psit_mG, P_ani, H, dH) if world.rank == 0: for n in range(N): assert abs(H_nn[n, n] - eps_n[n]) < 1.5e-8 assert not H_nn[n + 1:, n].round(8).any() return psit_mG
def test_trivial_cholesky(self): # Set up Hermitian overlap operator: S = lambda x: x dS = lambda a, P_ni: np.dot(P_ni, self.setups[a].dO_ii) nblocks = self.get_optimal_number_of_blocks(self.blocking) overlap = MatrixOperator(self.ksl, nblocks, self.async, True) S_nn = overlap.calculate_matrix_elements(self.psit_nG, self.P_ani, S, dS) # Known starting point of SI_nn = <psit_m|S+alpha*I|psit_n> I_nn = self.ksl.nndescriptor.empty(dtype=S_nn.dtype) scalapack_set(self.ksl.nndescriptor, I_nn, 0.0, 1.0, 'L') alpha = 1e-3 # shift eigenvalues away from zero C_nn = S_nn + alpha * I_nn self.ksl.nndescriptor.inverse_cholesky(C_nn, 'L') self.psit_nG = overlap.matrix_multiply(C_nn, self.psit_nG, self.P_ani) D_nn = overlap.calculate_matrix_elements(self.psit_nG, self.P_ani, S, dS) D_NN = self.ksl.nndescriptor.collect_on_master(D_nn) if self.bd.comm.rank == 0 and self.gd.comm.rank == 0: assert D_NN.shape == (self.bd.nbands,) * 2 D_NN = D_NN.T.copy() # Fortran -> C indexing tri2full(D_NN, 'U') # upper to lower.. else: assert D_NN.nbytes == 0 D_NN = np.empty((self.bd.nbands,) * 2, dtype=D_NN.dtype) if self.bd.comm.rank == 0: self.gd.comm.broadcast(D_NN, 0) self.bd.comm.broadcast(D_NN, 0) # D_NN = C_NN^dag * S_NN * C_NN = I_NN - alpha * C_NN^dag * C_NN I_NN = np.eye(self.bd.nbands) C0_NN = np.linalg.inv(np.linalg.cholesky(self.S0_nn + alpha*I_NN).T.conj()) D0_NN = I_NN - alpha * np.dot(C0_NN.T.conj(), C0_NN) self.check_and_plot(D_NN, D0_NN, 6, 'trivial,cholesky') #XXX precision