def tdm_k2s(imds):
tdm = imds.tdm
e = imds.td_e
nk = len(tdm)
orb_order = []
offsets = numpy.cumsum((0, ) + imds.eri.nmo[:-1])
for k in range(nk):
orb_order.append(
numpy.arange(offsets[k], offsets[k] + imds.eri.nocc[k]))
for k in range(nk):
orb_order.append(
numpy.arange(offsets[k] + imds.eri.nocc[k],
offsets[k] + imds.eri.nmo[k]))
orb_order = numpy.concatenate(orb_order)
tdm = numpy.array(tdm)
# Dims: k_transfer, (xy), ki, roots, i, a
result = []
for k_transfer, tdm_t in enumerate(tdm):
fw, bw, _, _ = get_block_k_ix(imds.eri, k_transfer)
x, y = tdm_t
# X
x_s = []
for k1 in range(nk):
x_s.append([])
for k2 in range(nk):
if k2 == bw[k1]:
x_s[-1].append(x[k1])
else:
x_s[-1].append(numpy.zeros_like(x[k1]))
x_s = numpy.array(x_s).transpose(2, 0, 3, 1, 4)
s = x_s.shape
x_s = x_s.reshape((s[0], s[1] * s[2], s[3] * s[4]))
# Y
y_s = []
for k1 in range(nk):
y_s.append([])
for k2 in range(nk):
if k2 == fw[k1]:
y_s[-1].append(y[k1])
else:
y_s[-1].append(numpy.zeros_like(y[k1]))
y_s = numpy.array(y_s).transpose(2, 0, 3, 1, 4)
s = y_s.shape
y_s = y_s.reshape((s[0], s[1] * s[2], s[3] * s[4]))
result.append([x_s, y_s])
result = numpy.array(result)
result = result.transpose(0, 2, 1, 3, 4)
s = result.shape
result = result.reshape(s[0] * s[1], s[2], s[3], s[4])
# k_transfer, (xy), o, v
order = numpy.argsort(numpy.concatenate(tuple(e[i] for i in range(nk))))
# return result[order, ...]
return result[numpy.ix_(order, (0, 1), orb_order, orb_order)]
def convert_k2s(vectors, k, eri):
"""Converts vectors from k-representation to the supercell space by padding with zeros."""
nv, _, nk, nocc, nvirt = vectors.shape
# _ = 2
result = numpy.zeros((nv, _, nk, nk, nocc, nvirt), dtype=vectors.dtype)
r1, r2, _, _ = get_block_k_ix(eri, k)
for k1 in range(nk):
k2_x = r1[k1]
result[:, 0, k1, k2_x] = vectors[:, 0, k1]
k2_y = r2[k1]
result[:, 1, k1, k2_y] = vectors[:, 1, k1]
return result
def test_class(self):
"""Tests container behavior."""
model = kproxy.TDProxy(self.model_krks, "dft", [self.k, 1, 1], KRKS)
model.nroots = self.td_model_rks_supercell.nroots
assert not model.fast
for k in range(self.k):
model.kernel(k)
try:
testing.assert_allclose(model.eri.proxy_vind.ratio, self.call_ratios[k])
except Exception:
print("During efficiency check @k={:d} the following exception occurred:".format(k))
raise
model.eri.proxy_vind.reset()
o = v = 4
# Concatenate everything
ks = numpy.array(sum(([i] * len(model.e[i]) for i in range(self.k)), []))
vals = numpy.concatenate(tuple(model.e[i] for i in range(self.k))).real
vecs = numpy.concatenate(tuple(model.xy[i] for i in range(self.k)), axis=0)
# Obtain order
order = numpy.argsort(vals)
# Sort
vals = vals[order]
vecs = vecs[order]
ks = ks[order]
# Verify
testing.assert_allclose(vals, self.td_model_rks_supercell.e, atol=1e-7)
# Following makes sense only for non-degenerate roots
if self.k < 3:
for k in range(self.k):
# Prepare indexes
r1, r2, c1, c2 = krhf_slow.get_block_k_ix(model.eri, k)
r = k2k(r1, r2)
c = k2k(c1, c2)
# Select roots
selection = ks == k
vecs_ref = self.td_model_rks_supercell.xy[selection]
vecs_test = vecs[selection]
vecs_test_padded = numpy.zeros((len(vecs_test), 2 * self.k * self.k, o, v), dtype=vecs_test.dtype)
vecs_test_padded[:, c] = vecs_test.reshape((len(vecs_test), 2 * self.k, o, v))
vecs_test_padded = vecs_test_padded.reshape(vecs_ref.shape)
testing.assert_equal(vecs_test.shape, (self.k * o * v, 2, self.k, o, v))
testing.assert_equal(vecs_test_padded.shape, (self.k * o * v, 2, self.k, self.k, o, v))
assert_vectors_close(vecs_test_padded, vecs_ref, atol=1e-7)
def tdhf_primary_form(self, k):
"""
A primary form of TD matrices (full).
Args:
k (tuple, int): momentum transfer: either a pair of k-point indexes specifying the momentum transfer
vector or a single integer with the second index assuming the first index being zero;
Returns:
Output type: "full", and the corresponding matrix.
"""
r1, r2, c1, c2 = krhf_slow.get_block_k_ix(self, k)
(a, _), (_, b), (_, b_star), (a_star,
_) = self.proxy_response_ov_batch(
(r1, r1, r2, r2), (c1, c2, c1, c2))
return "full", numpy.block([[a, b], [-b_star.conj(), -a_star.conj()]])
def test_eig_kernel(self):
"""Tests container behavior (supercell)."""
model = ktd.TDRHF(self.model_krhf)
assert not model.fast
model.kernel()
o = v = 4
# Concatenate everything
ks = numpy.array(
sum(([i] * len(model.e[i]) for i in range(self.k)), []))
vals = numpy.concatenate(tuple(model.e[i] for i in range(self.k))).real
vecs = numpy.concatenate(tuple(model.xy[i] for i in range(self.k)),
axis=0)
# Obtain order
order = numpy.argsort(vals)
# Sort
vals = vals[order]
vecs = vecs[order]
ks = ks[order]
# Verify
testing.assert_allclose(vals, self.td_model_rhf_supercell.e, atol=1e-7)
for k in range(self.k):
# Prepare indexes
r1, r2, c1, c2 = ktd.get_block_k_ix(model.eri, k)
r = k2k(r1, r2)
c = k2k(c1, c2)
# Select roots
selection = ks == k
vecs_ref = self.td_model_rhf_supercell.xy[selection]
vecs_test = vecs[selection]
vecs_test_padded = numpy.zeros(
(len(vecs_test), 2 * self.k * self.k, o, v),
dtype=vecs_test.dtype)
vecs_test_padded[:, c] = vecs_test.reshape(
(len(vecs_test), 2 * self.k, o, v))
vecs_test_padded = vecs_test_padded.reshape(vecs_ref.shape)
testing.assert_equal(vecs_test.shape,
(self.k * o * v, 2, self.k, o, v))
testing.assert_equal(vecs_test_padded.shape,
(self.k * o * v, 2, self.k, self.k, o, v))
assert_vectors_close(vecs_test_padded, vecs_ref, atol=1e-7)
def retrieve_m_khf(eri, k):
"""Retrieves TDHF matrix directly (K-version)."""
r1, r2, c1, c2 = get_block_k_ix(eri, k)
d1 = eri.tdhf_diag(r1)
d2 = eri.tdhf_diag(r2)
m11 = d1 + 2 * eri["knmj", r1, c1] - eri["knjm", r1, c1]
m12 = 2 * eri["kjmn", r1, c2] - eri["kjnm", r1, c2]
m21 = 2 * eri["mnkj", r2, c1] - eri["mnjk", r2, c1]
m22 = d2 + 2 * eri["mjkn", r2, c2] - eri["mjnk", r2, c2]
m = numpy.array([[m11, m12], [-m21, -m22]])
return m.transpose((0, 2, 1, 3)).reshape(
(m.shape[0] * m.shape[2], m.shape[1] * m.shape[3])
)
def test_eri(self):
"""Tests all ERI implementations: with and without symmetries (supercell)."""
o = v = 4
for eri in (ktd.PhysERI, ktd.PhysERI4, ktd.PhysERI8):
if eri is not ktd.PhysERI8 or self.test8:
try:
e = eri(self.model_krhf)
if eri is ktd.PhysERI:
self.assertIn("__full_eri_k__", dir(e))
else:
self.assertNotIn("__full_eri_k__", dir(e))
self.assertNotIn("__full_eri__", dir(e))
s = (2 * self.k * self.k, 2 * self.k * self.k, o * v,
o * v)
m = numpy.zeros(s, dtype=complex)
for k in range(self.k):
# Prepare indexes
r1, r2, c1, c2 = ktd.get_block_k_ix(e, k)
r = k2k(r1, r2)
c = k2k(c1, c2)
# Build matrix
_m = e.tdhf_full_form(k)
# Assign the submatrix
m[numpy.ix_(r, c)] = _m.reshape(
(2 * self.k, o * v, 2 * self.k,
o * v)).transpose(0, 2, 1, 3)
m = m.transpose(0, 2, 1,
3).reshape(self.ref_m_supercell.shape)
testing.assert_allclose(self.ref_m_supercell,
m,
atol=1e-11)
except Exception:
print("When testing {} the following exception occurred:".
format(eri))
raise
def get_sigma_element(self, omega, p, eta, vir_sgn=1):
k, p = p
# Molecular implementation
# ------------------------
# tdm = self.tdm.sum(axis=1)
# evi = direct_sum('v-i->vi', self.td_e, self.o)
# eva = direct_sum('v+a->va', self.td_e, self.v)
# sigma = numpy.sum(tdm[:, :self.nocc, p] ** 2 / (omega + evi - 1j * eta))
# sigma += numpy.sum(tdm[:, self.nocc:, p] ** 2 / (omega - eva + vir_sgn * 1j * eta))
# return sigma
sigma = 0
for k_transfer, tdm_k in enumerate(self.tdm):
fw, bw, _, _ = get_block_k_ix(self.eri, k_transfer)
same = fw == bw
different = numpy.logical_not(same)
terms = []
if same.sum() > 0:
terms.append((tdm_k[0] + tdm_k[1], fw[same], fw[same]))
if different.sum() > 0:
terms.append((tdm_k[0], bw[different], fw[different]))
terms.append((tdm_k[1], fw[different], bw[different]))
for tdm_kx, ix_fw, ix_bw in terms:
k1, k2 = ix_bw[k], k
evi = direct_sum('v-i->vi', self.td_e[k_transfer], self.o[k1])
eva = direct_sum('v+a->va', self.td_e[k_transfer], self.v[k1])
sigma += numpy.sum(tdm_kx[k1][:, :self.nocc[k1], p]**2 /
(omega + evi - 1j * eta))
sigma += numpy.sum(tdm_kx[k1][:, self.nocc[k1]:, p]**2 /
(omega - eva + vir_sgn * 1j * eta))
return sigma
def construct_tdm(self):
# Indexes of td_x:
# - k_transfer
# - k
# - o: k
# - v: fw[k]
# Indexes of td_y:
# - k_transfer
# - k
# - o: k
# - v: bw[k]
# Original code:
# tdm_oo = einsum('vxia,ipaq->vxpq', td_xy, self["oovo"])
# tdm_ov = einsum('vxia,ipaq->vxpq', td_xy, self["oovv"])
# tdm_vv = einsum('vxia,ipaq->vxpq', td_xy, self["ovvv"])
# ERI k:
# ki, kp, ka=?, kq
# Now fw[kp] = kq, bw[kq] = kp -> bw[ki] = ka, fw[ka] = ki
# for x amplitudes, the transfer is bw[0] such that ov -> k, bw[k]
# for y amplitudes, the transfer is also fw[0] such that ov -> k, bw[k]
result = []
for k_transfer in range(self.nk):
xy_k = self.td_xy[k_transfer]
fw, bw, _, _ = get_block_k_ix(self.eri, k_transfer)
result.append([[], []])
for xy_kx, ix_fw, ix_bw, storage in (
(xy_k[:, 0], fw, bw, result[k_transfer][0]), # X
(xy_k[:, 1], bw, fw, result[k_transfer][1]), # Y
):
for kp in range(self.nk):
tdm_oo = tdm_ov = tdm_vv = 0
tdm_vo = 0
for ki in range(self.nk):
x = xy_kx[:, ki] * 2
tdm_oo = tdm_oo + einsum(
'via,ipaq->vpq', x, self["oovo", ki, kp, ix_fw[ki],
ix_bw[kp]])
tdm_ov = tdm_ov + einsum(
'via,ipaq->vpq', x, self["oovv", ki, kp, ix_fw[ki],
ix_bw[kp]])
tdm_vv = tdm_vv + einsum(
'via,ipaq->vpq', x, self["ovvv", ki, kp, ix_fw[ki],
ix_bw[kp]])
tdm_vo = tdm_vo + einsum(
'via,ipaq->vpq', x, self["ovvo", ki, kp, ix_fw[ki],
ix_bw[kp]])
tdm = numpy.concatenate(
(numpy.concatenate((tdm_oo, tdm_ov), axis=2),
numpy.concatenate((tdm_vo, tdm_vv), axis=2)),
axis=1,
)
storage.append(tdm)
# The output is the following:
# for each kp, kq pair two 3-tensors are given
# The last two indexes in each tensor correspond to kp, kq
# Given fw, bw, _, _ = get_block_k_ix(self.eri, (kp, kq)),
# The first index of the two tensors will correspond to tdhf.e[bw[0]], tdhf.e[fw[0]] correspondingly
return numpy.array(result)
