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)
