1#!/usr/bin/env python 2# Copyright 2017-2021 The PySCF Developers. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15# 16# Authors: James D. McClain 17# Timothy Berkelbach <tim.berkelbach@gmail.com> 18# 19 20 21import numpy 22from functools import reduce 23 24from pyscf import lib 25from pyscf.lib import logger 26from pyscf.pbc import scf 27from pyscf.cc import gccsd 28from pyscf.cc import ccsd 29from pyscf.pbc.mp.kmp2 import (get_frozen_mask, get_nmo, get_nocc, 30 padded_mo_coeff, padding_k_idx) # noqa 31from pyscf.pbc.cc import kintermediates as imdk 32from pyscf.lib.parameters import LOOSE_ZERO_TOL, LARGE_DENOM # noqa 33from pyscf.pbc.lib import kpts_helper 34 35DEBUG = False 36 37# 38# FIXME: When linear dependence is found in KHF and handled by function 39# pyscf.scf.addons.remove_linear_dep_, different k-point may have different 40# number of orbitals. 41# 42 43#einsum = numpy.einsum 44einsum = lib.einsum 45 46 47def energy(cc, t1, t2, eris): 48 nkpts, nocc, nvir = t1.shape 49 fock = eris.fock 50 eris_oovv = eris.oovv.copy() 51 e = 0.0 + 0j 52 for ki in range(nkpts): 53 e += einsum('ia,ia', fock[ki, :nocc, nocc:], t1[ki, :, :]) 54 t1t1 = numpy.zeros(shape=t2.shape, dtype=t2.dtype) 55 for ki in range(nkpts): 56 ka = ki 57 for kj in range(nkpts): 58 #kb = kj 59 t1t1[ki, kj, ka, :, :, :, :] = einsum('ia,jb->ijab', t1[ki, :, :], t1[kj, :, :]) 60 tau = t2 + 2 * t1t1 61 e += 0.25 * numpy.dot(tau.flatten(), eris_oovv.flatten()) 62 e /= nkpts 63 if abs(e.imag) > 1e-4: 64 logger.warn(cc, 'Non-zero imaginary part found in KCCSD energy %s', e) 65 return e.real 66 67 68def update_amps(cc, t1, t2, eris): 69 time0 = logger.process_clock(), logger.perf_counter() 70 log = logger.Logger(cc.stdout, cc.verbose) 71 nkpts, nocc, nvir = t1.shape 72 fock = eris.fock 73 mo_e_o = [e[:nocc] for e in eris.mo_energy] 74 mo_e_v = [e[nocc:] + cc.level_shift for e in eris.mo_energy] 75 76 # Get location of padded elements in occupied and virtual space 77 nonzero_opadding, nonzero_vpadding = padding_k_idx(cc, kind="split") 78 79 fov = fock[:, :nocc, nocc:].copy() 80 81 # Get the momentum conservation array 82 # Note: chemist's notation for momentum conserving t2(ki,kj,ka,kb), even though 83 # integrals are in physics notation 84 kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts) 85 86 tau = imdk.make_tau(cc, t2, t1, t1, kconserv) 87 88 Fvv = imdk.cc_Fvv(cc, t1, t2, eris, kconserv) 89 Foo = imdk.cc_Foo(cc, t1, t2, eris, kconserv) 90 Fov = imdk.cc_Fov(cc, t1, t2, eris, kconserv) 91 Woooo = imdk.cc_Woooo(cc, t1, t2, eris, kconserv) 92 Wvvvv = imdk.cc_Wvvvv(cc, t1, t2, eris, kconserv) 93 Wovvo = imdk.cc_Wovvo(cc, t1, t2, eris, kconserv) 94 95 # Move energy terms to the other side 96 for k in range(nkpts): 97 Foo[k][numpy.diag_indices(nocc)] -= mo_e_o[k] 98 Fvv[k][numpy.diag_indices(nvir)] -= mo_e_v[k] 99 100 eris_ovvo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nocc, nvir, nvir, nocc), dtype=t2.dtype) 101 eris_oovo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nocc, nocc, nvir, nocc), dtype=t2.dtype) 102 eris_vvvo = numpy.zeros(shape=(nkpts, nkpts, nkpts, nvir, nvir, nvir, nocc), dtype=t2.dtype) 103 for km, kb, ke in kpts_helper.loop_kkk(nkpts): 104 kj = kconserv[km, ke, kb] 105 # <mb||je> -> -<mb||ej> 106 eris_ovvo[km, kb, ke] = -eris.ovov[km, kb, kj].transpose(0, 1, 3, 2) 107 # <mn||je> -> -<mn||ej> 108 # let kb = kn as a dummy variable 109 eris_oovo[km, kb, ke] = -eris.ooov[km, kb, kj].transpose(0, 1, 3, 2) 110 # <ma||be> -> - <be||am>* 111 # let kj = ka as a dummy variable 112 kj = kconserv[km, ke, kb] 113 eris_vvvo[ke, kj, kb] = -eris.ovvv[km, kb, ke].transpose(2, 3, 1, 0).conj() 114 115 # T1 equation 116 t1new = numpy.zeros(shape=t1.shape, dtype=t1.dtype) 117 for ka in range(nkpts): 118 ki = ka 119 t1new[ka] += numpy.array(fov[ka, :, :]).conj() 120 t1new[ka] += einsum('ie,ae->ia', t1[ka], Fvv[ka]) 121 t1new[ka] += -einsum('ma,mi->ia', t1[ka], Foo[ka]) 122 for km in range(nkpts): 123 t1new[ka] += einsum('imae,me->ia', t2[ka, km, ka], Fov[km]) 124 t1new[ka] += -einsum('nf,naif->ia', t1[km], eris.ovov[km, ka, ki]) 125 for kn in range(nkpts): 126 ke = kconserv[km, ki, kn] 127 t1new[ka] += -0.5 * einsum('imef,maef->ia', t2[ki, km, ke], eris.ovvv[km, ka, ke]) 128 t1new[ka] += -0.5 * einsum('mnae,nmei->ia', t2[km, kn, ka], eris_oovo[kn, km, ke]) 129 130 # T2 equation 131 t2new = numpy.array(eris.oovv).conj() 132 for ki, kj, ka in kpts_helper.loop_kkk(nkpts): 133 # Chemist's notation for momentum conserving t2(ki,kj,ka,kb) 134 kb = kconserv[ki, ka, kj] 135 136 Ftmp = Fvv[kb] - 0.5 * einsum('mb,me->be', t1[kb], Fov[kb]) 137 tmp = einsum('ijae,be->ijab', t2[ki, kj, ka], Ftmp) 138 t2new[ki, kj, ka] += tmp 139 140 #t2new[ki,kj,kb] -= tmp.transpose(0,1,3,2) 141 Ftmp = Fvv[ka] - 0.5 * einsum('ma,me->ae', t1[ka], Fov[ka]) 142 tmp = einsum('ijbe,ae->ijab', t2[ki, kj, kb], Ftmp) 143 t2new[ki, kj, ka] -= tmp 144 145 Ftmp = Foo[kj] + 0.5 * einsum('je,me->mj', t1[kj], Fov[kj]) 146 tmp = einsum('imab,mj->ijab', t2[ki, kj, ka], Ftmp) 147 t2new[ki, kj, ka] -= tmp 148 149 #t2new[kj,ki,ka] += tmp.transpose(1,0,2,3) 150 Ftmp = Foo[ki] + 0.5 * einsum('ie,me->mi', t1[ki], Fov[ki]) 151 tmp = einsum('jmab,mi->ijab', t2[kj, ki, ka], Ftmp) 152 t2new[ki, kj, ka] += tmp 153 154 for km in range(nkpts): 155 # Wminj 156 # - km - kn + ka + kb = 0 157 # => kn = ka - km + kb 158 kn = kconserv[ka, km, kb] 159 t2new[ki, kj, ka] += 0.5 * einsum('mnab,mnij->ijab', tau[km, kn, ka], Woooo[km, kn, ki]) 160 ke = km 161 t2new[ki, kj, ka] += 0.5 * einsum('ijef,abef->ijab', tau[ki, kj, ke], Wvvvv[ka, kb, ke]) 162 163 # Wmbej 164 # - km - kb + ke + kj = 0 165 # => ke = km - kj + kb 166 ke = kconserv[km, kj, kb] 167 tmp = einsum('imae,mbej->ijab', t2[ki, km, ka], Wovvo[km, kb, ke]) 168 # - km - kb + ke + kj = 0 169 # => ke = km - kj + kb 170 # 171 # t[i,e] => ki = ke 172 # t[m,a] => km = ka 173 if km == ka and ke == ki: 174 tmp -= einsum('ie,ma,mbej->ijab', t1[ki], t1[km], eris_ovvo[km, kb, ke]) 175 t2new[ki, kj, ka] += tmp 176 t2new[ki, kj, kb] -= tmp.transpose(0, 1, 3, 2) 177 t2new[kj, ki, ka] -= tmp.transpose(1, 0, 2, 3) 178 t2new[kj, ki, kb] += tmp.transpose(1, 0, 3, 2) 179 180 ke = ki 181 tmp = einsum('ie,abej->ijab', t1[ki], eris_vvvo[ka, kb, ke]) 182 t2new[ki, kj, ka] += tmp 183 # P(ij) term 184 ke = kj 185 tmp = einsum('je,abei->ijab', t1[kj], eris_vvvo[ka, kb, ke]) 186 t2new[ki, kj, ka] -= tmp 187 188 km = ka 189 tmp = einsum('ma,mbij->ijab', t1[ka], eris.ovoo[km, kb, ki]) 190 t2new[ki, kj, ka] -= tmp 191 # P(ab) term 192 km = kb 193 tmp = einsum('mb,maij->ijab', t1[kb], eris.ovoo[km, ka, ki]) 194 t2new[ki, kj, ka] += tmp 195 196 for ki in range(nkpts): 197 ka = ki 198 # Remove zero/padded elements from denominator 199 eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype) 200 n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka]) 201 eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia] 202 t1new[ki] /= eia 203 204 kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts) 205 for ki, kj, ka in kpts_helper.loop_kkk(nkpts): 206 kb = kconserv[ki, ka, kj] 207 # For LARGE_DENOM, see t1new update above 208 eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype) 209 n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka]) 210 eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia] 211 212 ejb = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype) 213 n0_ovp_jb = numpy.ix_(nonzero_opadding[kj], nonzero_vpadding[kb]) 214 ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb] 215 eijab = eia[:, None, :, None] + ejb[:, None, :] 216 217 t2new[ki, kj, ka] /= eijab 218 219 time0 = log.timer_debug1('update t1 t2', *time0) 220 221 return t1new, t2new 222 223def spatial2spin(tx, orbspin, kconserv): 224 '''Convert T1/T2 of spatial orbital representation to T1/T2 of 225 spin-orbital representation 226 ''' 227 if isinstance(tx, numpy.ndarray) and tx.ndim == 3: 228 # KRCCSD t1 amplitudes 229 return spatial2spin((tx,tx), orbspin, kconserv) 230 elif isinstance(tx, numpy.ndarray) and tx.ndim == 7: 231 # KRCCSD t2 amplitudes 232 t2aa = numpy.zeros_like(tx) 233 nkpts = t2aa.shape[2] 234 for ki, kj, ka in kpts_helper.loop_kkk(nkpts): 235 kb = kconserv[ki,ka,kj] 236 t2aa[ki,kj,ka] = tx[ki,kj,ka] - tx[ki,kj,kb].transpose(0,1,3,2) 237 return spatial2spin((t2aa,tx,t2aa), orbspin, kconserv) 238 elif len(tx) == 2: # KUCCSD t1 239 t1a, t1b = tx 240 nocc_a, nvir_a = t1a.shape[1:] 241 nocc_b, nvir_b = t1b.shape[1:] 242 else: # KUCCSD t2 243 t2aa, t2ab, t2bb = tx 244 nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape[3:] 245 246 nkpts = len(orbspin) 247 nocc = nocc_a + nocc_b 248 nvir = nvir_a + nvir_b 249 idxoa = [numpy.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)] 250 idxob = [numpy.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)] 251 idxva = [numpy.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)] 252 idxvb = [numpy.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)] 253 254 if len(tx) == 2: # t1 255 t1 = numpy.zeros((nkpts,nocc,nvir), dtype=t1a.dtype) 256 for k in range(nkpts): 257 lib.takebak_2d(t1[k], t1a[k], idxoa[k], idxva[k]) 258 lib.takebak_2d(t1[k], t1b[k], idxob[k], idxvb[k]) 259 t1 = lib.tag_array(t1, orbspin=orbspin) 260 return t1 261 262 else: 263 t2 = numpy.zeros((nkpts,nkpts,nkpts,nocc**2,nvir**2), dtype=t2aa.dtype) 264 for ki, kj, ka in kpts_helper.loop_kkk(nkpts): 265 kb = kconserv[ki,ka,kj] 266 idxoaa = idxoa[ki][:,None] * nocc + idxoa[kj] 267 idxoab = idxoa[ki][:,None] * nocc + idxob[kj] 268 idxoba = idxob[kj][:,None] * nocc + idxoa[ki] 269 idxobb = idxob[ki][:,None] * nocc + idxob[kj] 270 idxvaa = idxva[ka][:,None] * nvir + idxva[kb] 271 idxvab = idxva[ka][:,None] * nvir + idxvb[kb] 272 idxvba = idxvb[kb][:,None] * nvir + idxva[ka] 273 idxvbb = idxvb[ka][:,None] * nvir + idxvb[kb] 274 tmp2aa = t2aa[ki,kj,ka].reshape(nocc_a*nocc_a,nvir_a*nvir_a) 275 tmp2bb = t2bb[ki,kj,ka].reshape(nocc_b*nocc_b,nvir_b*nvir_b) 276 tmp2ab = t2ab[ki,kj,ka].reshape(nocc_a*nocc_b,nvir_a*nvir_b) 277 lib.takebak_2d(t2[ki,kj,ka], tmp2aa, idxoaa.ravel() , idxvaa.ravel() ) 278 lib.takebak_2d(t2[ki,kj,ka], tmp2bb, idxobb.ravel() , idxvbb.ravel() ) 279 lib.takebak_2d(t2[ki,kj,ka], tmp2ab, idxoab.ravel() , idxvab.ravel() ) 280 lib.takebak_2d(t2[kj,ki,kb], tmp2ab, idxoba.T.ravel(), idxvba.T.ravel()) 281 282 abba = -tmp2ab 283 lib.takebak_2d(t2[ki,kj,kb], abba, idxoab.ravel() , idxvba.T.ravel()) 284 lib.takebak_2d(t2[kj,ki,ka], abba, idxoba.T.ravel(), idxvab.ravel() ) 285 t2 = t2.reshape(nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir) 286 t2 = lib.tag_array(t2, orbspin=orbspin) 287 return t2 288 289def spin2spatial(tx, orbspin, kconserv): 290 if tx.ndim == 3: # t1 291 nocc, nvir = tx.shape[1:] 292 else: 293 nocc, nvir = tx.shape[4:6] 294 nkpts = len(tx) 295 296 idxoa = [numpy.where(orbspin[k][:nocc] == 0)[0] for k in range(nkpts)] 297 idxob = [numpy.where(orbspin[k][:nocc] == 1)[0] for k in range(nkpts)] 298 idxva = [numpy.where(orbspin[k][nocc:] == 0)[0] for k in range(nkpts)] 299 idxvb = [numpy.where(orbspin[k][nocc:] == 1)[0] for k in range(nkpts)] 300 nocc_a = len(idxoa[0]) 301 nocc_b = len(idxob[0]) 302 nvir_a = len(idxva[0]) 303 nvir_b = len(idxvb[0]) 304 305 if tx.ndim == 3: # t1 306 t1a = numpy.zeros((nkpts,nocc_a,nvir_a), dtype=tx.dtype) 307 t1b = numpy.zeros((nkpts,nocc_b,nvir_b), dtype=tx.dtype) 308 for k in range(nkpts): 309 lib.take_2d(tx[k], idxoa[k], idxva[k], out=t1a[k]) 310 lib.take_2d(tx[k], idxob[k], idxvb[k], out=t1b[k]) 311 return t1a, t1b 312 313 else: 314 t2aa = numpy.zeros((nkpts,nkpts,nkpts,nocc_a,nocc_a,nvir_a,nvir_a), dtype=tx.dtype) 315 t2ab = numpy.zeros((nkpts,nkpts,nkpts,nocc_a,nocc_b,nvir_a,nvir_b), dtype=tx.dtype) 316 t2bb = numpy.zeros((nkpts,nkpts,nkpts,nocc_b,nocc_b,nvir_b,nvir_b), dtype=tx.dtype) 317 t2 = tx.reshape(nkpts,nkpts,nkpts,nocc**2,nvir**2) 318 for ki, kj, ka in kpts_helper.loop_kkk(nkpts): 319 kb = kconserv[ki,ka,kj] 320 idxoaa = idxoa[ki][:,None] * nocc + idxoa[kj] 321 idxoab = idxoa[ki][:,None] * nocc + idxob[kj] 322 idxobb = idxob[ki][:,None] * nocc + idxob[kj] 323 idxvaa = idxva[ka][:,None] * nvir + idxva[kb] 324 idxvab = idxva[ka][:,None] * nvir + idxvb[kb] 325 idxvbb = idxvb[ka][:,None] * nvir + idxvb[kb] 326 lib.take_2d(t2[ki,kj,ka], idxoaa.ravel(), idxvaa.ravel(), out=t2aa[ki,kj,ka]) 327 lib.take_2d(t2[ki,kj,ka], idxobb.ravel(), idxvbb.ravel(), out=t2bb[ki,kj,ka]) 328 lib.take_2d(t2[ki,kj,ka], idxoab.ravel(), idxvab.ravel(), out=t2ab[ki,kj,ka]) 329 return t2aa, t2ab, t2bb 330 331 332class GCCSD(gccsd.GCCSD): 333 def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None): 334 assert (isinstance(mf, scf.khf.KSCF)) 335 if not isinstance(mf, scf.kghf.KGHF): 336 mf = scf.addons.convert_to_ghf(mf) 337 self.kpts = mf.kpts 338 self.khelper = kpts_helper.KptsHelper(mf.cell, mf.kpts) 339 gccsd.GCCSD.__init__(self, mf, frozen, mo_coeff, mo_occ) 340 341 @property 342 def nkpts(self): 343 return len(self.kpts) 344 345 get_nocc = get_nocc 346 get_nmo = get_nmo 347 get_frozen_mask = get_frozen_mask 348 349 def dump_flags(self, verbose=None): 350 logger.info(self, '\n') 351 logger.info(self, '******** PBC CC flags ********') 352 gccsd.GCCSD.dump_flags(self, verbose) 353 return self 354 355 def init_amps(self, eris): 356 time0 = logger.process_clock(), logger.perf_counter() 357 nocc = self.nocc 358 nvir = self.nmo - nocc 359 nkpts = self.nkpts 360 mo_e_o = [eris.mo_energy[k][:nocc] for k in range(nkpts)] 361 mo_e_v = [eris.mo_energy[k][nocc:] for k in range(nkpts)] 362 t1 = numpy.zeros((nkpts, nocc, nvir), dtype=numpy.complex128) 363 t2 = numpy.zeros((nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir), dtype=numpy.complex128) 364 self.emp2 = 0 365 eris_oovv = eris.oovv.copy() 366 367 # Get location of padded elements in occupied and virtual space 368 nonzero_opadding, nonzero_vpadding = padding_k_idx(self, kind="split") 369 370 kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts) 371 for ki, kj, ka in kpts_helper.loop_kkk(nkpts): 372 kb = kconserv[ki, ka, kj] 373 # For LARGE_DENOM, see t1new update above 374 eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype) 375 n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka]) 376 eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia] 377 378 ejb = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype) 379 n0_ovp_jb = numpy.ix_(nonzero_opadding[kj], nonzero_vpadding[kb]) 380 ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb] 381 eijab = eia[:, None, :, None] + ejb[:, None, :] 382 383 t2[ki, kj, ka] = eris_oovv[ki, kj, ka] / eijab 384 385 t2 = numpy.conj(t2) 386 self.emp2 = 0.25 * numpy.einsum('pqrijab,pqrijab', t2, eris_oovv).real 387 self.emp2 /= nkpts 388 389 logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2.real) 390 logger.timer(self, 'init mp2', *time0) 391 return self.emp2, t1, t2 392 393 def ccsd(self, t1=None, t2=None, eris=None, **kwargs): 394 if eris is None: eris = self.ao2mo(self.mo_coeff) 395 e_corr, self.t1, self.t2 = ccsd.CCSD.ccsd(self, t1, t2, eris) 396 if getattr(eris, 'orbspin', None) is not None: 397 self.t1 = lib.tag_array(self.t1, orbspin=eris.orbspin) 398 self.t2 = lib.tag_array(self.t2, orbspin=eris.orbspin) 399 return e_corr, self.t1, self.t2 400 401 update_amps = update_amps 402 403 energy = energy 404 405 def ao2mo(self, mo_coeff=None): 406 nkpts = self.nkpts 407 nmo = self.nmo 408 mem_incore = nkpts**3 * nmo**4 * 8 / 1e6 409 mem_now = lib.current_memory()[0] 410 411 if (mem_incore + mem_now < self.max_memory) or self.mol.incore_anyway: 412 return _make_eris_incore(self, mo_coeff) 413 else: 414 raise NotImplementedError 415 416 def ccsd_t(self, t1=None, t2=None, eris=None): 417 from pyscf.pbc.cc import kccsd_t 418 if t1 is None: t1 = self.t1 419 if t2 is None: t2 = self.t2 420 if eris is None: eris = self.ao2mo(self.mo_coeff) 421 return kccsd_t.kernel(self, eris, t1, t2, self.verbose) 422 423 def amplitudes_to_vector(self, t1, t2): 424 return numpy.hstack((t1.ravel(), t2.ravel())) 425 426 def vector_to_amplitudes(self, vec, nmo=None, nocc=None): 427 if nocc is None: nocc = self.nocc 428 if nmo is None: nmo = self.nmo 429 nvir = nmo - nocc 430 nkpts = self.nkpts 431 nov = nkpts * nocc * nvir 432 t1 = vec[:nov].reshape(nkpts, nocc, nvir) 433 t2 = vec[nov:].reshape(nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir) 434 return t1, t2 435 436 def spatial2spin(self, tx, orbspin=None, kconserv=None): 437 if orbspin is None: 438 if getattr(self.mo_coeff[0], 'orbspin', None) is not None: 439 orbspin = [self.mo_coeff[k].orbspin[idx] 440 for k, idx in enumerate(self.get_frozen_mask())] 441 else: 442 orbspin = numpy.zeros((self.nkpts,self.nmo), dtype=int) 443 orbspin[:,1::2] = 1 444 if kconserv is None: 445 kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts) 446 return spatial2spin(tx, orbspin, kconserv) 447 448 def spin2spatial(self, tx, orbspin=None, kconserv=None): 449 if orbspin is None: 450 if getattr(self.mo_coeff[0], 'orbspin', None) is not None: 451 orbspin = [self.mo_coeff[k].orbspin[idx] 452 for k, idx in enumerate(self.get_frozen_mask())] 453 else: 454 orbspin = numpy.zeros((self.nkpts,self.nmo), dtype=int) 455 orbspin[:,1::2] = 1 456 if kconserv is None: 457 kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts) 458 return spin2spatial(tx, orbspin, kconserv) 459 460 def from_uccsd(self, t1, t2, orbspin=None): 461 return self.spatial2spin(t1, orbspin), self.spatial2spin(t2, orbspin) 462 463 def to_uccsd(self, t1, t2, orbspin=None): 464 return spin2spatial(t1, orbspin), spin2spatial(t2, orbspin) 465 466CCSD = KCCSD = KGCCSD = GCCSD 467 468 469def _make_eris_incore(cc, mo_coeff=None): 470 from pyscf.pbc import tools 471 from pyscf.pbc.cc.ccsd import _adjust_occ 472 473 log = logger.Logger(cc.stdout, cc.verbose) 474 cput0 = (logger.process_clock(), logger.perf_counter()) 475 eris = gccsd._PhysicistsERIs() 476 cell = cc._scf.cell 477 kpts = cc.kpts 478 nkpts = cc.nkpts 479 nocc = cc.nocc 480 nmo = cc.nmo 481 eris.nocc = nocc 482 483 #if any(nocc != numpy.count_nonzero(cc._scf.mo_occ[k] > 0) for k in range(nkpts)): 484 # raise NotImplementedError('Different occupancies found for different k-points') 485 486 if mo_coeff is None: 487 mo_coeff = cc.mo_coeff 488 489 nao = mo_coeff[0].shape[0] 490 dtype = mo_coeff[0].dtype 491 492 moidx = get_frozen_mask(cc) 493 nocc_per_kpt = numpy.asarray(get_nocc(cc, per_kpoint=True)) 494 nmo_per_kpt = numpy.asarray(get_nmo(cc, per_kpoint=True)) 495 496 padded_moidx = [] 497 for k in range(nkpts): 498 kpt_nocc = nocc_per_kpt[k] 499 kpt_nvir = nmo_per_kpt[k] - kpt_nocc 500 kpt_padded_moidx = numpy.concatenate((numpy.ones(kpt_nocc, dtype=numpy.bool), 501 numpy.zeros(nmo - kpt_nocc - kpt_nvir, dtype=numpy.bool), 502 numpy.ones(kpt_nvir, dtype=numpy.bool))) 503 padded_moidx.append(kpt_padded_moidx) 504 505 eris.mo_coeff = [] 506 eris.orbspin = [] 507 # Generate the molecular orbital coefficients with the frozen orbitals masked. 508 # Each MO is tagged with orbspin, a list of 0's and 1's that give the overall 509 # spin of each MO. 510 # 511 # Here we will work with two index arrays; one is for our original (small) moidx 512 # array while the next is for our new (large) padded array. 513 for k in range(nkpts): 514 kpt_moidx = moidx[k] 515 kpt_padded_moidx = padded_moidx[k] 516 517 mo = numpy.zeros((nao, nmo), dtype=dtype) 518 mo[:, kpt_padded_moidx] = mo_coeff[k][:, kpt_moidx] 519 if getattr(mo_coeff[k], 'orbspin', None) is not None: 520 orbspin_dtype = mo_coeff[k].orbspin[kpt_moidx].dtype 521 orbspin = numpy.zeros(nmo, dtype=orbspin_dtype) 522 orbspin[kpt_padded_moidx] = mo_coeff[k].orbspin[kpt_moidx] 523 mo = lib.tag_array(mo, orbspin=orbspin) 524 eris.orbspin.append(orbspin) 525 # FIXME: What if the user freezes all up spin orbitals in 526 # an RHF calculation? The number of electrons will still be 527 # even. 528 else: # guess orbital spin - assumes an RHF calculation 529 assert (numpy.count_nonzero(kpt_moidx) % 2 == 0) 530 orbspin = numpy.zeros(mo.shape[1], dtype=int) 531 orbspin[1::2] = 1 532 mo = lib.tag_array(mo, orbspin=orbspin) 533 eris.orbspin.append(orbspin) 534 eris.mo_coeff.append(mo) 535 536 # Re-make our fock MO matrix elements from density and fock AO 537 dm = cc._scf.make_rdm1(cc.mo_coeff, cc.mo_occ) 538 with lib.temporary_env(cc._scf, exxdiv=None): 539 # _scf.exxdiv affects eris.fock. HF exchange correction should be 540 # excluded from the Fock matrix. 541 vhf = cc._scf.get_veff(cell, dm) 542 fockao = cc._scf.get_hcore() + vhf 543 eris.fock = numpy.asarray([reduce(numpy.dot, (mo.T.conj(), fockao[k], mo)) 544 for k, mo in enumerate(eris.mo_coeff)]) 545 eris.e_hf = cc._scf.energy_tot(dm=dm, vhf=vhf) 546 547 eris.mo_energy = [eris.fock[k].diagonal().real for k in range(nkpts)] 548 # Add HFX correction in the eris.mo_energy to improve convergence in 549 # CCSD iteration. It is useful for the 2D systems since their occupied and 550 # the virtual orbital energies may overlap which may lead to numerical 551 # issue in the CCSD iterations. 552 # FIXME: Whether to add this correction for other exxdiv treatments? 553 # Without the correction, MP2 energy may be largely off the correct value. 554 madelung = tools.madelung(cell, kpts) 555 eris.mo_energy = [_adjust_occ(mo_e, nocc, -madelung) 556 for k, mo_e in enumerate(eris.mo_energy)] 557 558 # Get location of padded elements in occupied and virtual space. 559 nocc_per_kpt = get_nocc(cc, per_kpoint=True) 560 nonzero_padding = padding_k_idx(cc, kind="joint") 561 562 # Check direct and indirect gaps for possible issues with CCSD convergence. 563 mo_e = [eris.mo_energy[kp][nonzero_padding[kp]] for kp in range(nkpts)] 564 mo_e = numpy.sort([y for x in mo_e for y in x]) # Sort de-nested array 565 gap = mo_e[numpy.sum(nocc_per_kpt)] - mo_e[numpy.sum(nocc_per_kpt)-1] 566 if gap < 1e-5: 567 logger.warn(cc, 'HOMO-LUMO gap %s too small for KCCSD. ' 568 'May cause issues in convergence.', gap) 569 570 kconserv = kpts_helper.get_kconserv(cell, kpts) 571 if getattr(mo_coeff[0], 'orbspin', None) is None: 572 # The bottom nao//2 coefficients are down (up) spin while the top are up (down). 573 mo_a_coeff = [mo[:nao // 2] for mo in eris.mo_coeff] 574 mo_b_coeff = [mo[nao // 2:] for mo in eris.mo_coeff] 575 576 eri = numpy.empty((nkpts, nkpts, nkpts, nmo, nmo, nmo, nmo), dtype=numpy.complex128) 577 fao2mo = cc._scf.with_df.ao2mo 578 for kp, kq, kr in kpts_helper.loop_kkk(nkpts): 579 ks = kconserv[kp, kq, kr] 580 eri_kpt = fao2mo( 581 (mo_a_coeff[kp], mo_a_coeff[kq], mo_a_coeff[kr], mo_a_coeff[ks]), 582 (kpts[kp], kpts[kq], kpts[kr], kpts[ks]), 583 compact=False) 584 eri_kpt += fao2mo( 585 (mo_b_coeff[kp], mo_b_coeff[kq], mo_b_coeff[kr], mo_b_coeff[ks]), 586 (kpts[kp], kpts[kq], kpts[kr], kpts[ks]), 587 compact=False) 588 eri_kpt += fao2mo( 589 (mo_a_coeff[kp], mo_a_coeff[kq], mo_b_coeff[kr], mo_b_coeff[ks]), 590 (kpts[kp], kpts[kq], kpts[kr], kpts[ks]), 591 compact=False) 592 eri_kpt += fao2mo( 593 (mo_b_coeff[kp], mo_b_coeff[kq], mo_a_coeff[kr], mo_a_coeff[ks]), 594 (kpts[kp], kpts[kq], kpts[kr], kpts[ks]), 595 compact=False) 596 597 eri_kpt = eri_kpt.reshape(nmo, nmo, nmo, nmo) 598 eri[kp, kq, kr] = eri_kpt 599 else: 600 mo_a_coeff = [mo[:nao // 2] + mo[nao // 2:] for mo in eris.mo_coeff] 601 602 eri = numpy.empty((nkpts, nkpts, nkpts, nmo, nmo, nmo, nmo), dtype=numpy.complex128) 603 fao2mo = cc._scf.with_df.ao2mo 604 for kp, kq, kr in kpts_helper.loop_kkk(nkpts): 605 ks = kconserv[kp, kq, kr] 606 eri_kpt = fao2mo( 607 (mo_a_coeff[kp], mo_a_coeff[kq], mo_a_coeff[kr], mo_a_coeff[ks]), 608 (kpts[kp], kpts[kq], kpts[kr], kpts[ks]), 609 compact=False) 610 611 eri_kpt[(eris.orbspin[kp][:, None] != eris.orbspin[kq]).ravel()] = 0 612 eri_kpt[:, (eris.orbspin[kr][:, None] != eris.orbspin[ks]).ravel()] = 0 613 eri_kpt = eri_kpt.reshape(nmo, nmo, nmo, nmo) 614 eri[kp, kq, kr] = eri_kpt 615 616 # Check some antisymmetrized properties of the integrals 617 if DEBUG: 618 check_antisymm_3412(cc, cc.kpts, eri) 619 620 # Antisymmetrizing (pq|rs)-(ps|rq), where the latter integral is equal to 621 # (rq|ps); done since we aren't tracking the kpoint of orbital 's' 622 eri = eri - eri.transpose(2, 1, 0, 5, 4, 3, 6) 623 # Chemist -> physics notation 624 eri = eri.transpose(0, 2, 1, 3, 5, 4, 6) 625 626 # Set the various integrals 627 eris.dtype = eri.dtype 628 eris.oooo = eri[:, :, :, :nocc, :nocc, :nocc, :nocc].copy() / nkpts 629 eris.ooov = eri[:, :, :, :nocc, :nocc, :nocc, nocc:].copy() / nkpts 630 eris.ovoo = eri[:, :, :, :nocc, nocc:, :nocc, :nocc].copy() / nkpts 631 eris.oovv = eri[:, :, :, :nocc, :nocc, nocc:, nocc:].copy() / nkpts 632 eris.ovov = eri[:, :, :, :nocc, nocc:, :nocc, nocc:].copy() / nkpts 633 eris.ovvv = eri[:, :, :, :nocc, nocc:, nocc:, nocc:].copy() / nkpts 634 eris.vvvv = eri[:, :, :, nocc:, nocc:, nocc:, nocc:].copy() / nkpts 635 636 log.timer('CCSD integral transformation', *cput0) 637 return eris 638 639 640def check_antisymm_3412(cc, kpts, integrals): 641 kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts) 642 nkpts = len(kpts) 643 diff = 0.0 644 for kp, kq, kr in kpts_helper.loop_kkk(nkpts): 645 ks = kconserv[kp, kr, kq] 646 for p in range(integrals.shape[3]): 647 for q in range(integrals.shape[4]): 648 for r in range(integrals.shape[5]): 649 for s in range(integrals.shape[6]): 650 pqrs = integrals[kp, kq, kr, p, q, r, s] 651 rspq = integrals[kq, kp, kr, q, p, r, s] 652 cdiff = numpy.linalg.norm(pqrs - rspq).real 653 if diff > 1e-5: 654 print("AS diff = %.15g" % cdiff, pqrs, rspq, kp, kq, kr, ks, p, q, r, s) 655 diff = max(diff, cdiff) 656 print("antisymmetrization : max diff = %.15g" % diff) 657 if diff > 1e-5: 658 print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.") 659 return diff 660 661 662def check_antisymm_12(cc, kpts, integrals): 663 kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts) 664 nkpts = len(kpts) 665 diff = 0.0 666 for kp, kq, kr in kpts_helper.loop_kkk(nkpts): 667 ks = kconserv[kp, kr, kq] 668 for p in range(integrals.shape[3]): 669 for q in range(integrals.shape[4]): 670 for r in range(integrals.shape[5]): 671 for s in range(integrals.shape[6]): 672 pqrs = integrals[kp, kq, kr, p, q, r, s] 673 qprs = integrals[kq, kp, kr, q, p, r, s] 674 cdiff = numpy.linalg.norm(pqrs + qprs).real 675 if diff > 1e-5: 676 print("AS diff = %.15g" % cdiff, pqrs, qprs, kp, kq, kr, ks, p, q, r, s) 677 diff = max(diff, cdiff) 678 print("antisymmetrization : max diff = %.15g" % diff) 679 if diff > 1e-5: 680 print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.") 681 682 683def check_antisymm_34(cc, kpts, integrals): 684 kconserv = kpts_helper.get_kconserv(cc._scf.cell, cc.kpts) 685 nkpts = len(kpts) 686 diff = 0.0 687 for kp, kq, kr in kpts_helper.loop_kkk(nkpts): 688 ks = kconserv[kp, kr, kq] 689 for p in range(integrals.shape[3]): 690 for q in range(integrals.shape[4]): 691 for r in range(integrals.shape[5]): 692 for s in range(integrals.shape[6]): 693 pqrs = integrals[kp, kq, kr, p, q, r, s] 694 pqsr = integrals[kp, kq, ks, p, q, s, r] 695 cdiff = numpy.linalg.norm(pqrs + pqsr).real 696 if diff > 1e-5: 697 print("AS diff = %.15g" % cdiff, pqrs, pqsr, kp, kq, kr, ks, p, q, r, s) 698 diff = max(diff, cdiff) 699 print("antisymmetrization : max diff = %.15g" % diff) 700 if diff > 1e-5: 701 print("Energy cutoff (or cell.mesh) is not enough to converge AO integrals.") 702 703imd = imdk 704class _IMDS: 705 # Identical to molecular rccsd_slow 706 def __init__(self, cc): 707 self.verbose = cc.verbose 708 self.stdout = cc.stdout 709 self.t1 = cc.t1 710 self.t2 = cc.t2 711 self.eris = cc.eris 712 self.kconserv = cc.khelper.kconserv 713 self.made_ip_imds = False 714 self.made_ea_imds = False 715 self._made_shared_2e = False 716 self._fimd = None 717 718 def _make_shared_1e(self): 719 cput0 = (logger.process_clock(), logger.perf_counter()) 720 log = logger.Logger(self.stdout, self.verbose) 721 722 t1,t2,eris = self.t1, self.t2, self.eris 723 kconserv = self.kconserv 724 self.Loo = imd.Loo(t1,t2,eris,kconserv) 725 self.Lvv = imd.Lvv(t1,t2,eris,kconserv) 726 self.Fov = imd.cc_Fov(t1,t2,eris,kconserv) 727 728 log.timer('EOM-CCSD shared one-electron intermediates', *cput0) 729 730 def _make_shared_2e(self): 731 cput0 = (logger.process_clock(), logger.perf_counter()) 732 log = logger.Logger(self.stdout, self.verbose) 733 734 t1,t2,eris = self.t1, self.t2, self.eris 735 kconserv = self.kconserv 736 737 # TODO: check whether to hold Wovov Wovvo in memory 738 if self._fimd is None: 739 self._fimd = lib.H5TmpFile() 740 nkpts, nocc, nvir = t1.shape 741 self._fimd.create_dataset('ovov', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), t1.dtype.char) 742 self._fimd.create_dataset('ovvo', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), t1.dtype.char) 743 744 # 2 virtuals 745 self.Wovov = imd.Wovov(t1,t2,eris,kconserv, self._fimd['ovov']) 746 self.Wovvo = imd.Wovvo(t1,t2,eris,kconserv, self._fimd['ovvo']) 747 self.Woovv = eris.oovv 748 749 log.timer('EOM-CCSD shared two-electron intermediates', *cput0) 750 751 def make_ip(self, ip_partition=None): 752 self._make_shared_1e() 753 if self._made_shared_2e is False and ip_partition != 'mp': 754 self._make_shared_2e() 755 self._made_shared_2e = True 756 757 cput0 = (logger.process_clock(), logger.perf_counter()) 758 log = logger.Logger(self.stdout, self.verbose) 759 760 t1,t2,eris = self.t1, self.t2, self.eris 761 kconserv = self.kconserv 762 763 nkpts, nocc, nvir = t1.shape 764 self._fimd.create_dataset('oooo', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nocc), t1.dtype.char) 765 self._fimd.create_dataset('ooov', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), t1.dtype.char) 766 self._fimd.create_dataset('ovoo', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nocc), t1.dtype.char) 767 768 # 0 or 1 virtuals 769 if ip_partition != 'mp': 770 self.Woooo = imd.Woooo(t1,t2,eris,kconserv, self._fimd['oooo']) 771 self.Wooov = imd.Wooov(t1,t2,eris,kconserv, self._fimd['ooov']) 772 self.Wovoo = imd.Wovoo(t1,t2,eris,kconserv, self._fimd['ovoo']) 773 self.made_ip_imds = True 774 log.timer('EOM-CCSD IP intermediates', *cput0) 775 776 def make_ea(self, ea_partition=None): 777 self._make_shared_1e() 778 if self._made_shared_2e is False and ea_partition != 'mp': 779 self._make_shared_2e() 780 self._made_shared_2e = True 781 782 cput0 = (logger.process_clock(), logger.perf_counter()) 783 log = logger.Logger(self.stdout, self.verbose) 784 785 t1,t2,eris = self.t1, self.t2, self.eris 786 kconserv = self.kconserv 787 788 nkpts, nocc, nvir = t1.shape 789 self._fimd.create_dataset('vovv', (nkpts,nkpts,nkpts,nvir,nocc,nvir,nvir), t1.dtype.char) 790 self._fimd.create_dataset('vvvo', (nkpts,nkpts,nkpts,nvir,nvir,nvir,nocc), t1.dtype.char) 791 self._fimd.create_dataset('vvvv', (nkpts,nkpts,nkpts,nvir,nvir,nvir,nvir), t1.dtype.char) 792 793 # 3 or 4 virtuals 794 self.Wvovv = imd.Wvovv(t1,t2,eris,kconserv, self._fimd['vovv']) 795 if ea_partition == 'mp' and numpy.all(t1 == 0): 796 self.Wvvvo = imd.Wvvvo(t1,t2,eris,kconserv, self._fimd['vvvo']) 797 else: 798 self.Wvvvv = imd.Wvvvv(t1,t2,eris,kconserv, self._fimd['vvvv']) 799 self.Wvvvo = imd.Wvvvo(t1,t2,eris,kconserv,self.Wvvvv, self._fimd['vvvo']) 800 self.made_ea_imds = True 801 log.timer('EOM-CCSD EA intermediates', *cput0) 802 803 804scf.kghf.KGHF.CCSD = lib.class_as_method(KGCCSD) 805 806 807if __name__ == '__main__': 808 from pyscf.pbc import gto 809 810 cell = gto.Cell() 811 cell.atom=''' 812 C 0.000000000000 0.000000000000 0.000000000000 813 C 1.685068664391 1.685068664391 1.685068664391 814 ''' 815 cell.basis = 'gth-szv' 816 cell.pseudo = 'gth-pade' 817 cell.a = ''' 818 0.000000000, 3.370137329, 3.370137329 819 3.370137329, 0.000000000, 3.370137329 820 3.370137329, 3.370137329, 0.000000000''' 821 cell.unit = 'B' 822 cell.verbose = 5 823 cell.build() 824 825 # Running HF and CCSD with 1x1x2 Monkhorst-Pack k-point mesh 826 kmf = scf.KRHF(cell, kpts=cell.make_kpts([1,1,2]), exxdiv=None) 827 ehf = kmf.kernel() 828 kmf = scf.addons.convert_to_ghf(kmf) 829 830 mycc = KGCCSD(kmf) 831 ecc, t1, t2 = mycc.kernel() 832 print(ecc - -0.155298393321855) 833