1#!/usr/bin/env python 2# Copyright 2014-2020 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# Author: Qiming Sun <osirpt.sun@gmail.com> 17# 18 19import warnings 20import numpy 21from pyscf import lib 22from pyscf import ao2mo 23from pyscf.ao2mo import _ao2mo 24from pyscf.ao2mo.incore import iden_coeffs, _conc_mos 25from pyscf.pbc.df.df_jk import zdotNN, zdotNC 26from pyscf.pbc.df.fft_ao2mo import _format_kpts, _iskconserv 27from pyscf.pbc.lib import kpts_helper 28from pyscf.pbc.lib.kpts_helper import is_zero, gamma_point, unique 29from pyscf import __config__ 30 31 32def get_eri(mydf, kpts=None, 33 compact=getattr(__config__, 'pbc_df_ao2mo_get_eri_compact', True)): 34 if mydf._cderi is None: 35 mydf.build() 36 37 cell = mydf.cell 38 nao = cell.nao_nr() 39 kptijkl = _format_kpts(kpts) 40 if not _iskconserv(cell, kptijkl): 41 lib.logger.warn(cell, 'df_ao2mo: momentum conservation not found in ' 42 'the given k-points %s', kptijkl) 43 return numpy.zeros((nao,nao,nao,nao)) 44 45 kpti, kptj, kptk, kptl = kptijkl 46 nao_pair = nao * (nao+1) // 2 47 max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6) 48 49#################### 50# gamma point, the integral is real and with s4 symmetry 51 if gamma_point(kptijkl): 52 eriR = numpy.zeros((nao_pair,nao_pair)) 53 for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, True): 54 lib.ddot(LpqR.T, LpqR, sign, eriR, 1) 55 LpqR = LpqI = None 56 if not compact: 57 eriR = ao2mo.restore(1, eriR, nao).reshape(nao**2,-1) 58 return eriR 59 60 elif is_zero(kpti-kptk) and is_zero(kptj-kptl): 61 eriR = numpy.zeros((nao*nao,nao*nao)) 62 eriI = numpy.zeros((nao*nao,nao*nao)) 63 for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False): 64 zdotNN(LpqR.T, LpqI.T, LpqR, LpqI, sign, eriR, eriI, 1) 65 LpqR = LpqI = None 66 return eriR + eriI*1j 67 68#################### 69# (kpt) i == j == k == l != 0 70# 71# (kpt) i == l && j == k && i != j && j != k => 72# both vbar and ovlp are zero. It corresponds to the exchange integral. 73# 74# complex integrals, N^4 elements 75 elif is_zero(kpti-kptl) and is_zero(kptj-kptk): 76 eriR = numpy.zeros((nao*nao,nao*nao)) 77 eriI = numpy.zeros((nao*nao,nao*nao)) 78 for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False): 79 zdotNC(LpqR.T, LpqI.T, LpqR, LpqI, sign, eriR, eriI, 1) 80 LpqR = LpqI = None 81# transpose(0,1,3,2) because 82# j == k && i == l => 83# (L|ij).transpose(0,2,1).conj() = (L^*|ji) = (L^*|kl) => (M|kl) 84 eri = lib.transpose((eriR+eriI*1j).reshape(-1,nao,nao), axes=(0,2,1)) 85 return eri.reshape(nao**2,-1) 86 87#################### 88# aosym = s1, complex integrals 89# 90# kpti == kptj => kptl == kptk 91# If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave 92# vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition. 93# So kptl/b - kptk/b must be -1 < k/b < 1. 94# 95 else: 96 eriR = numpy.zeros((nao*nao,nao*nao)) 97 eriI = numpy.zeros((nao*nao,nao*nao)) 98 blksize = int(max_memory*.4e6/16/nao**2) 99 for (LpqR, LpqI, sign), (LrsR, LrsI, sign1) in \ 100 lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False, blksize), 101 mydf.sr_loop(kptijkl[2:], max_memory, False, blksize)): 102 zdotNN(LpqR.T, LpqI.T, LrsR, LrsI, sign, eriR, eriI, 1) 103 LpqR = LpqI = LrsR = LrsI = None 104 return eriR + eriI*1j 105 106 107def general(mydf, mo_coeffs, kpts=None, 108 compact=getattr(__config__, 'pbc_df_ao2mo_general_compact', True)): 109 warn_pbc2d_eri(mydf) 110 if mydf._cderi is None: 111 mydf.build() 112 113 cell = mydf.cell 114 kptijkl = _format_kpts(kpts) 115 kpti, kptj, kptk, kptl = kptijkl 116 if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2: 117 mo_coeffs = (mo_coeffs,) * 4 118 if not _iskconserv(cell, kptijkl): 119 lib.logger.warn(cell, 'df_ao2mo: momentum conservation not found in ' 120 'the given k-points %s', kptijkl) 121 return numpy.zeros([mo.shape[1] for mo in mo_coeffs]) 122 123 all_real = not any(numpy.iscomplexobj(mo) for mo in mo_coeffs) 124 max_memory = max(2000, (mydf.max_memory - lib.current_memory()[0])) 125 126#################### 127# gamma point, the integral is real and with s4 symmetry 128 if gamma_point(kptijkl) and all_real: 129 ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact) 130 klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact) 131 eri_mo = numpy.zeros((nij_pair,nkl_pair)) 132 sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and 133 iden_coeffs(mo_coeffs[1], mo_coeffs[3])) 134 ijR = klR = None 135 for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, True): 136 ijR, klR = _dtrans(LpqR, ijR, ijmosym, moij, ijslice, 137 LpqR, klR, klmosym, mokl, klslice, sym) 138 lib.ddot(ijR.T, klR, sign, eri_mo, 1) 139 LpqR = LpqI = None 140 return eri_mo 141 142 elif is_zero(kpti-kptk) and is_zero(kptj-kptl): 143 mo_coeffs = _mo_as_complex(mo_coeffs) 144 nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] 145 nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3])[1:] 146 eri_mo = numpy.zeros((nij_pair,nkl_pair), dtype=numpy.complex128) 147 sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and 148 iden_coeffs(mo_coeffs[1], mo_coeffs[3])) 149 150 zij = zkl = None 151 for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False): 152 buf = LpqR+LpqI*1j 153 zij, zkl = _ztrans(buf, zij, moij, ijslice, 154 buf, zkl, mokl, klslice, sym) 155 lib.dot(zij.T, zkl, sign, eri_mo, 1) 156 LpqR = LpqI = buf = None 157 return eri_mo 158 159#################### 160# (kpt) i == j == k == l != 0 161# (kpt) i == l && j == k && i != j && j != k => 162# 163 elif is_zero(kpti-kptl) and is_zero(kptj-kptk): 164 mo_coeffs = _mo_as_complex(mo_coeffs) 165 nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] 166 nlk_pair, molk, lkslice = _conc_mos(mo_coeffs[3], mo_coeffs[2])[1:] 167 eri_mo = numpy.zeros((nij_pair,nlk_pair), dtype=numpy.complex128) 168 sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[3]) and 169 iden_coeffs(mo_coeffs[1], mo_coeffs[2])) 170 171 zij = zlk = None 172 for LpqR, LpqI, sign in mydf.sr_loop(kptijkl[:2], max_memory, False): 173 buf = LpqR+LpqI*1j 174 zij, zlk = _ztrans(buf, zij, moij, ijslice, 175 buf, zlk, molk, lkslice, sym) 176 lib.dot(zij.T, zlk.conj(), sign, eri_mo, 1) 177 LpqR = LpqI = buf = None 178 nmok = mo_coeffs[2].shape[1] 179 nmol = mo_coeffs[3].shape[1] 180 eri_mo = lib.transpose(eri_mo.reshape(-1,nmol,nmok), axes=(0,2,1)) 181 return eri_mo.reshape(nij_pair,nlk_pair) 182 183#################### 184# aosym = s1, complex integrals 185# 186# If kpti == kptj, (kptl-kptk)*a has to be multiples of 2pi because of the wave 187# vector symmetry. k is a fraction of reciprocal basis, 0 < k/b < 1, by definition. 188# So kptl/b - kptk/b must be -1 < k/b < 1. => kptl == kptk 189# 190 else: 191 mo_coeffs = _mo_as_complex(mo_coeffs) 192 nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1])[1:] 193 nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3])[1:] 194 nao = mo_coeffs[0].shape[0] 195 eri_mo = numpy.zeros((nij_pair,nkl_pair), dtype=numpy.complex128) 196 197 blksize = int(min(max_memory*.3e6/16/nij_pair, 198 max_memory*.3e6/16/nkl_pair, 199 max_memory*.3e6/16/nao**2)) 200 zij = zkl = None 201 for (LpqR, LpqI, sign), (LrsR, LrsI, sign1) in \ 202 lib.izip(mydf.sr_loop(kptijkl[:2], max_memory, False, blksize), 203 mydf.sr_loop(kptijkl[2:], max_memory, False, blksize)): 204 zij, zkl = _ztrans(LpqR+LpqI*1j, zij, moij, ijslice, 205 LrsR+LrsI*1j, zkl, mokl, klslice, False) 206 lib.dot(zij.T, zkl, sign, eri_mo, 1) 207 LpqR = LpqI = LrsR = LrsI = None 208 return eri_mo 209 210def ao2mo_7d(mydf, mo_coeff_kpts, kpts=None, factor=1, out=None): 211 cell = mydf.cell 212 if kpts is None: 213 kpts = mydf.kpts 214 nkpts = len(kpts) 215 216 if isinstance(mo_coeff_kpts, numpy.ndarray) and mo_coeff_kpts.ndim == 3: 217 mo_coeff_kpts = [mo_coeff_kpts] * 4 218 else: 219 mo_coeff_kpts = list(mo_coeff_kpts) 220 221 # Shape of the orbitals can be different on different k-points. The 222 # orbital coefficients must be formatted (padded by zeros) so that the 223 # shape of the orbital coefficients are the same on all k-points. This can 224 # be achieved by calling pbc.mp.kmp2.padded_mo_coeff function 225 nmoi, nmoj, nmok, nmol = [x.shape[2] for x in mo_coeff_kpts] 226 eri_shape = (nkpts, nkpts, nkpts, nmoi, nmoj, nmok, nmol) 227 if gamma_point(kpts): 228 dtype = numpy.result_type(*mo_coeff_kpts) 229 else: 230 dtype = numpy.complex128 231 232 if out is None: 233 out = numpy.empty(eri_shape, dtype=dtype) 234 else: 235 assert(out.shape == eri_shape) 236 237 kptij_lst = numpy.array([(ki, kj) for ki in kpts for kj in kpts]) 238 kptis_lst = kptij_lst[:,0] 239 kptjs_lst = kptij_lst[:,1] 240 kpt_ji = kptjs_lst - kptis_lst 241 uniq_kpts, uniq_index, uniq_inverse = unique(kpt_ji) 242 243 nao = cell.nao_nr() 244 max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]-nao**4*16/1e6) * .5 245 246 tao = [] 247 ao_loc = None 248 kconserv = kpts_helper.get_kconserv(cell, kpts) 249 for uniq_id, kpt in enumerate(uniq_kpts): 250 adapted_ji_idx = numpy.where(uniq_inverse == uniq_id)[0] 251 252 for ji, ji_idx in enumerate(adapted_ji_idx): 253 ki = ji_idx // nkpts 254 kj = ji_idx % nkpts 255 256 moij, ijslice = _conc_mos(mo_coeff_kpts[0][ki], mo_coeff_kpts[1][kj])[2:] 257 zij = [] 258 for LpqR, LpqI, sign in mydf.sr_loop(kpts[[ki,kj]], max_memory, False, mydf.blockdim): 259 zij.append(_ao2mo.r_e2(LpqR+LpqI*1j, moij, ijslice, tao, ao_loc)) 260 261 for kk in range(nkpts): 262 kl = kconserv[ki, kj, kk] 263 mokl, klslice = _conc_mos(mo_coeff_kpts[2][kk], mo_coeff_kpts[3][kl])[2:] 264 eri_mo = numpy.zeros((nmoi*nmoj,nmok*nmol), dtype=numpy.complex128) 265 for i, (LrsR, LrsI, sign) in \ 266 enumerate(mydf.sr_loop(kpts[[kk,kl]], max_memory, False, mydf.blockdim)): 267 zkl = _ao2mo.r_e2(LrsR+LrsI*1j, mokl, klslice, tao, ao_loc) 268 lib.dot(zij[i].T, zkl, sign*factor, eri_mo, 1) 269 270 if dtype == numpy.double: 271 eri_mo = eri_mo.real 272 out[ki,kj,kk] = eri_mo.reshape(eri_shape[3:]) 273 return out 274 275 276def _mo_as_complex(mo_coeffs): 277 mos = [] 278 for c in mo_coeffs: 279 if c.dtype == numpy.float64: 280 mos.append(c+0j) 281 else: 282 mos.append(c) 283 return mos 284 285def _dtrans(Lpq, Lij, ijmosym, moij, ijslice, 286 Lrs, Lkl, klmosym, mokl, klslice, sym): 287 Lij = _ao2mo.nr_e2(Lpq, moij, ijslice, aosym='s2', mosym=ijmosym, out=Lij) 288 if sym: 289 Lkl = Lij 290 else: 291 Lkl = _ao2mo.nr_e2(Lrs, mokl, klslice, aosym='s2', mosym=klmosym, out=Lkl) 292 return Lij, Lkl 293 294def _ztrans(Lpq, zij, moij, ijslice, Lrs, zkl, mokl, klslice, sym): 295 tao = [] 296 ao_loc = None 297 zij = _ao2mo.r_e2(Lpq, moij, ijslice, tao, ao_loc, out=zij) 298 if sym: 299 zkl = zij 300 else: 301 zkl = _ao2mo.r_e2(Lrs, mokl, klslice, tao, ao_loc, out=zkl) 302 return zij, zkl 303 304 305class PBC2DIntegralsWarning(RuntimeWarning): 306 pass 307def warn_pbc2d_eri(mydf): 308 cell = mydf.cell 309 if cell.dimension == 2 and cell.low_dim_ft_type == 'inf_vacuum': 310 with warnings.catch_warnings(): 311 warnings.simplefilter('once', PBC2DIntegralsWarning) 312 warnings.warn('\nERIs of PBC-2D systems with infinity vacuum are ' 313 'singular. cell.low_dim_ft_type = None should be ' 314 'set.\n') 315 316 317if __name__ == '__main__': 318 from pyscf.pbc import gto as pgto 319 from pyscf.pbc.df import DF 320 321 L = 5. 322 n = 11 323 cell = pgto.Cell() 324 cell.a = numpy.diag([L,L,L]) 325 cell.mesh = numpy.array([n,n,n]) 326 327 cell.atom = '''He 3. 2. 3. 328 He 1. 1. 1.''' 329 #cell.basis = {'He': [[0, (1.0, 1.0)]]} 330 #cell.basis = '631g' 331 #cell.basis = {'He': [[0, (2.4, 1)], [1, (1.1, 1)]]} 332 cell.basis = 'ccpvdz' 333 cell.verbose = 0 334 cell.build(0,0) 335 336 nao = cell.nao_nr() 337 numpy.random.seed(1) 338 kpts = numpy.random.random((4,3)) 339 kpts[3] = -numpy.einsum('ij->j', kpts[:3]) 340 with_df = DF(cell, kpts) 341 with_df.auxbasis = 'weigend' 342 with_df.mesh = [n] * 3 343 mo =(numpy.random.random((nao,nao)) + 344 numpy.random.random((nao,nao))*1j) 345 eri = with_df.get_eri(kpts).reshape((nao,)*4) 346 eri0 = numpy.einsum('pjkl,pi->ijkl', eri , mo.conj()) 347 eri0 = numpy.einsum('ipkl,pj->ijkl', eri0, mo ) 348 eri0 = numpy.einsum('ijpl,pk->ijkl', eri0, mo.conj()) 349 eri0 = numpy.einsum('ijkp,pl->ijkl', eri0, mo ) 350 eri1 = with_df.ao2mo(mo, kpts) 351 print(abs(eri1-eri0).sum()) 352