1#!/usr/bin/env python 2# Copyright 2014-2019 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 19'''Density expansion on plane waves''' 20 21 22import copy 23import numpy 24from pyscf import lib 25from pyscf import gto 26from pyscf.lib import logger 27from pyscf.pbc import tools 28from pyscf.pbc.gto import pseudo, estimate_ke_cutoff, error_for_ke_cutoff 29from pyscf.pbc import gto as pbcgto 30from pyscf.pbc.df import ft_ao 31from pyscf.pbc.df import incore 32from pyscf.pbc.lib.kpts_helper import is_zero, gamma_point 33from pyscf.pbc.df import aft_jk 34from pyscf.pbc.df import aft_ao2mo 35from pyscf import __config__ 36 37 38CUTOFF = getattr(__config__, 'pbc_df_aft_estimate_eta_cutoff', 1e-12) 39ETA_MIN = getattr(__config__, 'pbc_df_aft_estimate_eta_min', 0.2) 40PRECISION = getattr(__config__, 'pbc_df_aft_estimate_eta_precision', 1e-8) 41KE_SCALING = getattr(__config__, 'pbc_df_aft_ke_cutoff_scaling', 0.75) 42 43def estimate_eta(cell, cutoff=CUTOFF): 44 '''The exponent of the smooth gaussian model density, requiring that at 45 boundary, density ~ 4pi rmax^2 exp(-eta/2*rmax^2) ~ 1e-12 46 ''' 47 lmax = min(numpy.max(cell._bas[:,gto.ANG_OF]), 4) 48 # If lmax=3 (r^5 for radial part), this expression guarantees at least up 49 # to f shell the convergence at boundary 50 eta = max(numpy.log(4*numpy.pi*cell.rcut**(lmax+2)/cutoff)/cell.rcut**2*2, 51 ETA_MIN) 52 return eta 53 54def estimate_eta_for_ke_cutoff(cell, ke_cutoff, precision=PRECISION): 55 '''Given ke_cutoff, the upper bound of eta to produce the required 56 precision in AFTDF Coulomb integrals. 57 ''' 58 # search eta for interaction between GTO(eta) and point charge at the same 59 # location so that 60 # \sum_{k^2/2 > ke_cutoff} weight*4*pi/k^2 GTO(eta, k) < precision 61 # GTO(eta, k) = Fourier transform of Gaussian e^{-eta r^2} 62 63 lmax = numpy.max(cell._bas[:,gto.ANG_OF]) 64 kmax = (ke_cutoff*2)**.5 65 # The interaction between two s-type density distributions should be 66 # enough for the error estimation. Put lmax here to increate Ecut for 67 # slightly better accuracy 68 log_rest = numpy.log(precision / (32*numpy.pi**2 * kmax**(lmax-1))) 69 log_eta = -1 70 eta = kmax**2/4 / (-log_eta - log_rest) 71 return eta 72 73def estimate_ke_cutoff_for_eta(cell, eta, precision=PRECISION): 74 '''Given eta, the lower bound of ke_cutoff to produce the required 75 precision in AFTDF Coulomb integrals. 76 ''' 77 # estimate ke_cutoff for interaction between GTO(eta) and point charge at 78 # the same location so that 79 # \sum_{k^2/2 > ke_cutoff} weight*4*pi/k^2 GTO(eta, k) < precision 80 # \sum_{k^2/2 > ke_cutoff} weight*4*pi/k^2 GTO(eta, k) 81 # ~ \int_kmax^infty 4*pi/k^2 GTO(eta,k) dk^3 82 # = (4*pi)^2 *2*eta/kmax^{n-1} e^{-kmax^2/4eta} + ... < precision 83 84 # The magic number 0.2 comes from AFTDF.__init__ and GDF.__init__ 85 eta = max(eta, 0.2) 86 log_k0 = 3 + numpy.log(eta) / 2 87 log_rest = numpy.log(precision / (32*numpy.pi**2*eta)) 88 # The interaction between two s-type density distributions should be 89 # enough for the error estimation. Put lmax here to increate Ecut for 90 # slightly better accuracy 91 lmax = numpy.max(cell._bas[:,gto.ANG_OF]) 92 Ecut = 2*eta * (log_k0*(lmax-1) - log_rest) 93 Ecut = max(Ecut, .5) 94 return Ecut 95 96# \sum_{k^2/2 > ke_cutoff} weight*4*pi/k^2 exp(-k^2/(4 omega^2)) rho(k) < precision 97# ~ 16 pi^2 int_cutoff^infty exp(-k^2/(4*omega^2)) dk 98# = 16 pi^{5/2} omega erfc(sqrt(ke_cutoff/(2*omega^2))) 99# ~ 16 pi^2 exp(-ke_cutoff/(2*omega^2))) 100def estimate_ke_cutoff_for_omega(cell, omega, precision=None): 101 '''Energy cutoff to converge attenuated Coulomb in moment space 102 ''' 103 if precision is None: 104 precision = cell.precision 105 ke_cutoff = -2*omega**2 * numpy.log(precision / (32*omega**2*numpy.pi**2)) 106 return ke_cutoff 107 108def estimate_omega_for_ke_cutoff(cell, ke_cutoff, precision=None): 109 '''The minimal omega in attenuated Coulombl given energy cutoff 110 ''' 111 if precision is None: 112 precision = cell.precision 113 omega = (-.5 * ke_cutoff / numpy.log(precision / (32*numpy.pi**2)))**.5 114 return omega 115 116def get_nuc(mydf, kpts=None): 117 # Pseudopotential is ignored when computing just the nuclear attraction 118 t0 = (logger.process_clock(), logger.perf_counter()) 119 with lib.temporary_env(mydf.cell, _pseudo={}): 120 nuc = get_pp_loc_part1(mydf, kpts) 121 logger.timer(mydf, 'get_nuc', *t0) 122 return nuc 123 124def get_pp_loc_part1(mydf, kpts=None): 125 cell = mydf.cell 126 if kpts is None: 127 kpts_lst = numpy.zeros((1,3)) 128 else: 129 kpts_lst = numpy.reshape(kpts, (-1,3)) 130 131 log = logger.Logger(mydf.stdout, mydf.verbose) 132 t0 = t1 = (logger.process_clock(), logger.perf_counter()) 133 134 mesh = numpy.asarray(mydf.mesh) 135 nkpts = len(kpts_lst) 136 nao = cell.nao_nr() 137 nao_pair = nao * (nao+1) // 2 138 charges = cell.atom_charges() 139 140 kpt_allow = numpy.zeros(3) 141 if mydf.eta == 0: 142 if cell.dimension > 0: 143 ke_guess = estimate_ke_cutoff(cell, cell.precision) 144 mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess) 145 if numpy.any(mesh[:cell.dimension] < mesh_guess[:cell.dimension]*.8): 146 logger.warn(mydf, 'mesh %s is not enough for AFTDF.get_nuc function ' 147 'to get integral accuracy %g.\nRecommended mesh is %s.', 148 mesh, cell.precision, mesh_guess) 149 Gv, Gvbase, kws = cell.get_Gv_weights(mesh) 150 151 vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv) 152 vpplocG = -numpy.einsum('ij,ij->j', cell.get_SI(Gv), vpplocG) 153 154 vpplocG *= kws 155 vG = vpplocG 156 vj = numpy.zeros((nkpts,nao_pair), dtype=numpy.complex128) 157 158 else: 159 if cell.dimension > 0: 160 ke_guess = estimate_ke_cutoff_for_eta(cell, mydf.eta, cell.precision) 161 mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess) 162 if numpy.any(mesh < mesh_guess*.8): 163 logger.warn(mydf, 'mesh %s is not enough for AFTDF.get_nuc function ' 164 'to get integral accuracy %g.\nRecommended mesh is %s.', 165 mesh, cell.precision, mesh_guess) 166 mesh_min = numpy.min((mesh_guess, mesh), axis=0) 167 if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum': 168 mesh[:cell.dimension] = mesh_min[:cell.dimension] 169 else: 170 mesh = mesh_min 171 Gv, Gvbase, kws = cell.get_Gv_weights(mesh) 172 173 nuccell = _compensate_nuccell(mydf) 174 # PP-loc part1 is handled by fakenuc in _int_nuc_vloc 175 vj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst)) 176 t0 = t1 = log.timer_debug1('vnuc pass1: analytic int', *t0) 177 178 coulG = tools.get_coulG(cell, kpt_allow, mesh=mesh, Gv=Gv) * kws 179 aoaux = ft_ao.ft_ao(nuccell, Gv) 180 vG = numpy.einsum('i,xi->x', -charges, aoaux) * coulG 181 182 max_memory = max(2000, mydf.max_memory-lib.current_memory()[0]) 183 for aoaoks, p0, p1 in mydf.ft_loop(mesh, kpt_allow, kpts_lst, 184 max_memory=max_memory, aosym='s2'): 185 for k, aoao in enumerate(aoaoks): 186 # rho_ij(G) nuc(-G) / G^2 187 # = [Re(rho_ij(G)) + Im(rho_ij(G))*1j] [Re(nuc(G)) - Im(nuc(G))*1j] / G^2 188 if gamma_point(kpts_lst[k]): 189 vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real) 190 vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag) 191 else: 192 vj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao) 193 t1 = log.timer_debug1('contracting Vnuc [%s:%s]'%(p0, p1), *t1) 194 log.timer_debug1('contracting Vnuc', *t0) 195 196 vj_kpts = [] 197 for k, kpt in enumerate(kpts_lst): 198 if gamma_point(kpt): 199 vj_kpts.append(lib.unpack_tril(vj[k].real.copy())) 200 else: 201 vj_kpts.append(lib.unpack_tril(vj[k])) 202 203 if kpts is None or numpy.shape(kpts) == (3,): 204 vj_kpts = vj_kpts[0] 205 return numpy.asarray(vj_kpts) 206 207def _int_nuc_vloc(mydf, nuccell, kpts, intor='int3c2e', aosym='s2', comp=1): 208 '''Vnuc - Vloc''' 209 cell = mydf.cell 210 nkpts = len(kpts) 211 212 # Use the 3c2e code with steep s gaussians to mimic nuclear density 213 fakenuc = _fake_nuc(cell) 214 fakenuc._atm, fakenuc._bas, fakenuc._env = \ 215 gto.conc_env(nuccell._atm, nuccell._bas, nuccell._env, 216 fakenuc._atm, fakenuc._bas, fakenuc._env) 217 218 kptij_lst = numpy.hstack((kpts,kpts)).reshape(-1,2,3) 219 buf = incore.aux_e2(cell, fakenuc, intor, aosym=aosym, comp=comp, 220 kptij_lst=kptij_lst) 221 222 charge = cell.atom_charges() 223 charge = numpy.append(charge, -charge) # (charge-of-nuccell, charge-of-fakenuc) 224 nao = cell.nao_nr() 225 nchg = len(charge) 226 if aosym == 's1': 227 nao_pair = nao**2 228 else: 229 nao_pair = nao*(nao+1)//2 230 if comp == 1: 231 buf = buf.reshape(nkpts,nao_pair,nchg) 232 mat = numpy.einsum('kxz,z->kx', buf, charge) 233 else: 234 buf = buf.reshape(nkpts,comp,nao_pair,nchg) 235 mat = numpy.einsum('kcxz,z->kcx', buf, charge) 236 237 # vbar is the interaction between the background charge 238 # and the compensating function. 0D, 1D, 2D do not have vbar. 239 if cell.dimension == 3 and intor in ('int3c2e', 'int3c2e_sph', 240 'int3c2e_cart'): 241 assert(comp == 1) 242 charge = -cell.atom_charges() 243 244 nucbar = sum([z/nuccell.bas_exp(i)[0] for i,z in enumerate(charge)]) 245 nucbar *= numpy.pi/cell.vol 246 247 ovlp = cell.pbc_intor('int1e_ovlp', 1, lib.HERMITIAN, kpts) 248 for k in range(nkpts): 249 if aosym == 's1': 250 mat[k] -= nucbar * ovlp[k].reshape(nao_pair) 251 else: 252 mat[k] -= nucbar * lib.pack_tril(ovlp[k]) 253 return mat 254 255def get_pp(mydf, kpts=None): 256 '''Get the periodic pseudotential nuc-el AO matrix, with G=0 removed. 257 ''' 258 t0 = (logger.process_clock(), logger.perf_counter()) 259 cell = mydf.cell 260 if kpts is None: 261 kpts_lst = numpy.zeros((1,3)) 262 else: 263 kpts_lst = numpy.reshape(kpts, (-1,3)) 264 nkpts = len(kpts_lst) 265 266 vloc1 = get_pp_loc_part1(mydf, kpts_lst) 267 t1 = logger.timer_debug1(mydf, 'get_pp_loc_part1', *t0) 268 vloc2 = pseudo.pp_int.get_pp_loc_part2(cell, kpts_lst) 269 t1 = logger.timer_debug1(mydf, 'get_pp_loc_part2', *t1) 270 vpp = pseudo.pp_int.get_pp_nl(cell, kpts_lst) 271 for k in range(nkpts): 272 vpp[k] += vloc1[k] + vloc2[k] 273 t1 = logger.timer_debug1(mydf, 'get_pp_nl', *t1) 274 275 if kpts is None or numpy.shape(kpts) == (3,): 276 vpp = vpp[0] 277 logger.timer(mydf, 'get_pp', *t0) 278 return vpp 279 280def weighted_coulG(mydf, kpt=numpy.zeros(3), exx=False, mesh=None): 281 cell = mydf.cell 282 if mesh is None: 283 mesh = mydf.mesh 284 Gv, Gvbase, kws = cell.get_Gv_weights(mesh) 285 coulG = tools.get_coulG(cell, kpt, exx, mydf, mesh, Gv) 286 coulG *= kws 287 return coulG 288 289 290class AFTDF(lib.StreamObject): 291 '''Density expansion on plane waves 292 ''' 293 def __init__(self, cell, kpts=numpy.zeros((1,3))): 294 self.cell = cell 295 self.stdout = cell.stdout 296 self.verbose = cell.verbose 297 self.max_memory = cell.max_memory 298 self.mesh = cell.mesh 299# For nuclear attraction integrals using Ewald-like technique. 300# Set to 0 to switch off Ewald tech and use the regular reciprocal space 301# method (solving Poisson equation of nuclear charges in reciprocal space). 302 if cell.dimension == 0: 303 self.eta = 0.2 304 else: 305 ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(), self.mesh) 306 ke_cutoff = ke_cutoff[:cell.dimension].min() 307 self.eta = max(estimate_eta_for_ke_cutoff(cell, ke_cutoff, cell.precision), 308 estimate_eta(cell, cell.precision)) 309 self.kpts = kpts 310 311 # to mimic molecular DF object 312 self.blockdim = getattr(__config__, 'pbc_df_df_DF_blockdim', 240) 313 314 # The following attributes are not input options. 315 self._rsh_df = {} # Range separated Coulomb DF objects 316 self._keys = set(self.__dict__.keys()) 317 318 def dump_flags(self, verbose=None): 319 logger.info(self, '\n') 320 logger.info(self, '******** %s ********', self.__class__) 321 logger.info(self, 'mesh = %s (%d PWs)', self.mesh, numpy.prod(self.mesh)) 322 logger.info(self, 'eta = %s', self.eta) 323 logger.info(self, 'len(kpts) = %d', len(self.kpts)) 324 logger.debug1(self, ' kpts = %s', self.kpts) 325 return self 326 327 def reset(self, cell=None): 328 if cell is not None: 329 self.cell = cell 330 self._rsh_df = {} 331 return self 332 333 def check_sanity(self): 334 lib.StreamObject.check_sanity(self) 335 cell = self.cell 336 if not cell.has_ecp(): 337 logger.warn(self, 'AFTDF integrals are found in all-electron ' 338 'calculation. It often causes huge error.\n' 339 'Recommended methods are DF or MDF. In SCF calculation, ' 340 'they can be initialized as\n' 341 ' mf = mf.density_fit()\nor\n' 342 ' mf = mf.mix_density_fit()') 343 344 if cell.dimension > 0: 345 if cell.ke_cutoff is None: 346 ke_cutoff = tools.mesh_to_cutoff(cell.lattice_vectors(), self.mesh) 347 ke_cutoff = ke_cutoff[:cell.dimension].min() 348 else: 349 ke_cutoff = numpy.min(cell.ke_cutoff) 350 ke_guess = estimate_ke_cutoff(cell, cell.precision) 351 mesh_guess = tools.cutoff_to_mesh(cell.lattice_vectors(), ke_guess) 352 if ke_cutoff < ke_guess * KE_SCALING: 353 logger.warn(self, 'ke_cutoff/mesh (%g / %s) is not enough for AFTDF ' 354 'to get integral accuracy %g.\nCoulomb integral error ' 355 'is ~ %.2g Eh.\nRecommended ke_cutoff/mesh are %g / %s.', 356 ke_cutoff, self.mesh, cell.precision, 357 error_for_ke_cutoff(cell, ke_cutoff), ke_guess, mesh_guess) 358 else: 359 mesh_guess = numpy.copy(self.mesh) 360 361 if cell.dimension < 3: 362 err = numpy.exp(-0.436392335*min(self.mesh[cell.dimension:]) - 2.99944305) 363 err *= cell.nelectron 364 meshz = pbcgto.cell._mesh_inf_vaccum(cell) 365 mesh_guess[cell.dimension:] = int(meshz) 366 if err > cell.precision*10: 367 logger.warn(self, 'mesh %s of AFTDF may not be enough to get ' 368 'integral accuracy %g for %dD PBC system.\n' 369 'Coulomb integral error is ~ %.2g Eh.\n' 370 'Recommended mesh is %s.', 371 self.mesh, cell.precision, cell.dimension, err, mesh_guess) 372 if any(x/meshz > 1.1 for x in cell.mesh[cell.dimension:]): 373 meshz = pbcgto.cell._mesh_inf_vaccum(cell) 374 logger.warn(self, 'setting mesh %s of AFTDF too high in non-periodic direction ' 375 '(=%s) can result in an unnecessarily slow calculation.\n' 376 'For coulomb integral error of ~ %.2g Eh in %dD PBC, \n' 377 'a recommended mesh for non-periodic direction is %s.', 378 self.mesh, self.mesh[cell.dimension:], cell.precision, 379 cell.dimension, mesh_guess[cell.dimension:]) 380 return self 381 382# TODO: Put Gv vector in the arguments 383 def pw_loop(self, mesh=None, kpti_kptj=None, q=None, shls_slice=None, 384 max_memory=2000, aosym='s1', blksize=None, 385 intor='GTO_ft_ovlp', comp=1, bvk_kmesh=None): 386 ''' 387 Fourier transform iterator for AO pair 388 ''' 389 cell = self.cell 390 if mesh is None: 391 mesh = self.mesh 392 if kpti_kptj is None: 393 kpti = kptj = numpy.zeros(3) 394 else: 395 kpti, kptj = kpti_kptj 396 if q is None: 397 q = kptj - kpti 398 399 ao_loc = cell.ao_loc_nr() 400 Gv, Gvbase, kws = cell.get_Gv_weights(mesh) 401 b = cell.reciprocal_vectors() 402 gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) 403 ngrids = gxyz.shape[0] 404 405 if shls_slice is None: 406 shls_slice = (0, cell.nbas, 0, cell.nbas) 407 if aosym == 's2': 408 assert(shls_slice[2] == 0) 409 i0 = ao_loc[shls_slice[0]] 410 i1 = ao_loc[shls_slice[1]] 411 nij = i1*(i1+1)//2 - i0*(i0+1)//2 412 else: 413 ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]] 414 nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]] 415 nij = ni*nj 416 417 if blksize is None: 418 blksize = min(max(64, int(max_memory*1e6*.75/(nij*16*comp))), 16384) 419 sublk = int(blksize//4) 420 else: 421 sublk = blksize 422 buf = numpy.empty(nij*blksize*comp, dtype=numpy.complex128) 423 pqkRbuf = numpy.empty(nij*sublk*comp) 424 pqkIbuf = numpy.empty(nij*sublk*comp) 425 426 for p0, p1 in self.prange(0, ngrids, blksize): 427 #aoao = ft_ao.ft_aopair(cell, Gv[p0:p1], shls_slice, aosym, 428 # b, Gvbase, gxyz[p0:p1], mesh, (kpti, kptj), q) 429 aoao = ft_ao.ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, 430 b, gxyz[p0:p1], Gvbase, q, 431 kptj.reshape(1,3), intor, comp, 432 bvk_kmesh=bvk_kmesh, out=buf)[0] 433 aoao = aoao.reshape(p1-p0,nij) 434 435 for i0, i1 in lib.prange(0, p1-p0, sublk): 436 nG = i1 - i0 437 if comp == 1: 438 pqkR = numpy.ndarray((nij,nG), buffer=pqkRbuf) 439 pqkI = numpy.ndarray((nij,nG), buffer=pqkIbuf) 440 pqkR[:] = aoao[i0:i1].real.T 441 pqkI[:] = aoao[i0:i1].imag.T 442 else: 443 pqkR = numpy.ndarray((comp,nij,nG), buffer=pqkRbuf) 444 pqkI = numpy.ndarray((comp,nij,nG), buffer=pqkIbuf) 445 pqkR[:] = aoao[i0:i1].real.transpose(0,2,1) 446 pqkI[:] = aoao[i0:i1].imag.transpose(0,2,1) 447 yield (pqkR, pqkI, p0+i0, p0+i1) 448 449 def ft_loop(self, mesh=None, q=numpy.zeros(3), kpts=None, shls_slice=None, 450 max_memory=4000, aosym='s1', intor='GTO_ft_ovlp', comp=1, 451 bvk_kmesh=None): 452 ''' 453 Fourier transform iterator for all kpti which satisfy 454 2pi*N = (kpts - kpti - q)*a, N = -1, 0, 1 455 ''' 456 cell = self.cell 457 if mesh is None: 458 mesh = self.mesh 459 if kpts is None: 460 assert(is_zero(q)) 461 kpts = self.kpts 462 kpts = numpy.asarray(kpts) 463 nkpts = len(kpts) 464 465 ao_loc = cell.ao_loc_nr() 466 b = cell.reciprocal_vectors() 467 Gv, Gvbase, kws = cell.get_Gv_weights(mesh) 468 gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase]) 469 ngrids = gxyz.shape[0] 470 471 if shls_slice is None: 472 shls_slice = (0, cell.nbas, 0, cell.nbas) 473 if aosym == 's2': 474 assert(shls_slice[2] == 0) 475 i0 = ao_loc[shls_slice[0]] 476 i1 = ao_loc[shls_slice[1]] 477 nij = i1*(i1+1)//2 - i0*(i0+1)//2 478 else: 479 ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]] 480 nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]] 481 nij = ni*nj 482 483 blksize = max(16, int(max_memory*.9e6/(nij*nkpts*16*comp))) 484 blksize = min(blksize, ngrids, 16384) 485 buf = numpy.empty(nkpts*nij*blksize*comp, dtype=numpy.complex128) 486 487 for p0, p1 in self.prange(0, ngrids, blksize): 488 dat = ft_ao.ft_aopair_kpts(cell, Gv[p0:p1], shls_slice, aosym, 489 b, gxyz[p0:p1], Gvbase, q, kpts, 490 intor, comp, bvk_kmesh=bvk_kmesh, out=buf) 491 yield dat, p0, p1 492 493 weighted_coulG = weighted_coulG 494 _int_nuc_vloc = _int_nuc_vloc 495 get_nuc = get_nuc 496 get_pp = get_pp 497 498 # Note: Special exxdiv by default should not be used for an arbitrary 499 # input density matrix. When the df object was used with the molecular 500 # post-HF code, get_jk was often called with an incomplete DM (e.g. the 501 # core DM in CASCI). An SCF level exxdiv treatment is inadequate for 502 # post-HF methods. 503 def get_jk(self, dm, hermi=1, kpts=None, kpts_band=None, 504 with_j=True, with_k=True, omega=None, exxdiv=None): 505 if omega is not None: # J/K for RSH functionals 506 return _sub_df_jk_(self, dm, hermi, kpts, kpts_band, 507 with_j, with_k, omega, exxdiv) 508 509 if kpts is None: 510 if numpy.all(self.kpts == 0): 511 # Gamma-point calculation by default 512 kpts = numpy.zeros(3) 513 else: 514 kpts = self.kpts 515 kpts = numpy.asarray(kpts) 516 517 if kpts.shape == (3,): 518 return aft_jk.get_jk(self, dm, hermi, kpts, kpts_band, with_j, 519 with_k, exxdiv) 520 521 vj = vk = None 522 if with_k: 523 vk = aft_jk.get_k_kpts(self, dm, hermi, kpts, kpts_band, exxdiv) 524 if with_j: 525 vj = aft_jk.get_j_kpts(self, dm, hermi, kpts, kpts_band) 526 return vj, vk 527 528 get_eri = get_ao_eri = aft_ao2mo.get_eri 529 ao2mo = get_mo_eri = aft_ao2mo.general 530 ao2mo_7d = aft_ao2mo.ao2mo_7d 531 get_ao_pairs_G = get_ao_pairs = aft_ao2mo.get_ao_pairs_G 532 get_mo_pairs_G = get_mo_pairs = aft_ao2mo.get_mo_pairs_G 533 534 def update_mf(self, mf): 535 mf = copy.copy(mf) 536 mf.with_df = self 537 return mf 538 539 def prange(self, start, stop, step): 540 '''This is a hook for MPI parallelization. DO NOT use it out of the 541 scope of AFTDF/GDF/MDF. 542 ''' 543 return lib.prange(start, stop, step) 544 545################################################################################ 546# With this function to mimic the molecular DF.loop function, the pbc gamma 547# point DF object can be used in the molecular code 548 def loop(self, blksize=None): 549 cell = self.cell 550 if cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum': 551 raise RuntimeError('ERIs of PBC-2D systems are not positive ' 552 'definite. Current API only supports postive ' 553 'definite ERIs.') 554 555 if blksize is None: 556 blksize = self.blockdim 557 # coulG of 1D and 2D has negative elements. 558 coulG = self.weighted_coulG() 559 Lpq = None 560 for pqkR, pqkI, p0, p1 in self.pw_loop(aosym='s2', blksize=blksize): 561 vG = numpy.sqrt(coulG[p0:p1]) 562 pqkR *= vG 563 pqkI *= vG 564 Lpq = lib.transpose(pqkR, out=Lpq) 565 yield Lpq 566 Lpq = lib.transpose(pqkI, out=Lpq) 567 yield Lpq 568 569 def get_naoaux(self): 570 mesh = numpy.asarray(self.mesh) 571 ngrids = numpy.prod(mesh) 572 return ngrids * 2 573 574 575# Since the real-space lattice-sum for nuclear attraction is not implemented, 576# use the 3c2e code with steep gaussians to mimic nuclear density 577def _fake_nuc(cell): 578 fakenuc = copy.copy(cell) 579 fakenuc._atm = cell._atm.copy() 580 fakenuc._atm[:,gto.PTR_COORD] = numpy.arange(gto.PTR_ENV_START, 581 gto.PTR_ENV_START+cell.natm*3,3) 582 _bas = [] 583 _env = [0]*gto.PTR_ENV_START + [cell.atom_coords().ravel()] 584 ptr = gto.PTR_ENV_START + cell.natm * 3 585 half_sph_norm = .5/numpy.sqrt(numpy.pi) 586 for ia in range(cell.natm): 587 symb = cell.atom_symbol(ia) 588 if symb in cell._pseudo: 589 pp = cell._pseudo[symb] 590 rloc, nexp, cexp = pp[1:3+1] 591 eta = .5 / rloc**2 592 else: 593 eta = 1e16 594 norm = half_sph_norm/gto.gaussian_int(2, eta) 595 _env.extend([eta, norm]) 596 _bas.append([ia, 0, 1, 1, 0, ptr, ptr+1, 0]) 597 ptr += 2 598 fakenuc._bas = numpy.asarray(_bas, dtype=numpy.int32) 599 fakenuc._env = numpy.asarray(numpy.hstack(_env), dtype=numpy.double) 600 fakenuc.rcut = cell.rcut 601 return fakenuc 602 603def _compensate_nuccell(mydf): 604 '''A cell of the compensated Gaussian charges for nucleus''' 605 cell = mydf.cell 606 nuccell = copy.copy(cell) 607 half_sph_norm = .5/numpy.sqrt(numpy.pi) 608 norm = half_sph_norm/gto.gaussian_int(2, mydf.eta) 609 chg_env = [mydf.eta, norm] 610 ptr_eta = cell._env.size 611 ptr_norm = ptr_eta + 1 612 chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0] for ia in range(cell.natm)] 613 nuccell._atm = cell._atm 614 nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32) 615 nuccell._env = numpy.hstack((cell._env, chg_env)) 616 return nuccell 617 618def _sub_df_jk_(dfobj, dm, hermi=1, kpts=None, kpts_band=None, 619 with_j=True, with_k=True, omega=None, exxdiv=None): 620 key = '%.6f' % omega 621 if key in dfobj._rsh_df: 622 rsh_df = dfobj._rsh_df[key] 623 else: 624 rsh_df = dfobj._rsh_df[key] = copy.copy(dfobj).reset() 625 logger.info(dfobj, 'Create RSH-%s object %s for omega=%s', 626 dfobj.__class__.__name__, rsh_df, omega) 627 with rsh_df.cell.with_range_coulomb(omega): 628 return rsh_df.get_jk(dm, hermi, kpts, kpts_band, with_j, with_k, 629 omega=None, exxdiv=exxdiv) 630 631del(CUTOFF, PRECISION) 632 633 634if __name__ == '__main__': 635 cell = pbcgto.Cell() 636 cell.verbose = 0 637 cell.atom = 'C 0 0 0; C 1 1 1' 638 cell.a = numpy.diag([4, 4, 4]) 639 cell.basis = 'gth-szv' 640 cell.pseudo = 'gth-pade' 641 cell.mesh = [20]*3 642 cell.build() 643 k = numpy.ones(3)*.25 644 v1 = AFTDF(cell).get_pp(k) 645 print(abs(v1).sum(), 21.7504294462) 646