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)