def test_einsum_misc(self): # The iterator had an issue with buffering this reduction a0 = numpy.ones((5, 12, 4, 2, 3)) b0 = numpy.ones((5, 12, 11)) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue(allclose(ctf.einsum('ijklm,ijn->', a1, b1), numpy.einsum('ijklm,ijn->', a0, b0))) self.assertTrue(allclose(ctf.einsum('ijklm,ijn,ijn->', a1, b1, b1), #numpy.einsum('ijklm,ijn,ijn->', a0, b0, b0))) numpy.einsum('ijklm,ijn->', a0, b0))) # inner loop implementation a0 = numpy.arange(1., 3.) b0 = numpy.arange(1., 5.).reshape(2, 2) c0 = numpy.arange(1., 9.).reshape(4, 2) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) c1 = ctf.astensor(c0) self.assertTrue(allclose(ctf.einsum('x,yx,zx->xzy', a1, b1, c1), numpy.einsum('x,yx,zx->xzy', a0, b0, c0))) a0 = numpy.random.normal(0, 1, (5, 5, 5, 5)) a1 = ctf.astensor(a0) self.assertTrue(allclose(ctf.einsum('aabb->ab', a1), numpy.einsum('aabb->ab', a0))) a0 = numpy.arange(25.).reshape(5, 5) a1 = ctf.astensor(a0) self.assertTrue(allclose(ctf.einsum('mi,mi,mi->m', a1, a1, a1), numpy.einsum('mi,mi,mi->m', a0, a0, a0)))
def test_einsum_sums(self): # outer(a,b) for n in range(1, 17): a0 = numpy.arange(3, dtype=numpy.double)+1 b0 = numpy.arange(n, dtype=numpy.double)+1 a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue(allclose(ctf.einsum("i,j", a1, b1), numpy.outer(a0, b0))) # matvec(a,b) / a.dot(b) where a is matrix, b is vector for n in range(1, 17): a0 = numpy.arange(4*n, dtype=numpy.double).reshape(n, 4) b0 = numpy.arange(n, dtype=numpy.double) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue(allclose(ctf.einsum("ji,j", a1, b1), numpy.dot(b0.T, a0))) self.assertTrue(allclose(ctf.einsum("ji,j->", a1, b1), numpy.dot(b0.T, a0).sum())) # matmat(a,b) / a.dot(b) where a is matrix, b is matrix for n in range(1, 17): a0 = numpy.arange(4*n, dtype=numpy.double).reshape(n, 4) b0 = numpy.arange(6*n, dtype=numpy.double).reshape(n, 6) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue(allclose(ctf.einsum("ji,jk", a1, b1), numpy.dot(a0.T, b0))) self.assertTrue(allclose(ctf.einsum("ji,jk->", a1, b1), numpy.dot(a0.T, b0).sum())) # matrix triple product (note this is not currently an efficient # way to multiply 3 matrices) a0 = numpy.arange(12.).reshape(3, 4) b0 = numpy.arange(20.).reshape(4, 5) c0 = numpy.arange(30.).reshape(5, 6) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) c1 = ctf.astensor(c0) self.assertTrue(allclose(ctf.einsum("ij,jk,kl", a1, b1, c1), numpy.einsum("ij,jk,kl", a0, b0, c0))) # tensordot(a, b) a0 = numpy.arange(27.).reshape(3, 3, 3) b0 = numpy.arange(27.).reshape(3, 3, 3) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue(allclose(ctf.einsum("ijk, jli -> kl", a1, b1), numpy.einsum("ijk, jli -> kl", a0, b0))) self.assertTrue(allclose(ctf.einsum("ijk, jli -> lk", a1, b1), numpy.einsum("ijk, jli -> lk", a0, b0))) self.assertTrue(allclose(ctf.einsum("ikj, jli -> kl", a1, b1), numpy.einsum("ikj, jli -> kl", a0, b0))) self.assertTrue(allclose(ctf.einsum("kij, lij -> lk", a1, b1), numpy.einsum("kij, lij -> lk", a0, b0)))
def test_noslice_sym_3d(self): n = 5 SY = ctf.SYM.SY NS = ctf.SYM.NS a0 = ctf.tensor([n,n,n], sym=[NS,SY,NS]) ctf.random.seed(1) a0.fill_random() mo = ctf.tensor([n,n]) mo.fill_random() dat = ctf.einsum('qrs,ri,sj->qij', a0, mo, mo) a1 = ctf.tensor(a0.shape, sym=[NS,NS,NS]) a1.i('ijkl') << a0.i('ijkl') ref = ctf.einsum('qrs,ri,sj->qij', a1, mo, mo) self.assertTrue(allclose(ref, dat))
def test_einsum_views(self): a0 = numpy.arange(27.).reshape(3,3,3) a1 = ctf.astensor(a0) self.assertTrue(allclose(ctf.einsum("jii->ij", a1), numpy.einsum("jii->ij", a0))) self.assertTrue(allclose(ctf.einsum("iii->i", a1), numpy.einsum("iii->i", a0))) self.assertTrue(allclose(ctf.einsum("iii", a1), numpy.einsum("iii", a0))) a0 = numpy.arange(6.) a1 = ctf.astensor(a0) self.assertTrue(allclose(ctf.einsum("i,i,i->i", a1, a1, a1), numpy.einsum("i,i,i->i", a0, a0, a0))) # swap axes a0 = numpy.arange(24.).reshape(4,3,2) a1 = ctf.astensor(a0) self.assertTrue(allclose(ctf.einsum("ijk->kji", a1), numpy.einsum("ijk->kji", a0)))
def test_einsum_mix_types(self): a0 = numpy.random.random((5, 1, 4, 2, 3)).astype(numpy.complex)+1j b0 = numpy.random.random((5, 1, 11)).astype(numpy.float32) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) c0 = numpy.einsum('ijklm,ijn->', a0, b0) c1 = ctf.einsum('ijklm,ijn->', a1, b1).to_nparray() self.assertTrue(c1.dtype == numpy.complex) self.assertTrue(allclose(c0, c1))
def test_sp_TTTP_mat(self): A = ctf.tensor((5, 1, 4, 2, 3), sp=True) A.fill_sp_random(0., 1., .2) u = ctf.random.random((5, 3)) v = ctf.random.random((1, 3)) w = ctf.random.random((4, 3)) x = ctf.random.random((2, 3)) y = ctf.random.random((3, 3)) ans = ctf.einsum("ijklm,ia,ja,ka,la,ma->ijklm", A, u, v, w, x, y) cans = ctf.TTTP(A, [u, v, w, x, y]) self.assertTrue(allclose(ans, cans))
def _einsum(self, expr, operands): inputs_indices = [operand.indices for operand in operands] inputs_shapes = [operand.tsr.shape for operand in operands] expanded_expr = indices_utils.expand_einsum(expr, inputs_indices, inputs_shapes) if expanded_expr is not None: result = ctf.einsum(expanded_expr.indices_string, *(operand.tsr for operand in operands)) if isinstance(result, ctf.tensor): newshape = expanded_expr.outputs[0].newshape(result.shape) return self.tensor(result).reshape(*newshape) else: return result else: result = ctf.einsum(expr.indices_string, *(operand.unwrap() for operand in operands)) if isinstance(result, ctf.tensor): newshape = expr.outputs[0].newshape(result.shape) return self.tensor(result).reshape(*newshape) else: return result
def test_einsum_views(self): a0 = numpy.arange(27.).reshape(3, 3, 3) a1 = ctf.astensor(a0) self.assertTrue( allclose(ctf.einsum("jii->ij", a1), numpy.einsum("jii->ij", a0))) self.assertTrue( allclose(ctf.einsum("iii->i", a1), numpy.einsum("iii->i", a0))) self.assertTrue( allclose(ctf.einsum("iii", a1), numpy.einsum("iii", a0))) a0 = numpy.arange(6.) a1 = ctf.astensor(a0) self.assertTrue( allclose(ctf.einsum("i,i,i->i", a1, a1, a1), numpy.einsum("i,i,i->i", a0, a0, a0))) # swap axes a0 = numpy.arange(24.).reshape(4, 3, 2) a1 = ctf.astensor(a0) self.assertTrue( allclose(ctf.einsum("ijk->kji", a1), numpy.einsum("ijk->kji", a0)))
def test_einsum_misc(self): # The iterator had an issue with buffering this reduction a0 = numpy.ones((5, 12, 4, 2, 3)) b0 = numpy.ones((5, 12, 11)) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue( allclose(ctf.einsum('ijklm,ijn->', a1, b1), numpy.einsum('ijklm,ijn->', a0, b0))) self.assertTrue( allclose( ctf.einsum('ijklm,ijn,ijn->', a1, b1, b1), #numpy.einsum('ijklm,ijn,ijn->', a0, b0, b0))) numpy.einsum('ijklm,ijn->', a0, b0))) # inner loop implementation a0 = numpy.arange(1., 3.) b0 = numpy.arange(1., 5.).reshape(2, 2) c0 = numpy.arange(1., 9.).reshape(4, 2) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) c1 = ctf.astensor(c0) self.assertTrue( allclose(ctf.einsum('x,yx,zx->xzy', a1, b1, c1), numpy.einsum('x,yx,zx->xzy', a0, b0, c0))) a0 = numpy.random.normal(0, 1, (5, 5, 5, 5)) a1 = ctf.astensor(a0) self.assertTrue( allclose(ctf.einsum('aabb->ab', a1), numpy.einsum('aabb->ab', a0))) a0 = numpy.arange(25.).reshape(5, 5) a1 = ctf.astensor(a0) self.assertTrue( allclose(ctf.einsum('mi,mi,mi->m', a1, a1, a1), numpy.einsum('mi,mi,mi->m', a0, a0, a0)))
def einsum(string, *args): if "..." in string: left = string.split(",") left[-1], right = left[-1].split("->") symbols = "".join( list( set([chr(i) for i in range(48, 127)]) - set( string.replace(".", "").replace(",", "").replace("->", "")) )) symbol_idx = 0 for i, (s, tsr) in enumerate(zip(left, args)): num_missing = tsr.ndim - len(s.replace("...", "")) left[i] = s.replace("...", symbols[symbol_idx:symbol_idx + num_missing]) symbol_idx += num_missing right = right.replace("...", symbols[:symbol_idx]) string = ",".join(left) + "->" + right return ctf.einsum(string, *args)
def einsum(string, *args): # doesn't check for invalid use of ellipses if '...' in string: left = string.split(',') left[-1], right = left[-1].split('->') symbols = ''.join( list( set(chr(i) for i in range(48, 127)) - set( string.replace('.', '').replace(',', '').replace('->', '')) )) symbol_idx = 0 for i, (s, tsr) in enumerate(zip(left, args)): num_missing = tsr.ndim - len(s.replace('...', '')) left[i] = s.replace("...", symbols[symbol_idx:symbol_idx + num_missing]) symbol_idx += num_missing right = right.replace('...', symbols[:symbol_idx]) string = ','.join(left) + '->' + right return ctf.einsum(string, *args)
def einsvd(einstr, A, r=None, transpose=True, compute_uv=True, full_matrices=True, mult_sv=False): str_a, str_uv = einstr.split("->") str_u, str_v = str_uv.split(",") U, S, Vh = A.i(str_a).svd(str_u, str_v, rank=r) if not compute_uv: return S if mult_sv: char_i = list(set(str_v) - set(str_a))[0] char_s = list(set(str_a) - set(str_v))[0] Vh = ctf.einsum( char_s + char_i + "," + str_v + "->" + str_v.replace(char_i, char_s), ctf.diag(S), Vh) if not transpose: Vh = Vh.T return U, S, Vh
def fast_gemm(A, B, n, b): A = A.reshape((n, b, n, b)).transpose([0, 2, 1, 3]) B = B.reshape((n, b, n, b)).transpose([0, 2, 1, 3]) sA = -ctf.einsum("ijxy->ixy", A) sB = -ctf.einsum("ijxy->ixy", B) for i in range(n): sA[i, :, :] += 2. * A[i, i, :, :] sB[i, :, :] += 2. * B[i, i, :, :] nn = n * (n - 1) * (n - 2) // 6 idPC = ctf.zeros((nn, n, n)) idPA = ctf.zeros((nn, n, n)) idPB = ctf.zeros((nn, n, n)) l = 0 for i in range(n): for j in range(i): for k in range(j): idPC[l, i, j] = 1 idPA[l, i, k] = 1 idPB[l, j, k] = 1 l += 1 oA = ctf.einsum("aij,ijxy->axy", idPC, A) + ctf.einsum( "aik,ikxy->axy", idPA, A) + ctf.einsum("ajk,jkxy->axy", idPB, A) oB = ctf.einsum("aij,ijxy->axy", idPC, B) + ctf.einsum( "aik,ikxy->axy", idPA, B) + ctf.einsum("ajk,jkxy->axy", idPB, B) P = ctf.einsum("axy,ayz->axz", oA, oB) Z = ctf.einsum("aij,axy->ijxy", idPC, P) + ctf.einsum( "aik,axy->ikxy", idPA, P) + ctf.einsum("ajk,axy->jkxy", idPB, P) Z += Z.transpose([1, 0, 2, 3]) U = ctf.einsum("ijxy,iyz->ijxz", A, sB) U += U.transpose([1, 0, 2, 3]) V = ctf.einsum("ixy,ijyz->ijxz", sA, B) V += V.transpose([1, 0, 2, 3]) W = ctf.einsum("ijxy,ijyz->ijxz", A, B) sW = ctf.einsum("ijxy->ixy", W) for i in range(n): sW[i, :, :] -= W[i, i, :, :] U[i, i, :, :] = 0 V[i, i, :, :] = 0 #print(U.transpose([0,2,1,3]).reshape((n*b,n*b))) #print(V.transpose([0,2,1,3]).reshape((n*b,n*b))) #print("W",W.transpose([0,2,1,3]).reshape((n*b,n*b))) #print((- sW.reshape((1,n,b,b)) - sW.reshape((n,1,b,b))).transpose([0,2,1,3]).reshape((n*b,n*b))) C = Z - (n - 8) * W + U + V - sW.reshape((1, n, b, b)) - sW.reshape( (n, 1, b, b)) for i in range(n): C[i, i, :, :] = 2 * sW[i, :, :] + 2 * W[i, i, :, :] return C.transpose([0, 2, 1, 3]).reshape((n * b, n * b))
def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init): wrld = ctf.comm() s = s_start nnz = float(s_start * s_start * s_start) * sp_init agg_s = [] agg_avg_times = [] agg_min_times = [] agg_max_times = [] agg_min_95 = [] agg_max_95 = [] if num_iter > 1: if ctf.comm().rank() == 0: print("Performing MTTKRP WARMUP with s =", s, "nnz =", nnz, "sp =", sp, "sp_init =", sp_init) T = ctf.tensor((s, s, s), sp=sp) T.fill_sp_random(-1., 1., float(nnz) / float(s * s * s)) U = ctf.random.random((s, R)) V = ctf.random.random((s, R)) W = ctf.random.random((s, R)) U = ctf.einsum("ijk,jr,kr->ir", T, V, W) V = ctf.einsum("ijk,ir,kr->jr", T, U, W) W = ctf.einsum("ijk,ir,jr->kr", T, U, V) if ctf.comm().rank() == 0: print("Completed MTTKRP WARMUP with s =", s, "nnz =", nnz, "sp =", sp, "sp_init =", sp_init) while s <= s_end: agg_s.append(s) T = ctf.tensor((s, s, s), sp=sp) T.fill_sp_random(-1., 1., float(nnz) / float(s * s * s)) if ctf.comm().rank() == 0: print("Performing MTTKRP with s =", s, "nnz =", nnz, "sp =", sp, "sp_init =", sp_init) U = ctf.random.random((s, R)) V = ctf.random.random((s, R)) W = ctf.random.random((s, R)) te1 = 0. te2 = 0. te3 = 0. avg_times = [] for i in range(num_iter): t0 = time.time() U = ctf.einsum("ijk,jr,kr->ir", T, V, W) t1 = time.time() ite1 = t1 - t0 te1 += ite1 t0 = time.time() V = ctf.einsum("ijk,ir,kr->jr", T, U, W) t1 = time.time() ite2 = t1 - t0 te2 += ite2 t0 = time.time() W = ctf.einsum("ijk,ir,jr->kr", T, U, V) t1 = time.time() ite3 = t1 - t0 te3 += ite3 if ctf.comm().rank() == 0: print(ite1, ite2, ite3, "avg:", (ite1 + ite2 + ite3) / 3.) avg_times.append((ite1 + ite2 + ite3) / 3.) if ctf.comm().rank() == 0: print("Completed", num_iter, "iterations, took", te1 / num_iter, te2 / num_iter, te3 / num_iter, "seconds on average for 3 variants.") avg_time = (te1 + te2 + te3) / (3 * num_iter) agg_avg_times.append(avg_time) print("MTTKRP took", avg_times, "seconds on average across variants with s =", s, "nnz =", nnz, "sp =", sp, "sp_init =", sp_init) min_time = np.min(avg_times) max_time = np.max(avg_times) agg_min_times.append(min_time) agg_max_times.append(max_time) print("min/max interval is [", min_time, ",", max_time, "]") stddev = np.std(avg_times) min_95 = (te1 + te2 + te3) / (3 * num_iter) - 2 * stddev max_95 = (te1 + te2 + te3) / (3 * num_iter) + 2 * stddev agg_min_95.append(min_95) agg_max_95.append(max_95) print("95% confidence interval is [", min_95, ",", max_95, "]") s = int(s * mult) if ctf.comm().rank() == 0: print("s min_time min_95 avg_time max_95 max_time") for i in range(len(agg_s)): print(agg_s[i], agg_min_times[i], agg_min_95[i], agg_avg_times[i], agg_max_95[i], agg_max_times[i])
[z, np.exp(-s) * Sp, -n, I]]) W.append(site_i) # Last Site site_N = ctf.astensor([[I], [Sm], [v], [beta * (np.exp(-s) * Sp - n)]]) W.append(site_N) ############################################## # Canonicalize MPS ########################### for i in range(int(N) - 1, 0, -1): M_reshape = ctf.transpose(M[i], axes=[1, 0, 2]) (n1, n2, n3) = M_reshape.shape M_reshape = M_reshape.reshape(n1, n2 * n3) (U, S, V) = ctf.svd(M_reshape) M_reshape = V.reshape(n1, n2, n3) M[i] = ctf.transpose(M_reshape, axes=[1, 0, 2]) M[i - 1] = ctf.einsum('klj,ji,i->kli', M[i - 1], U, S) ############################################## # Create Environment ######################### print('Generating Environment') # Allocate empty environment F = [] tmp = np.array([[[1.]]]) + 0.j F.append(ctf.from_nparray(tmp)) for i in range(int(N / 2)): tmp = np.zeros( (min(2**(i + 1), maxBondDim), 4, min(2**(i + 1), maxBondDim))) + 0.j F.append(ctf.from_nparray(tmp)) for i in range(int(N / 2) - 1, 0, -1): tmp = np.zeros((min(2**(i), maxBondDim), 4, min(2**i, maxBondDim))) + 0.j F.append(ctf.from_nparray(tmp))
def mult_lists(list_A, list_B): s = 0 for i in range(len(list_A)): s += ctf.einsum('ij,ij->', list_A[i], list_B[i]) return s
site_N = ctf.astensor([[I ], [Sm ], [v ], [beta*(np.exp(-s)*Sp-n)]]) W.append(site_N) ############################################## # Canonicalize MPS ########################### for i in range(int(N)-1,0,-1): M_reshape = ctf.transpose(M[i],axes=[1,0,2]) (n1,n2,n3) = M_reshape.shape M_reshape = M_reshape.reshape(n1,n2*n3) (U,S,V) = ctf.svd(M_reshape) M_reshape = V.reshape(n1,n2,n3) M[i] = ctf.transpose(M_reshape,axes=[1,0,2]) M[i-1] = ctf.einsum('klj,ji,i->kli',M[i-1],U,s) ############################################## # Create Environment ######################### # Allocate empty environment F = [] tmp = np.array([[[1.]]]) F.append(ctf.from_nparray(tmp)) print(F[0]) for i in range(int(N/2)): tmp = np.zeros((min(2**(i+1),maxBondDim),4,min(2**(i+1),maxBondDim))) F.append(ctf.from_nparray(tmp)) for i in range(int(N/2)-1,0,-1): tmp = np.zeros((min(2**(i),maxBondDim),4,min(2**i,maxBondDim))) F.append(ctf.from_nparray(tmp)) tmp = np.array([[[1.]]])
def __init__(self, cc, mo_coeff=None): self.lib = cc.lib log = Logger(cc.stdout, cc.verbose) nocc, nmo = cc.nocc, cc.nmo nvir = nmo - nocc cput1 = cput0 = (time.clock(), time.time()) if mo_coeff is None: self.mo_coeff = mo_coeff = ccsd._mo_without_core(cc, cc.mo_coeff) else: self.mo_coeff = mo_coeff = ccsd._mo_without_core(cc, mo_coeff) fock = None if rank == 0: dm = cc._scf.make_rdm1(cc.mo_coeff, cc.mo_occ) fockao = cc._scf.get_hcore() + cc._scf.get_veff(cc.mol, dm) fock = reduce(np.dot, (mo_coeff.T, fockao, mo_coeff)) fock = comm.bcast(fock, root=0) self.dtype = dtype = np.result_type(fock) self.foo = zeros([nocc, nocc], dtype, verbose=cc.SYMVERBOSE) self.fov = zeros([nocc, nvir], dtype, verbose=cc.SYMVERBOSE) self.fvv = zeros([nvir, nvir], dtype, verbose=cc.SYMVERBOSE) self.eia = zeros([nocc, nvir], verbose=cc.SYMVERBOSE) if rank == 0: mo_e = fock.diagonal().real eia = mo_e[:nocc, None] - mo_e[None, nocc:] self.foo.write(range(nocc * nocc), fock[:nocc, :nocc].ravel()) self.fov.write(range(nocc * nvir), fock[:nocc, nocc:].ravel()) self.fvv.write(range(nvir * nvir), fock[nocc:, nocc:].ravel()) self.eia.write(range(nocc * nvir), eia.ravel()) else: self.foo.write([], []) self.fov.write([], []) self.fvv.write([], []) self.eia.write([], []) cput1 = log.timer("Writing Fock", *cput1) self._foo = self.foo.diagonal(preserve_shape=True) self._fvv = self.fvv.diagonal(preserve_shape=True) eijab = self.eia.array.reshape( nocc, 1, nvir, 1) + self.eia.array.reshape(1, nocc, 1, nvir) self.eijab = tensor(eijab, verbose=cc.SYMVERBOSE) ppoo, ppov, ppvv = _make_ao_ints(cc.mol, mo_coeff, nocc, dtype) cput1 = log.timer('making ao integrals', *cput1) mo = ctf.astensor(mo_coeff) orbo, orbv = mo[:, :nocc], mo[:, nocc:] tmp = ctf.einsum('uvmn,ui->ivmn', ppoo, orbo.conj()) oooo = ctf.einsum('ivmn,vj->ijmn', tmp, orbo) ooov = ctf.einsum('ivmn,va->mnia', tmp, orbv) tmp = ctf.einsum('uvma,vb->ubma', ppov, orbv) ovov = ctf.einsum('ubma,ui->ibma', tmp, orbo.conj()) tmp = ctf.einsum('uvma,ub->mabv', ppov, orbv.conj()) ovvo = ctf.einsum('mabv,vi->mabi', tmp, orbo) tmp = ctf.einsum('uvab,ui->ivab', ppvv, orbo.conj()) oovv = ctf.einsum('ivab,vj->ijab', tmp, orbo) tmp = ctf.einsum('uvab,vc->ucab', ppvv, orbv) ovvv = ctf.einsum('ucab,ui->icab', tmp, orbo.conj()) vvvv = ctf.einsum('ucab,ud->dcab', tmp, orbv.conj()) self.oooo = tensor(oooo, verbose=cc.SYMVERBOSE) self.ooov = tensor(ooov, verbose=cc.SYMVERBOSE) self.ovov = tensor(ovov, verbose=cc.SYMVERBOSE) self.oovv = tensor(oovv, verbose=cc.SYMVERBOSE) self.ovvo = tensor(ovvo, verbose=cc.SYMVERBOSE) self.ovvv = tensor(ovvv, verbose=cc.SYMVERBOSE) self.vvvv = tensor(vvvv, verbose=cc.SYMVERBOSE) log.timer('ao2mo transformation', *cput0)
def main(): t0 = time.time() ######## Inputs ############################## # SEP Model N = 50 alpha = 0.35 # In at left beta = 2. / 3. # Exit at right s = -1. # Exponential weighting gamma = 0. # Exit at left delta = 0. # In at right p = 1. # Jump right q = 0. # Jump Left # Optimization tol = 1e-5 maxIter = 0 maxBondDim = 10 useCTF = True ############################################## # Create MPS ################################# # PH - Make Isometries, Center Site mpiprint('Generating MPS') M = [] for i in range(int(N / 2)): tmp = np.random.rand(2, min(2**(i),maxBondDim), min(2**(i+1),maxBondDim))\ +0.j M.append(ctf.from_nparray(tmp)) for i in range(int(N / 2))[::-1]: tmp = np.random.rand(2, min(2**(i+1),maxBondDim), min(2**i,maxBondDim))\ +0.j M.append(ctf.from_nparray(tmp)) ############################################## # Create MPO ################################# mpiprint('Generating MPO') # Simple Operators Sp = np.array([[0., 1.], [0., 0.]]) Sm = np.array([[0., 0.], [1., 0.]]) n = np.array([[0., 0.], [0., 1.]]) v = np.array([[1., 0.], [0., 0.]]) I = np.array([[1., 0.], [0., 1.]]) z = np.array([[0., 0.], [0., 0.]]) # List to hold MPOs W = [] # First Site site_0 = ctf.astensor( [[alpha * (np.exp(-s) * Sm - v), np.exp(-s) * Sp, -n, I]]) W.append(site_0) # Central Sites for i in range(N - 2): site_i = ctf.astensor([[I, z, z, z], [Sm, z, z, z], [v, z, z, z], [z, np.exp(-s) * Sp, -n, I]]) W.append(site_i) # Last Site site_N = ctf.astensor([[I], [Sm], [v], [beta * (np.exp(-s) * Sp - n)]]) W.append(site_N) ############################################## # Canonicalize MPS ########################### for i in range(int(N) - 1, 0, -1): M_reshape = ctf.transpose(M[i], axes=[1, 0, 2]) (n1, n2, n3) = M_reshape.shape M_reshape = M_reshape.reshape(n1, n2 * n3) (U, S, V) = ctf.svd(M_reshape) M_reshape = V.reshape(n1, n2, n3) M[i] = ctf.transpose(M_reshape, axes=[1, 0, 2]) M[i - 1] = ctf.einsum('klj,ji,i->kli', M[i - 1], U, S) ############################################## # Canonicalize MPS ########################### def pick_eigs(w, v, nroots, x0): idx = np.argsort(np.real(w)) w = w[idx] v = v[:, idx] return w, v, idx ############################################## # Create Environment ######################### mpiprint('Generating Environment') # Allocate empty environment F = [] tmp = np.array([[[1.]]]) + 0.j F.append(ctf.from_nparray(tmp)) for i in range(int(N / 2)): tmp = np.zeros((min(2**(i + 1), maxBondDim), 4, min(2**(i + 1), maxBondDim))) + 0.j F.append(ctf.from_nparray(tmp)) for i in range(int(N / 2) - 1, 0, -1): tmp = np.zeros( (min(2**(i), maxBondDim), 4, min(2**i, maxBondDim))) + 0.j F.append(ctf.from_nparray(tmp)) tmp = np.array([[[1.]]]) + 0.j F.append(ctf.from_nparray(tmp)) # Calculate initial environment for i in range(int(N) - 1, 0, -1): tmp = ctf.einsum('eaf,cdf->eacd', M[i], F[i + 1]) tmp = ctf.einsum('ydbe,eacd->ybac', W[i], tmp) F[i] = ctf.einsum('bxc,ybac->xya', ctf.conj(M[i]), tmp) ############################################## # Optimization Sweeps ######################## converged = False iterCnt = 0 E_prev = 0 while not converged: # Right Sweep ---------------------------- tr = time.time() mpiprint('Start Right Sweep {}'.format(iterCnt)) for i in range(N - 1): (n1, n2, n3) = M[i].shape # Make Hx Function def Hfun(x): x_reshape = ctf.array(x) x_reshape = ctf.reshape(x_reshape, (n1, n2, n3)) tmp = ctf.einsum('ijk,lmk->ijlm', F[i + 1], x_reshape) tmp = ctf.einsum('njol,ijlm->noim', W[i], tmp) res = ctf.einsum('pnm,noim->opi', F[i], tmp) return -ctf.reshape(res, -1).to_nparray() def precond(dx, e, x0): return dx # Set up initial guess guess = ctf.reshape(M[i], -1).to_nparray() # Run eigenproblem u, v = eig(Hfun, guess, precond, pick=pick_eigs) E = -u v = ctf.array(v) M[i] = ctf.reshape(v, (n1, n2, n3)) # Print Results mpiprint('\tEnergy at site {} = {}'.format(i, E)) # Right Normalize M_reshape = ctf.reshape(M[i], (n1 * n2, n3)) (U, S, V) = ctf.svd(M_reshape) M[i] = ctf.reshape(U, (n1, n2, n3)) M[i + 1] = ctf.einsum('i,ij,kjl->kil', S, V, M[i + 1]) # Update F tmp = ctf.einsum('jlp,ijk->lpik', F[i], ctf.conj(M[i])) tmp = ctf.einsum('lmin,lpik->mpnk', W[i], tmp) F[i + 1] = ctf.einsum('npq,mpnk->kmq', M[i], tmp) mpiprint('Complete Right Sweep {}, {} sec'.format( iterCnt, time.time() - tr)) # Left Sweep ------------------------------ tl = time.time() mpiprint('Start Left Sweep {}'.format(iterCnt)) for i in range(N - 1, 0, -1): (n1, n2, n3) = M[i].shape # Make Hx Function def Hfun(x): x_reshape = ctf.array(x) x_reshape = ctf.reshape(x_reshape, (n1, n2, n3)) tmp = ctf.einsum('ijk,lmk->ijlm', F[i + 1], x_reshape) tmp = ctf.einsum('njol,ijlm->noim', W[i], tmp) res = ctf.einsum('pnm,noim->opi', F[i], tmp) return -ctf.reshape(res, -1).to_nparray() def precond(dx, e, x0): return dx # Set up initial guess guess = ctf.reshape(M[i], -1).to_nparray() # Run eigenproblem u, v = eig(Hfun, guess, precond, pick=pick_eigs) E = -u v = ctf.array(v) M[i] = ctf.reshape(v, (n1, n2, n3)) # Print Results mpiprint('\tEnergy at site {}= {}'.format(i, E)) # Right Normalize M_reshape = ctf.transpose(M[i], (1, 0, 2)) M_reshape = ctf.reshape(M_reshape, (n2, n1 * n3)) (U, S, V) = ctf.svd(M_reshape) M_reshape = ctf.reshape(V, (n2, n1, n3)) M[i] = ctf.transpose(M_reshape, (1, 0, 2)) M[i - 1] = ctf.einsum('klj,ji,i->kli', M[i - 1], U, S) # Update F tmp = ctf.einsum('eaf,cdf->eacd', M[i], F[i + 1]) tmp = ctf.einsum('ydbe,eacd->ybac', W[i], tmp) F[i] = ctf.einsum('bxc,ybac->xya', ctf.conj(M[i]), tmp) mpiprint('Left Sweep {}, {} sec'.format(iterCnt, time.time() - tl)) # Convergence Test ----------------------- if np.abs(E - E_prev) < tol: mpiprint('#' * 75 + '\nConverged at E = {}'.format(E) + '\n' + '#' * 75) converged = True elif iterCnt > maxIter: mpiprint('Convergence not acheived') converged = True else: iterCnt += 1 E_prev = E mpiprint('Total Time = {}'.format(time.time() - t0))
def _make_eris(mycc, mo_coeff=None): mol = mycc.mol NS = ctf.SYM.NS SY = ctf.SYM.SY eris = ccsd._ChemistsERIs() if mo_coeff is None: mo_coeff = mycc.mo_coeff eris.mo_coeff = ccsd._mo_without_core(mycc, mo_coeff) nao, nmo = eris.mo_coeff.shape nocc = mycc.nocc nvir = nmo - nocc nvir_pair = nvir * (nvir + 1) // 2 nao_pair = nao * (nao + 1) // 2 mo = ctf.astensor(eris.mo_coeff) ppoo, ppov, ppvv = _make_ao_ints(mol, eris.mo_coeff, nocc) eris.nocc = mycc.nocc eris.mol = mycc.mol eris.fock = ctf.tensor((nmo, nmo)) with omp(16): if rank == 0: dm = mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ) fockao = mycc._scf.get_hcore() + mycc._scf.get_veff(mycc.mol, dm) fock = reduce(numpy.dot, (eris.mo_coeff.T, fockao, eris.mo_coeff)) eris.fock.write(numpy.arange(nmo**2), fock.ravel()) else: eris.fock.write([], []) orbo = mo[:, :nocc] orbv = mo[:, nocc:] print 'before contraction', rank, lib.current_memory() tmp = ctf.tensor([nao, nocc, nvir, nvir], sym=[NS, NS, SY, NS]) otmp = ctf.einsum('pqrs,qj->pjrs', ppvv, orbo, out=tmp) eris.oovv = ctf.tensor([nocc, nocc, nvir, nvir], sym=[NS, NS, SY, NS]) eris.ovvv = ctf.einsum('pjrs,pi->ijrs', otmp, orbo, out=eris.oovv) otmp = tmp = None print '___________ vvvv', rank, lib.current_memory() tmp = ctf.tensor([nao, nvir, nvir, nvir], sym=[NS, NS, SY, NS]) print '___________ vvvv sub1', rank, lib.current_memory() vtmp = ctf.einsum('pqrs,qj->pjrs', ppvv, orbv, out=tmp) eris.ovvv = ctf.tensor([nocc, nvir, nvir, nvir], sym=[NS, NS, SY, NS]) eris.ovvv = ctf.einsum('pjrs,pi->ijrs', vtmp, orbo, out=eris.ovvv) ppvv = None tmp = ctf.tensor([nvir, nvir, nvir, nvir], sym=[NS, NS, SY, NS]) print '___________ vvvv sub2', rank, lib.current_memory() vtmp = ctf.einsum('pjrs,pi->ijrs', vtmp, orbv, out=tmp) eris.vvvv = ctf.tensor(sym=[SY, NS, SY, NS], copy=vtmp) vtmp = tmp = None vtmp = ctf.einsum('pqrs,qj->pjrs', ppov, orbv) eris.ovov = ctf.einsum('pjrs,pi->ijrs', vtmp, orbo) vtmp = None otmp = ctf.einsum('pqrs,qj->pjrs', ppov, orbo) eris.ooov = ctf.einsum('pjrs,pi->ijrs', otmp, orbo) ppov = otmp = None otmp = ctf.einsum('pqrs,qj->pjrs', ppoo, orbo) eris.oooo = ctf.einsum('pjrs,pi->ijrs', otmp, orbo) ppoo = otmp = None print '___________ fock', rank, lib.current_memory() eris.mo_energy = eris.fock.diagonal() eris.foo = eris.fock[:nocc, :nocc].copy() eris.foo.i('ii') << (eris.mo_energy[:nocc] * -1).i('i') eris.fvv = eris.fock[nocc:, nocc:].copy() eris.fvv.i('ii') << (eris.mo_energy[nocc:] * -1).i('i') eris.fov = ctf.astensor(eris.fock[:nocc, nocc:]) return eris
def update_amps(t1, t2, eris): time0 = time.clock(), time.time() nocc, nvir = t1.shape nov = nocc * nvir t1new = ctf.tensor(t1.shape) t2new = ctf.tensor(t2.shape) #tau = t2 + ctf.einsum('ia,jb->ijab', t1, t1) #t2new = ctf.einsum('acbd,ijab->ijcd', eris.vvvv, tau) ctf.einsum('acbd,ijab->ijcd', eris.vvvv, t2, out=t2new) ctf.einsum('acbd,ia,jb->ijcd', eris.vvvv, t1, t1, out=t2new) t2new *= .5 t1new += eris.fov foo = eris.foo + .5 * ctf.einsum('ia,ja->ij', eris.fov, t1) fvv = eris.fvv - .5 * ctf.einsum('ia,ib->ab', t1, eris.fov) ctf.einsum('kc,jikc->ij', 2 * t1, eris.ooov, out=foo) ctf.einsum('kc,jkic->ij', -1 * t1, eris.ooov, out=foo) woooo = ctf.einsum('ijka,la->ijkl', eris.ooov, t1) woooo = woooo + woooo.transpose(2, 3, 0, 1) woooo += eris.oooo #for p0, p1 in lib.prange(0, nvir, vblk): ctf.einsum('kc,kcba->ab', 2 * t1, eris.ovvv, out=fvv) ctf.einsum('kc,kbca->ab', -1 * t1, eris.ovvv, out=fvv) woVoV = ctf.einsum('ijka,kb->ijab', eris.ooov, t1) woVoV -= ctf.einsum('jc,icab->ijab', t1, eris.ovvv) #:eris_ovvv = ctf.einsum('ial,bcl->iabc', ovL, vvL) #:tmp = ctf.einsum('ijcd,kcdb->kijb', tau, eris_ovvv) #:t2new += ctf.einsum('ka,kijb->jiba', -t1, tmp) #tmp = ctf.einsum('ijcd,kcdb->kijb', tau, eris.ovvv) tmp = ctf.tensor([nocc, nocc, nocc, nvir]) ctf.einsum('ic,jd,kcdb->kijb', t1, t1, eris.ovvv, out=tmp) ctf.einsum('ijcd,kcdb->kijb', t2, eris.ovvv, out=tmp) t2new -= ctf.einsum('ka,kijb->jiba', t1, tmp) wOVov = ctf.einsum('ikjb,ka->ijba', eris.ooov, -1 * t1) ctf.einsum('jc,iabc->jiab', t1, eris.ovvv, out=wOVov) t2new += wOVov.transpose(0, 1, 3, 2) theta = t2.transpose(0, 1, 3, 2) * 2 theta -= t2 ctf.einsum('ijcb,jcba->ia', theta, eris.ovvv, out=t1new) t2new += eris.ovov.transpose(0, 2, 1, 3) * .5 fov = eris.fov.copy() ctf.einsum('kc,iakc->ia', t1, eris.ovov, out=fov) * 2 fov -= ctf.einsum('kc,icka->ia', t1, eris.ovov) ctf.einsum('jb,jiab->ia', fov, theta, out=t1new) t1new -= ctf.einsum('kijb,kjba->ia', eris.ooov, theta) theta = None wOVov += eris.ovov.transpose(0, 2, 3, 1) wOVov -= .5 * ctf.einsum('icka,jkbc->jiab', eris.ovov, t2) tau = t2.transpose(0, 2, 1, 3) * 2 - t2.transpose(0, 3, 1, 2) #tau -= ctf.einsum('ia,jb->ibja', t1*2, t1) #wOVov += .5 * ctf.einsum('iakc,jbkc->jiab', eris.ovov, tau) ctf.einsum('iakc,jc,kb->jiab', eris.ovov, t1, t1, out=wOVov) ctf.einsum('iakc,jbkc->jiab', eris.ovov, .5 * tau, out=wOVov) theta = t2 * 2 - t2.transpose(0, 1, 3, 2) ctf.einsum('ikac,jkcb->jiba', theta, wOVov, out=t2new) tau = theta = wOVov = None tau = ctf.einsum('ia,jb->ijab', t1 * .5, t1) + t2 theta = tau.transpose(0, 1, 3, 2) * 2 - tau fvv -= ctf.einsum('ijca,ibjc->ab', theta, eris.ovov) ctf.einsum('iakb,jkba->ij', eris.ovov, theta, out=foo) tau = theta = None tmp = ctf.einsum('ic,jkbc->jibk', t1, eris.oovv) t2new -= ctf.einsum('ka,jibk->jiab', t1, tmp) tmp = ctf.einsum('ic,jbkc->jibk', t1, eris.ovov) t2new -= ctf.einsum('ka,jibk->jiba', t1, tmp) tmp = None ctf.einsum('jb,iajb->ia', 2 * t1, eris.ovov, out=t1new) ctf.einsum('jb,ijba->ia', -1 * t1, eris.oovv, out=t1new) woVoV -= eris.oovv #tau = t2 + ctf.einsum('ia,jb->ijab', t1, t1) #woooo += ctf.einsum('iajb,klab->ikjl', eris.ovov, tau) #t2new += .5 * ctf.einsum('kilj,klab->ijab', woooo, tau) #tau -= t2 * .5 ctf.einsum('iajb,klab->ikjl', eris.ovov, t2, out=woooo) ctf.einsum('iajb,ka,lb->ikjl', eris.ovov, t1, t1, out=woooo) ctf.einsum('kilj,klab->ijab', woooo, .5 * t2, out=t2new) ctf.einsum('kilj,ka,lb->ijab', woooo, .5 * t1, t1, out=t2new) #tau -= t2 * .5 ctf.einsum('jkca,ickb->ijba', .5 * t2, eris.ovov, out=woVoV) ctf.einsum('jc,ka,ickb->ijba', t1, t1, eris.ovov, out=woVoV) ctf.einsum('kicb,kjac->ijab', woVoV, t2, out=t2new) ctf.einsum('kica,kjcb->ijab', woVoV, t2, out=t2new) woooo = tau = woVoV = None ft_ij = foo + ctf.einsum('ja,ia->ij', .5 * t1, fov) ft_ab = fvv - ctf.einsum('ia,ib->ab', .5 * t1, fov) ctf.einsum('ijac,bc->ijab', t2, ft_ab, out=t2new) t2new -= ctf.einsum('ki,kjab->ijab', ft_ij, t2) mo_e = ctf.tensor([t1.shape[0] + t1.shape[1]]) ctf.random.seed(42) mo_e.fill_random(0., 1.) eia = mo_e[:nocc].reshape(nocc, 1) - mo_e[nocc:].reshape(1, nvir) ctf.einsum('ib,ab->ia', t1, fvv, out=t1new) t1new -= ctf.einsum('ja,ji->ia', t1, foo) t1new /= eia t2new = t2new + t2new.transpose(1, 0, 3, 2) dijab = eia.reshape(nocc, 1, nvir, 1) + eia.reshape(1, nocc, 1, nvir) t2new /= dijab return t1new, t2new
def _make_eris(mycc, mo_coeff=None, cutoff=None): mol = mycc.mol NS = ctf.SYM.NS SY = ctf.SYM.SY eris = ccsd._ChemistsERIs() if mo_coeff is None: mo_coeff = mycc.mo_coeff eris.mo_coeff = ccsd._mo_without_core(mycc, mo_coeff) nao, nmo = eris.mo_coeff.shape nocc = mycc.nocc nvir = nmo - nocc nvir_pair = nvir * (nvir + 1) // 2 nao_pair = nao * (nao + 1) // 2 mo = ctf.astensor(eris.mo_coeff) ppoo, ppov, ppvv = _make_ao_ints(mol, eris.mo_coeff, nocc) eris.nocc = mycc.nocc eris.mol = mycc.mol eris.fock = ctf.tensor((nmo, nmo)) with omp(16): if rank == 0: dm = mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ) fockao = mycc._scf.get_hcore() + mycc._scf.get_veff(mycc.mol, dm) fock = reduce(numpy.dot, (eris.mo_coeff.T, fockao, eris.mo_coeff)) eris.fock.write(numpy.arange(nmo**2), fock.ravel()) else: eris.fock.write([], []) orbo = mo[:, :nocc] orbv = mo[:, nocc:] if (ctf.comm().rank() == 0): print 'before contraction', rank, lib.current_memory() tmp = ctf.tensor([nao, nocc, nvir, nvir], sym=[NS, NS, SY, NS]) otmp = ctf.einsum('pqrs,qj->pjrs', ppvv, orbo, out=tmp) eris.oovv = ctf.tensor([nocc, nocc, nvir, nvir], sym=[NS, NS, SY, NS]) eris.ovvv = ctf.einsum('pjrs,pi->ijrs', otmp, orbo, out=eris.oovv) otmp = tmp = None if (ctf.comm().rank() == 0): print '___________ vvvv', rank, lib.current_memory() tmp = ctf.tensor([nao, nvir, nvir, nvir], sym=[NS, NS, SY, NS]) if (ctf.comm().rank() == 0): print '___________ vvvv sub1', rank, lib.current_memory() vtmp = ctf.einsum('pqrs,qj->pjrs', ppvv, orbv, out=tmp) eris.ovvv = ctf.tensor([nocc, nvir, nvir, nvir], sym=[NS, NS, SY, NS]) eris.ovvv = ctf.einsum('pjrs,pi->ijrs', vtmp, orbo, out=eris.ovvv) ppvv = None tmp = ctf.tensor([nvir, nvir, nvir, nvir], sym=[NS, NS, SY, NS]) if (ctf.comm().rank() == 0): print '___________ vvvv sub2', rank, lib.current_memory() vtmp = ctf.einsum('pjrs,pi->ijrs', vtmp, orbv, out=tmp) eris.vvvv = ctf.tensor(sym=[SY, NS, SY, NS], copy=vtmp) vtmp = tmp = None vtmp = ctf.einsum('pqrs,qj->pjrs', ppov, orbv) eris.ovov = ctf.einsum('pjrs,pi->ijrs', vtmp, orbo) vtmp = None otmp = ctf.einsum('pqrs,qj->pjrs', ppov, orbo) eris.ooov = ctf.einsum('pjrs,pi->ijrs', otmp, orbo) ppov = otmp = None otmp = ctf.einsum('pqrs,qj->pjrs', ppoo, orbo) eris.oooo = ctf.einsum('pjrs,pi->ijrs', otmp, orbo) ppoo = otmp = None if cutoff != None: print("Using cutoff", cutoff) eris.ovvv = eris.ovvv.sparsify(ctf.vecnorm(eris.ovvv) * cutoff) eris.oovv = eris.oovv.sparsify(ctf.vecnorm(eris.oovv) * cutoff) eris.oooo = eris.oooo.sparsify(ctf.vecnorm(eris.oooo) * cutoff) eris.ooov = eris.ooov.sparsify(ctf.vecnorm(eris.ooov) * cutoff) eris.vvvv = eris.vvvv.sparsify(ctf.vecnorm(eris.vvvv) * cutoff) eris.ovov = eris.ovov.sparsify(ctf.vecnorm(eris.ovov) * cutoff) if (ctf.comm().rank() == 0): for e in [ eris.ovvv, eris.oovv, eris.oooo, eris.ooov, eris.vvvv, eris.ovov ]: print "For integral tensor with shape,", e.shape, "symmetry", e.sym, "number of nonzeros with cutoff", cutoff, "is ", ( int(100000 * e.nnz_tot / e.size)) / 1000, "%" if (ctf.comm().rank() == 0): print '___________ fock', rank, lib.current_memory() eris.mo_energy = eris.fock.diagonal() eris.foo = eris.fock[:nocc, :nocc].copy() eris.foo.i('ii') << (eris.mo_energy[:nocc] * -1).i('i') eris.fvv = eris.fock[nocc:, nocc:].copy() eris.fvv.i('ii') << (eris.mo_energy[nocc:] * -1).i('i') eris.fov = ctf.astensor(eris.fock[:nocc, nocc:]) return eris
def eaccsd_diag(eom, kshift, imds=None, diag=None): # Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1994) Eqs.(30)-(31) if imds is None: imds = eom.make_imds() kpts, gvec = eom.kpts, eom._cc._scf.cell.reciprocal_vectors() t1, t2 = imds.t1, imds.t2 dtype = t2.dtype nocc, nvir = t1.shape nkpts = len(kpts) kconserv = eom.kconserv sym1 = ['+', [ kpts, ], kpts[kshift], gvec] sym2 = ['-++', [ kpts, ] * 3, kpts[kshift], gvec] Hr1array = imds.Lvv.diagonal()[kshift] Hr1 = tensor(Hr1array, sym1, symlib=eom.symlib) Hr2 = zeros([nocc, nvir, nvir], dtype, sym2, symlib=eom.symlib) tasks = static_partition(range(nkpts**2)) ntasks = max(comm.allgather(len(tasks))) idx_jab = np.arange(nocc * nvir * nvir) if eom.partition == 'mp': foo = imds.eris.foo.diagonal().to_nparray() fvv = imds.eris.fvv.diagonal().to_nparray() for itask in range(ntasks): if itask >= len(tasks): Hr2.write([], []) continue kjab = tasks[itask] kj, ka = (kjab).__divmod__(nkpts) kb = kconserv[ki, ka, kshift] off = kj * nkpts + ka jab = np.zeros([nocc, nvir, nvir], dtype=dtype) jab += -foo[kj][:, None, None] jab += fvv[ka][None, :, None] jab += fvv[kb][None, None, :] Hr2.write(off * idx_jab.size + idx_jab, jab.ravel()) else: idx = np.arange(nvir) loo = imds.Loo.diagonal().to_nparray() lvv = imds.Lvv.diagonal().to_nparray() wab = ctf.einsum("ABAabab->ABab", imds.Wvvvv.array).to_nparray() wjb = ctf.einsum('JBJjbjb->JBjb', imds.Wovov.array).to_nparray() wjb2 = ctf.einsum('JBBjbbj->JBjb', imds.Wovvo.array).to_nparray() wja = ctf.einsum('JAJjaja->JAja', imds.Wovov.array).to_nparray() for itask in range(ntasks): if itask >= len(tasks): Hr2.write([], []) continue kjab = tasks[itask] kj, ka = (kjab).__divmod__(nkpts) kb = kconserv[kj, ka, kshift] jab = np.zeros([nocc, nvir, nvir], dtype=dtype) jab -= loo[kj][:, None, None] jab += lvv[ka][None, :, None] jab += lvv[kb][None, None, :] jab += wab[ka, kb][None, :, :] jab -= wjb[kj, kb][:, None, :] jab += 2 * wjb2[kj, kb][:, None, :] if ka == kb: jab[:, idx, idx] -= wjb2[kj, ka] jab -= wja[kj, ka][:, :, None] off = kj * nkpts + ka Hr2.write(off * idx_jab.size + idx_jab, jab.ravel()) Hr2 -= 2 * ctf.einsum('JAijab,JAijab->JAjab', t2[kshift], imds.Woovv[kshift]) Woovvtmp = imds.Woovv.transpose(0, 1, 3, 2)[kshift] Hr2 += ctf.einsum('JAijab,JAijab->JAjab', t2[kshift], Woovvtmp) return eom.amplitudes_to_vector(Hr1, Hr2)
def einsvd(einstr, A, rank=None, threshold=None, size_limit=None, criterion=None, mult_s=True): """ Perform Singular Value Decomposition according to the specified Einstein notation string. Will always preserve at least one singular value during the truncation. Parameters ---------- einstr: str A string of Einstein notations in the form of 'idxofA->idxofU,idxofV'. There must be one and only one contraction index. A: tensor_like The tensor to be decomposed. Should be of order 2 or more. rank: int or None, optional The minimum number of singular values/vectors to preserve. Will influence the actual truncation rank. threshold: float or None, optional The value used with criterion to decide the cutoff. Will influence the actual truncation rank. size_limit: int or tuple or None, optional The size limit(s) for both U and V (when specified as a int) or U and V respectively (when specified as a tuple). Will influence the actual truncation rank. criterion: int or None, optional The norm to be used together with threshold to decide the cutoff. Will influence the actual truncation rank. When being left as None, the threshold is treated as the plain cutoff value. Otherwise, cutoff rank is the largest int satisfies: threshold * norm(s) > norm(s[rank:]). mult_s: bool, optional Whether or not to multiply U and V by S**0.5 to decompose A into two tensors instead of three. True by default. Returns ------- u: tensor_like A unitary tensor with indices ordered by the Einstein notation string. s: 1d tensor_like A 1d tensor containing singular values sorted in descending order. v: tensor_like A unitary tensor with indices ordered by the Einstein notation string. """ str_a, str_uv = einstr.replace(' ', '').split('->') str_u, str_v = str_uv.split(',') char_i = list(set(str_v) - set(str_a))[0] shape_u = np.prod([A.shape[str_a.find(c)] for c in str_u if c != char_i]) shape_v = np.prod([A.shape[str_a.find(c)] for c in str_v if c != char_i]) rank = rank or min(shape_u, shape_v) if size_limit is not None: if np.isscalar(size_limit): size_limit = (size_limit, size_limit) if size_limit[0] is not None: rank = min(rank, int(size_limit[0] / shape_u) or 1) if size_limit[1] is not None: rank = min(rank, int(size_limit[1] / shape_v) or 1) if threshold is None or criterion is None: u, s, vh = A.i(str_a).svd(str_u, str_v, rank, threshold) else: u, s, vh = A.i(str_a).svd(str_u, str_v) threshold = threshold * ctf.norm(s, criterion) # will always preserve at least one singular value for i in range(rank, 0, -1): if ctf.norm(s[i - 1:], criterion) >= threshold: rank = i break if rank < s.size: u = u[tuple(slice(None) for i in range(str_u.find(char_i))) + (slice(0, rank), )] s = s[:rank] vh = vh[tuple(slice(None) for i in range(str_v.find(char_i))) + (slice(0, rank), )] if mult_s: char_s = list(set(string.ascii_letters) - set(str_v))[0] sqrtS = ctf.diag(s**0.5) vh = ctf.einsum( char_s + char_i + ',' + str_v + '->' + str_v.replace(char_i, char_s), sqrtS, vh) char_s = list(set(string.ascii_letters) - set(str_u))[0] u = ctf.einsum( str_u + ',' + char_s + char_i + '->' + str_u.replace(char_i, char_s), u, sqrtS) return u, s, vh
def main(): I = 300 J = 300 K = 300 r = 30 sparsity = .1 regParam = .1 ctf.random.seed(42) # # 3rd-order tensor # T = ctf.tensor((I,J,K)) # # T.read_from_file('tensor.txt') # T.fill_sp_random(0,1,sparsity) T = function_tensor(I, J, K, sparsity) t0 = time.time() omega = getOmega(T) if glob_comm.rank() == 0: print('getOmega takes {}'.format(time.time() - t0)) U = ctf.random.random((I, r)) V = ctf.random.random((J, r)) W = ctf.random.random((K, r)) # print(T) # T.write_to_file('tensor_out.txt') # assert(T.sp == 1) # exit(0) # print(U) # print(V) # print(W) ite = 0 objectives = [] t_before_loop = time.time() while True: t0 = time.time() R = ctf.copy(T) t1 = time.time() R -= ctf.einsum('ijk, ir, jr, kr -> ijk', omega, U, V, W) t2 = time.time() R += ctf.einsum('ijk, i, j, k -> ijk', omega, U[:, 0], V[:, 0], W[:, 0]) t3 = time.time() # print(R) # exit(0) t4 = time.time() objective = get_objective(T, U, V, W, I, J, K, omega, regParam) if glob_comm.rank() == 0: print('ctf.copy() takes {}'.format(t1 - t0)) print('ctf.einsum() takes {}'.format(t2 - t1)) print('ctf.einsum() takes {}'.format(t3 - t2)) print('get_objective takes {}'.format(time.time() - t4)) print('Objective: {}'.format(objective)) objectives.append(objective) for f in range(r): # update U[:,f] if glob_comm.rank() == 0: print('updating U[:,{}]'.format(f)) t0 = time.time() alphas = ctf.einsum('ijk, j, k -> i', R, V[:, f], W[:, f]) t1 = time.time() betas = ctf.einsum('ijk, j, j, k, k -> i', omega, V[:, f], V[:, f], W[:, f], W[:, f]) t2 = time.time() U[:, f] = alphas / (regParam + betas) objective = get_objective(T, U, V, W, I, J, K, omega, regParam) if glob_comm.rank() == 0: print('Objective: {}'.format(objective)) print('ctf.einsum() takes {}'.format(t1 - t0)) print('ctf.einsum() takes {}'.format(t2 - t1)) objectives.append(objective) # update V[:,f] if glob_comm.rank() == 0: print('updating V[:,{}]'.format(f)) alphas = ctf.einsum('ijk, i, k -> j', R, U[:, f], W[:, f]) betas = ctf.einsum('ijk, i, i, k, k -> j', omega, U[:, f], U[:, f], W[:, f], W[:, f]) V[:, f] = alphas / (regParam + betas) objective = get_objective(T, U, V, W, I, J, K, omega, regParam) if glob_comm.rank() == 0: print('Objective: {}'.format(objective)) objectives.append(objective) # exit(0) # update W[:,f] if glob_comm.rank() == 0: print('updating W[:,{}]'.format(f)) alphas = ctf.einsum('ijk, i, j -> k', R, U[:, f], V[:, f]) betas = ctf.einsum('ijk, i, i, j, j -> k', omega, U[:, f], U[:, f], V[:, f], V[:, f]) W[:, f] = alphas / (regParam + betas) objective = get_objective(T, U, V, W, I, J, K, omega, regParam) if glob_comm.rank() == 0: print('Objective: {}'.format(objective)) objectives.append(objective) # exit(0) # t0 = time.time() R -= ctf.einsum('ijk, i, j, k -> ijk', omega, U[:, f], V[:, f], W[:, f]) R += ctf.einsum('ijk, i, j, k -> ijk', omega, U[:, f + 1], V[:, f + 1], W[:, f + 1]) # print(time.time() - t0) # print(R) # exit(0) ite += 1 if ite == 1: break if glob_comm.rank() == 0: print('Time/Iteration: {}'.format((time.time() - t_before_loop) / 1))
def run_CCD(T,U,V,W,omega,regParam,num_iter,time_limit,objective_frequency,use_MTTKRP=True): U_vec_list = [] V_vec_list = [] W_vec_list = [] r = U.shape[1] for f in range(r): U_vec_list.append(U[:,f]) V_vec_list.append(V[:,f]) W_vec_list.append(W[:,f]) # print(T) # T.write_to_file('tensor_out.txt') # assert(T.sp == 1) ite = 0 objectives = [] t_before_loop = time.time() t_obj_calc = 0. t_CCD = ctf.timer_epoch("ccd_CCD") t_CCD.begin() while True: t_iR_upd = ctf.timer("ccd_init_R_upd") t_iR_upd.start() t0 = time.time() R = ctf.copy(T) t1 = time.time() # R -= ctf.einsum('ijk, ir, jr, kr -> ijk', omega, U, V, W) R -= ctf.TTTP(omega, [U,V,W]) t2 = time.time() # R += ctf.einsum('ijk, i, j, k -> ijk', omega, U[:,0], V[:,0], W[:,0]) R += ctf.TTTP(omega, [U[:,0], V[:,0], W[:,0]]) t3 = time.time() t_iR_upd.stop() t_b_obj = time.time() if ite % objective_frequency == 0: duration = time.time() - t_before_loop - t_obj_calc [objective, RMSE] = get_objective(T,U,V,W,omega,regParam) objectives.append(objective) if glob_comm.rank() == 0: print('Objective after',duration,'seconds (',ite,'iterations) is: {}'.format(objective)) print('RMSE after',duration,'seconds (',ite,'iterations) is: {}'.format(RMSE)) t_obj_calc += time.time() - t_b_obj if glob_comm.rank() == 0 and status_prints == True: print('ctf.copy() takes {}'.format(t1-t0)) print('ctf.TTTP() takes {}'.format(t2-t1)) print('ctf.TTTP() takes {}'.format(t3-t2)) for f in range(r): # update U[:,f] if glob_comm.rank() == 0 and status_prints == True: print('updating U[:,{}]'.format(f)) t0 = time.time() if use_MTTKRP: alphas = ctf.tensor(R.shape[0]) #ctf.einsum('ijk -> i', ctf.TTTP(R, [None, V_vec_list[f], W_vec_list[f]]),out=alphas) ctf.MTTKRP(R, [alphas, V_vec_list[f], W_vec_list[f]], 0) else: alphas = ctf.einsum('ijk, j, k -> i', R, V_vec_list[f], W_vec_list[f]) t1 = time.time() if use_MTTKRP: betas = ctf.tensor(R.shape[0]) #ctf.einsum('ijk -> i', ctf.TTTP(omega, [None, V_vec_list[f]*V_vec_list[f], W_vec_list[f]*W_vec_list[f]]),out=betas) ctf.MTTKRP(omega, [betas, V_vec_list[f]*V_vec_list[f], W_vec_list[f]*W_vec_list[f]], 0) else: betas = ctf.einsum('ijk, j, j, k, k -> i', omega, V_vec_list[f], V_vec_list[f], W_vec_list[f], W_vec_list[f]) t2 = time.time() U_vec_list[f] = alphas / (regParam + betas) U[:,f] = U_vec_list[f] if glob_comm.rank() == 0 and status_prints == True: print('ctf.einsum() takes {}'.format(t1-t0)) print('ctf.einsum() takes {}'.format(t2-t1)) # update V[:,f] if glob_comm.rank() == 0 and status_prints == True: print('updating V[:,{}]'.format(f)) if use_MTTKRP: alphas = ctf.tensor(R.shape[1]) #ctf.einsum('ijk -> j', ctf.TTTP(R, [U_vec_list[f], None, W_vec_list[f]]),out=alphas) ctf.MTTKRP(R, [U_vec_list[f], alphas, W_vec_list[f]], 1) else: alphas = ctf.einsum('ijk, i, k -> j', R, U_vec_list[f], W_vec_list[f]) if use_MTTKRP: betas = ctf.tensor(R.shape[1]) #ctf.einsum('ijk -> j', ctf.TTTP(omega, [U_vec_list[f]*U_vec_list[f], None, W_vec_list[f]*W_vec_list[f]]),out=betas) ctf.MTTKRP(omega, [U_vec_list[f]*U_vec_list[f], betas, W_vec_list[f]*W_vec_list[f]], 1) else: betas = ctf.einsum('ijk, i, i, k, k -> j', omega, U_vec_list[f], U_vec_list[f], W_vec_list[f], W_vec_list[f]) V_vec_list[f] = alphas / (regParam + betas) V[:,f] = V_vec_list[f] if glob_comm.rank() == 0 and status_prints == True: print('updating W[:,{}]'.format(f)) if use_MTTKRP: alphas = ctf.tensor(R.shape[2]) #ctf.einsum('ijk -> k', ctf.TTTP(R, [U_vec_list[f], V_vec_list[f], None]),out=alphas) ctf.MTTKRP(R, [U_vec_list[f], V_vec_list[f], alphas], 2) else: alphas = ctf.einsum('ijk, i, j -> k', R, U_vec_list[f], V_vec_list[f]) if use_MTTKRP: betas = ctf.tensor(R.shape[2]) #ctf.einsum('ijk -> k', ctf.TTTP(omega, [U_vec_list[f]*U_vec_list[f], V_vec_list[f]*V_vec_list[f], None]),out=betas) ctf.MTTKRP(omega, [U_vec_list[f]*U_vec_list[f], V_vec_list[f]*V_vec_list[f], betas], 2) else: betas = ctf.einsum('ijk, i, i, j, j -> k', omega, U_vec_list[f], U_vec_list[f], V_vec_list[f], V_vec_list[f]) W_vec_list[f] = alphas / (regParam + betas) W[:,f] = W_vec_list[f] t_tttp = ctf.timer("ccd_TTTP") t_tttp.start() R -= ctf.TTTP(omega, [U_vec_list[f], V_vec_list[f], W_vec_list[f]]) if f+1 < r: R += ctf.TTTP(omega, [U_vec_list[f+1], V_vec_list[f+1], W_vec_list[f+1]]) t_tttp.stop() t_iR_upd.stop() ite += 1 if ite == num_iter or time.time() - t_before_loop - t_obj_calc > time_limit: break t_CCD.end() duration = time.time() - t_before_loop - t_obj_calc [objective, RMSE] = get_objective(T,U,V,W,omega,regParam) if glob_comm.rank() == 0: print('CCD amortized seconds per sweep: {}'.format(duration/ite)) print('Time/CCD Iteration: {}'.format(duration/ite)) print('Objective after',duration,'seconds (',ite,'iterations) is: {}'.format(objective)) print('RMSE after',duration,'seconds (',ite,'iterations) is: {}'.format(RMSE))
def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init, use_tttp): wrld = ctf.comm() s = s_start nnz = float(s_start * s_start * s_start) * sp_init agg_s = [] agg_avg_times = [] agg_min_times = [] agg_max_times = [] agg_min_95 = [] agg_max_95 = [] if num_iter > 1: if ctf.comm().rank() == 0: print("Performing TTTP WARMUP with s =", s, "nnz =", nnz, "sp", sp, "sp_init is", sp_init, "use_tttp", use_tttp) T = ctf.tensor((s, s, s), sp=sp) T.fill_sp_random(-1., 1., float(nnz) / float(s * s * s)) U = ctf.random.random((s, R)) V = ctf.random.random((s, R)) W = ctf.random.random((s, R)) if use_tttp: S = ctf.TTTP(T, [U, V, W]) else: if sp: S = ctf.tensor((s, s, s), sp=sp) Z = ctf.tensor((s, s, s, R), sp=sp) Z.i("ijkr") << T.i("ijk") * U.i("ir") Z.i("ijkr") << Z.i("ijkr") * V.i("jr") S.i("ijk") << Z.i("ijkr") * W.i("kr") else: S = ctf.einsum("ijk,iR,jR,kR->ijk", T, U, V, W) if ctf.comm().rank() == 0: print("Completed TTTP WARMUP with s =", s, "nnz =", nnz, "sp", sp, "sp_init is", sp_init, "use_tttp", use_tttp) while s <= s_end: agg_s.append(s) if ctf.comm().rank() == 0: print("Performing TTTP with s =", s, "nnz =", nnz, "sp", sp, "sp_init is", sp_init, "use_tttp", use_tttp) T = ctf.tensor((s, s, s), sp=sp) T.fill_sp_random(-1., 1., float(nnz) / float(s * s * s)) te1 = 0. times = [] if R > 1: U = ctf.random.random((s, R)) V = ctf.random.random((s, R)) W = ctf.random.random((s, R)) for i in range(num_iter): t0 = time.time() if use_tttp: S = ctf.TTTP(T, [U, V, W]) else: if sp: S = ctf.tensor((s, s, s), sp=sp) Z = ctf.tensor((s, s, s, R), sp=sp) Z.i("ijkr") << T.i("ijk") * U.i("ir") Z.i("ijkr") << Z.i("ijkr") * V.i("jr") S.i("ijk") << Z.i("ijkr") * W.i("kr") #S.i("ijk") << T.i("ijk")*U.i("iR")*V.i("jR")*W.i("kR") else: S = ctf.einsum("ijk,iR,jR,kR->ijk", T, U, V, W) t1 = time.time() ite1 = t1 - t0 te1 += ite1 times.append(ite1) if ctf.comm().rank() == 0: print(ite1) else: U = ctf.random.random((s)) V = ctf.random.random((s)) W = ctf.random.random((s)) for i in range(num_iter): t0 = time.time() if use_tttp: S = ctf.TTTP(T, [U, V, W]) else: if sp: S = ctf.tensor((s, s, s), sp=sp) S.i("ijk") << T.i("ijk") * U.i("i") 0.0 * S.i("ijk") << S.i("ijk") * V.i("j") 0.0 * S.i("ijk") << S.i("ijk") * W.i("k") else: S = ctf.einsum("ijk,i,j,k->ijk", T, U, V, W) t1 = time.time() ite1 = t1 - t0 te1 += ite1 times.append(ite1) if ctf.comm().rank() == 0: print(ite1) if ctf.comm().rank() == 0: avg_time = (te1) / (num_iter) agg_avg_times.append(avg_time) print("TTTP", avg_time, "seconds on average with s =", s, "nnz =", nnz, "sp", sp, "sp_init is", sp_init, "use_tttp", use_tttp) min_time = np.min(times) max_time = np.max(times) agg_min_times.append(min_time) agg_max_times.append(max_time) print("min/max interval is [", min_time, ",", max_time, "]") stddev = np.std(times) min_95 = te1 / num_iter - 2 * stddev max_95 = te1 / num_iter + 2 * stddev agg_min_95.append(min_95) agg_max_95.append(max_95) print("95% confidence interval is [", min_95, ",", max_95, "]") s = int(s * mult) if ctf.comm().rank() == 0: print("s min_time min_95 avg_time max_95 max_time") for i in range(len(agg_s)): print(agg_s[i], agg_min_times[i], agg_min_95[i], agg_avg_times[i], agg_max_95[i], agg_max_times[i])
def test_einsum_sums(self): # outer(a,b) for n in range(1, 17): a0 = numpy.arange(3, dtype=numpy.double) + 1 b0 = numpy.arange(n, dtype=numpy.double) + 1 a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue( allclose(ctf.einsum("i,j", a1, b1), numpy.outer(a0, b0))) # matvec(a,b) / a.dot(b) where a is matrix, b is vector for n in range(1, 17): a0 = numpy.arange(4 * n, dtype=numpy.double).reshape(n, 4) b0 = numpy.arange(n, dtype=numpy.double) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue( allclose(ctf.einsum("ji,j", a1, b1), numpy.dot(b0.T, a0))) self.assertTrue( allclose(ctf.einsum("ji,j->", a1, b1), numpy.dot(b0.T, a0).sum())) # matmat(a,b) / a.dot(b) where a is matrix, b is matrix for n in range(1, 17): a0 = numpy.arange(4 * n, dtype=numpy.double).reshape(n, 4) b0 = numpy.arange(6 * n, dtype=numpy.double).reshape(n, 6) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue( allclose(ctf.einsum("ji,jk", a1, b1), numpy.dot(a0.T, b0))) self.assertTrue( allclose(ctf.einsum("ji,jk->", a1, b1), numpy.dot(a0.T, b0).sum())) # matrix triple product (note this is not currently an efficient # way to multiply 3 matrices) a0 = numpy.arange(12.).reshape(3, 4) b0 = numpy.arange(20.).reshape(4, 5) c0 = numpy.arange(30.).reshape(5, 6) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) c1 = ctf.astensor(c0) self.assertTrue( allclose(ctf.einsum("ij,jk,kl", a1, b1, c1), numpy.einsum("ij,jk,kl", a0, b0, c0))) # tensordot(a, b) a0 = numpy.arange(27.).reshape(3, 3, 3) b0 = numpy.arange(27.).reshape(3, 3, 3) a1 = ctf.astensor(a0) b1 = ctf.astensor(b0) self.assertTrue( allclose(ctf.einsum("ijk, jli -> kl", a1, b1), numpy.einsum("ijk, jli -> kl", a0, b0))) self.assertTrue( allclose(ctf.einsum("ijk, jli -> lk", a1, b1), numpy.einsum("ijk, jli -> lk", a0, b0))) self.assertTrue( allclose(ctf.einsum("ikj, jli -> kl", a1, b1), numpy.einsum("ikj, jli -> kl", a0, b0))) self.assertTrue( allclose(ctf.einsum("kij, lij -> lk", a1, b1), numpy.einsum("kij, lij -> lk", a0, b0)))
def ipccsd_diag(eom, kshift, imds=None): if imds is None: imds = eom.make_imds() kpts, gvec = eom.kpts, eom._cc._scf.cell.reciprocal_vectors() t1, t2 = imds.t1, imds.t2 dtype = t2.dtype nocc, nvir = t1.shape nkpts = len(kpts) kconserv = eom.kconserv Hr1array = -imds.Loo.diagonal()[kshift] sym1 = ['+', [ kpts, ], kpts[kshift], gvec] sym2 = ['++-', [ kpts, ] * 3, kpts[kshift], gvec] Hr1 = tensor(Hr1array, sym1, symlib=eom.symlib) Hr2 = zeros([nocc, nocc, nvir], dtype, sym2, symlib=eom.symlib) tasks = static_partition(range(nkpts**2)) ntasks = max(comm.allgather(len(tasks))) idx_ijb = np.arange(nocc * nocc * nvir) if eom.partition == 'mp': foo = eom.eris.foo.diagonal().to_nparray() fvv = eom.eris.fvv.diagonal().to_nparray() for itask in range(ntasks): if itask >= len(tasks): Hr2.write([], []) continue kijb = tasks[itask] ki, kj = (kijb).__divmod__(nkpts) kb = kconserv[ki, kshift, kj] off = ki * nkpts + kj ijb = np.zeros([nocc, nocc, nvir], dtype=dtype) ijb += fvv[kb].reshape(1, 1, -1) ijb -= foo[ki][:, None, None] ijb -= foo[kj][None, :, None] Hr2.write(off * idx_ijb.size + idx_ijb, ijb.ravel()) else: lvv = imds.Lvv.diagonal().to_nparray() loo = imds.Loo.diagonal().to_nparray() wij = ctf.einsum('IJIijij->IJij', imds.Woooo.array).to_nparray() wjb = ctf.einsum('JBJjbjb->JBjb', imds.Wovov.array).to_nparray() wjb2 = ctf.einsum('JBBjbbj->JBjb', imds.Wovvo.array).to_nparray() wib = ctf.einsum('IBIibib->IBib', imds.Wovov.array).to_nparray() idx = np.arange(nocc) for itask in range(ntasks): if itask >= len(tasks): Hr2.write([], []) continue kijb = tasks[itask] ki, kj = (kijb).__divmod__(nkpts) kb = kconserv[ki, kshift, kj] ijb = np.zeros([nocc, nocc, nvir], dtype=dtype) ijb += lvv[kb][None, None, :] ijb -= loo[ki][:, None, None] ijb -= loo[kj][None, :, None] ijb += wij[ki, kj][:, :, None] ijb -= wjb[kj, kb][None, :, :] ijb += 2 * wjb2[kj, kb][None, :, :] if ki == kj: ijb[idx, idx] -= wjb2[kj, kb] ijb -= wib[ki, kb][:, None, :] off = ki * nkpts + kj Hr2.write(off * idx_ijb.size + idx_ijb, ijb.ravel()) Woovvtmp = imds.Woovv.transpose(0, 1, 3, 2)[:, :, kshift] Hr2 -= 2. * ctf.einsum('IJijcb,JIjicb->IJijb', t2[:, :, kshift], Woovvtmp) Hr2 += ctf.einsum('IJijcb,IJijcb->IJijb', t2[:, :, kshift], Woovvtmp) return eom.amplitudes_to_vector(Hr1, Hr2)