Source code for pyscf.df.df_jk

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import ctypes
import numpy
import scipy.linalg
from pyscf import lib
from pyscf import scf
from pyscf.lib import logger
from pyscf.ao2mo import _ao2mo

DEBUG = False

libri = lib.load_library('libri')

[docs] def density_fit(mf, auxbasis=None, with_df=None, only_dfj=False): '''For the given SCF object, update the J, K matrix constructor with corresponding density fitting integrals. Args: mf : an SCF object Kwargs: auxbasis : str or basis dict Same format to the input attribute mol.basis. If auxbasis is None, optimal auxiliary basis based on AO basis (if possible) or even-tempered Gaussian basis will be used. only_dfj : str Compute Coulomb integrals only and no approximation for HF exchange. Same to RIJONX in ORCA Returns: An SCF object with a modified J, K matrix constructor which uses density fitting integrals to compute J and K Examples: >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0) >>> mf = scf.density_fit(scf.RHF(mol)) >>> mf.scf() -100.005306000435510 >>> mol.symmetry = 1 >>> mol.build(0, 0) >>> mf = scf.density_fit(scf.UHF(mol)) >>> mf.scf() -100.005306000435510 ''' from pyscf import df from pyscf.scf import dhf assert (isinstance(mf, scf.hf.SCF)) if with_df is None: if isinstance(mf, dhf.UHF): with_df = df.DF4C(mf.mol) else: with_df = df.DF(mf.mol) with_df.max_memory = mf.max_memory with_df.stdout = mf.stdout with_df.verbose = mf.verbose with_df.auxbasis = auxbasis if isinstance(mf, _DFHF): if mf.with_df is None: mf.with_df = with_df elif getattr(mf.with_df, 'auxbasis', None) != auxbasis: #logger.warn(mf, 'DF might have been initialized twice.') mf = mf.copy() mf.with_df = with_df mf.only_dfj = only_dfj return mf dfmf = _DFHF(mf, with_df, only_dfj) return lib.set_class(dfmf, (_DFHF, mf.__class__))
# 1. A tag to label the derived SCF class # 2. A hook to register DF specific methods, such as nuc_grad_method. class _DFHF: ''' Density fitting SCF class Attributes for density-fitting SCF: auxbasis : str or basis dict Same format to the input attribute mol.basis. The default basis 'weigend+etb' means weigend-coulomb-fit basis for light elements and even-tempered basis for heavy elements. with_df : DF object Set mf.with_df = None to switch off density fitting mode. See also the documents of class for other SCF attributes. ''' __name_mixin__ = 'DF' _keys = {'with_df', 'only_dfj'} def __init__(self, mf, df=None, only_dfj=None): self.__dict__.update(mf.__dict__) self._eri = None self.with_df = df self.only_dfj = only_dfj # Unless DF is used only for J matrix, disable direct_scf for K build. # It is more efficient to construct K matrix with MO coefficients than # the incremental method in direct_scf. self.direct_scf = only_dfj def undo_df(self): '''Remove the DFHF Mixin''' obj = lib.view(self, lib.drop_class(self.__class__, _DFHF)) del obj.with_df, obj.only_dfj return obj def reset(self, mol=None): self.with_df.reset(mol) return super().reset(mol) def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True, omega=None): if dm is None: dm = self.make_rdm1() if not self.with_df: return super().get_jk(mol, dm, hermi, with_j, with_k, omega) with_dfk = with_k and not self.only_dfj if isinstance(self, scf.ghf.GHF): def jkbuild(mol, dm, hermi, with_j, with_k, omega=None): vj, vk = self.with_df.get_jk(dm.real, hermi, with_j, with_k, self.direct_scf_tol, omega) if dm.dtype == numpy.complex128: vjI, vkI = self.with_df.get_jk(dm.imag, hermi, with_j, with_k, self.direct_scf_tol, omega) if with_j: vj = vj + vjI * 1j if with_k: vk = vk + vkI * 1j return vj, vk vj, vk = scf.ghf.get_jk(mol, dm, hermi, with_j, with_dfk, jkbuild, omega) else: vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_dfk, self.direct_scf_tol, omega) if with_k and not with_dfk: vk = super().get_jk(mol, dm, hermi, False, True, omega)[1] return vj, vk # for pyscf 1.0, 1.1 compatibility @property def _cderi(self): naux = self.with_df.get_naoaux() return next(self.with_df.loop(blksize=naux)) @_cderi.setter def _cderi(self, x): self.with_df._cderi = x @property def auxbasis(self): return getattr(self.with_df, 'auxbasis', None) def nuc_grad_method(self): from pyscf.df.grad import rhf, rohf, uhf, rks, roks, uks if isinstance(self, scf.uhf.UHF): if isinstance(self, scf.hf.KohnShamDFT): return uks.Gradients(self) else: return uhf.Gradients(self) elif isinstance(self, scf.rohf.ROHF): if isinstance(self, scf.hf.KohnShamDFT): return roks.Gradients(self) else: return rohf.Gradients(self) elif isinstance(self, scf.rhf.RHF): if isinstance(self, scf.hf.KohnShamDFT): return rks.Gradients(self) else: return rhf.Gradients(self) else: raise NotImplementedError Gradients = nuc_grad_method def Hessian(self): from pyscf.df.hessian import rhf, uhf, rks, uks if isinstance(self, (scf.uhf.UHF, scf.rohf.ROHF)): if isinstance(self, scf.hf.KohnShamDFT): return uks.Hessian(self) else: return uhf.Hessian(self) elif isinstance(self, scf.rohf.ROHF): raise NotImplementedError elif isinstance(self, scf.rhf.RHF): if isinstance(self, scf.hf.KohnShamDFT): return rks.Hessian(self) else: return rhf.Hessian(self) else: raise NotImplementedError def method_not_implemented(self, *args, **kwargs): raise NotImplementedError NMR = method_not_implemented NSR = method_not_implemented Polarizability = method_not_implemented RotationalGTensor = method_not_implemented MP2 = method_not_implemented CISD = method_not_implemented CCSD = method_not_implemented def CASCI(self, ncas, nelecas, auxbasis=None, ncore=None): from pyscf import mcscf return mcscf.DFCASCI(self, ncas, nelecas, auxbasis, ncore) def CASSCF(self, ncas, nelecas, auxbasis=None, ncore=None, frozen=None): from pyscf import mcscf return mcscf.DFCASSCF(self, ncas, nelecas, auxbasis, ncore, frozen) def to_gpu(self): obj = self.undo_df().to_gpu().density_fit() return lib.to_gpu(self, obj)
[docs] def get_jk(dfobj, dm, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-13): assert (with_j or with_k) if (not with_k and not dfobj.mol.incore_anyway and # 3-center integral tensor is not initialized dfobj._cderi is None): return get_j(dfobj, dm, hermi, direct_scf_tol), None t0 = t1 = (logger.process_clock(), logger.perf_counter()) log = logger.Logger(dfobj.stdout, dfobj.verbose) fmmm = _ao2mo.libao2mo.AO2MOmmm_bra_nr_s2 fdrv = _ao2mo.libao2mo.AO2MOnr_e2_drv ftrans = _ao2mo.libao2mo.AO2MOtranse2_nr_s2 null = lib.c_null_ptr() dms = numpy.asarray(dm) dm_shape = dms.shape nao = dm_shape[-1] dms = dms.reshape(-1,nao,nao) nset = dms.shape[0] vj = 0 vk = numpy.zeros_like(dms) if with_j: idx = numpy.arange(nao) dmtril = lib.pack_tril(dms + dms.conj().transpose(0,2,1)) dmtril[:,idx*(idx+1)//2+idx] *= .5 if not with_k: for eri1 in dfobj.loop(): # uses numpy.matmul vj += (dmtril @ eri1.T) @ eri1 elif getattr(dm, 'mo_coeff', None) is not None: #TODO: test whether dm.mo_coeff matching dm mo_coeff = numpy.asarray(dm.mo_coeff, order='F') mo_occ = numpy.asarray(dm.mo_occ) nmo = mo_occ.shape[-1] mo_coeff = mo_coeff.reshape(-1,nao,nmo) mo_occ = mo_occ.reshape(-1,nmo) if mo_occ.shape[0] * 2 == nset: # handle ROHF DM mo_coeff = numpy.vstack((mo_coeff, mo_coeff)) mo_occa = numpy.array(mo_occ> 0, dtype=numpy.double) mo_occb = numpy.array(mo_occ==2, dtype=numpy.double) assert (mo_occa.sum() + mo_occb.sum() == mo_occ.sum()) mo_occ = numpy.vstack((mo_occa, mo_occb)) orbo = [] for k in range(nset): c = numpy.einsum('pi,i->pi', mo_coeff[k][:,mo_occ[k]>0], numpy.sqrt(mo_occ[k][mo_occ[k]>0])) orbo.append(numpy.asarray(c, order='F')) max_memory = dfobj.max_memory - lib.current_memory()[0] blksize = max(4, int(min(dfobj.blockdim, max_memory*.3e6/8/nao**2))) buf = numpy.empty((blksize*nao,nao)) for eri1 in dfobj.loop(blksize): naux, nao_pair = eri1.shape assert (nao_pair == nao*(nao+1)//2) if with_j: # uses numpy.matmul vj += (dmtril @ eri1.T) @ eri1 for k in range(nset): nocc = orbo[k].shape[1] if nocc > 0: buf1 = buf[:naux*nocc] fdrv(ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p), eri1.ctypes.data_as(ctypes.c_void_p), orbo[k].ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux), ctypes.c_int(nao), (ctypes.c_int*4)(0, nocc, 0, nao), null, ctypes.c_int(0)) vk[k] += lib.dot(buf1.T, buf1) t1 = log.timer_debug1('jk', *t1) else: #:vk = numpy.einsum('pij,jk->pki', cderi, dm) #:vk = numpy.einsum('pki,pkj->ij', cderi, vk) rargs = (ctypes.c_int(nao), (ctypes.c_int*4)(0, nao, 0, nao), null, ctypes.c_int(0)) dms = [numpy.asarray(x, order='F') for x in dms] max_memory = dfobj.max_memory - lib.current_memory()[0] blksize = max(4, int(min(dfobj.blockdim, max_memory*.22e6/8/nao**2))) buf = numpy.empty((2,blksize,nao,nao)) for eri1 in dfobj.loop(blksize): naux, nao_pair = eri1.shape if with_j: # uses numpy.matmul vj += (dmtril @ eri1.T) @ eri1 for k in range(nset): buf1 = buf[0,:naux] fdrv(ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p), eri1.ctypes.data_as(ctypes.c_void_p), dms[k].ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux), *rargs) buf2 = lib.unpack_tril(eri1, out=buf[1]) vk[k] += lib.dot(buf1.reshape(-1,nao).T, buf2.reshape(-1,nao)) t1 = log.timer_debug1('jk', *t1) if with_j: vj = lib.unpack_tril(vj, 1).reshape(dm_shape) if with_k: vk = vk.reshape(dm_shape) logger.timer(dfobj, 'df vj and vk', *t0) return vj, vk
[docs] def get_j(dfobj, dm, hermi=1, direct_scf_tol=1e-13): from pyscf.scf import _vhf from pyscf.scf import jk from pyscf.df import addons t0 = t1 = (logger.process_clock(), logger.perf_counter()) mol = dfobj.mol if dfobj._vjopt is None: dfobj.auxmol = auxmol = addons.make_auxmol(mol, dfobj.auxbasis) opt = _vhf._VHFOpt(mol, 'int3c2e', 'CVHFnr3c2e_schwarz_cond', dmcondname='CVHFnr_dm_cond', direct_scf_tol=direct_scf_tol) # q_cond part 1: the regular int2e (ij|ij) for mol's basis opt.init_cvhf_direct(mol, 'int2e', 'CVHFnr_int2e_q_cond') # Update q_cond to include the 2e-integrals (auxmol|auxmol) j2c = auxmol.intor('int2c2e', hermi=1) j2c_diag = numpy.sqrt(abs(j2c.diagonal())) aux_loc = auxmol.ao_loc aux_q_cond = [j2c_diag[i0:i1].max() for i0, i1 in zip(aux_loc[:-1], aux_loc[1:])] q_cond = numpy.hstack((opt.q_cond.ravel(), aux_q_cond)) opt.q_cond = q_cond try: opt.j2c = j2c = scipy.linalg.cho_factor(j2c, lower=True) opt.j2c_type = 'cd' except scipy.linalg.LinAlgError: opt.j2c = j2c opt.j2c_type = 'regular' # jk.get_jk function supports 4-index integrals. Use bas_placeholder # (l=0, nctr=1, 1 function) to hold the last index. bas_placeholder = numpy.array([0, 0, 1, 1, 0, 0, 0, 0], dtype=numpy.int32) fakemol = mol + auxmol fakemol._bas = numpy.vstack((fakemol._bas, bas_placeholder)) opt.fakemol = fakemol dfobj._vjopt = opt t1 = logger.timer_debug1(dfobj, 'df-vj init_direct_scf', *t1) opt = dfobj._vjopt fakemol = opt.fakemol dm = numpy.asarray(dm, order='C') dm_shape = dm.shape nao = dm_shape[-1] dm = dm.reshape(-1,nao,nao) n_dm = dm.shape[0] # First compute the density in auxiliary basis # j3c = fauxe2(mol, auxmol) # jaux = numpy.einsum('ijk,ji->k', j3c, dm) # rho = numpy.linalg.solve(auxmol.intor('int2c2e'), jaux) nbas = mol.nbas nbas1 = mol.nbas + dfobj.auxmol.nbas shls_slice = (0, nbas, 0, nbas, nbas, nbas1, nbas1, nbas1+1) with lib.temporary_env(opt, prescreen='CVHFnr3c2e_vj_pass1_prescreen'): jaux = jk.get_jk(fakemol, dm, ['ijkl,ji->kl']*n_dm, 'int3c2e', aosym='s2ij', hermi=0, shls_slice=shls_slice, vhfopt=opt) # remove the index corresponding to bas_placeholder jaux = numpy.array(jaux)[:,:,0] t1 = logger.timer_debug1(dfobj, 'df-vj pass 1', *t1) if opt.j2c_type == 'cd': rho = scipy.linalg.cho_solve(opt.j2c, jaux.T) else: rho = scipy.linalg.solve(opt.j2c, jaux.T) # transform rho to shape (:,1,naux), to adapt to 3c2e integrals (ij|k) rho = rho.T[:,numpy.newaxis,:] t1 = logger.timer_debug1(dfobj, 'df-vj solve ', *t1) # Next compute the Coulomb matrix # j3c = fauxe2(mol, auxmol) # vj = numpy.einsum('ijk,k->ij', j3c, rho) # temporarily set "_dmcondname=None" to skip the call to set_dm method. with lib.temporary_env(opt, prescreen='CVHFnr3c2e_vj_pass2_prescreen', _dmcondname=None): # CVHFnr3c2e_vj_pass2_prescreen requires custom dm_cond aux_loc = dfobj.auxmol.ao_loc dm_cond = [abs(rho[:,:,i0:i1]).max() for i0, i1 in zip(aux_loc[:-1], aux_loc[1:])] opt.dm_cond = numpy.array(dm_cond) vj = jk.get_jk(fakemol, rho, ['ijkl,lk->ij']*n_dm, 'int3c2e', aosym='s2ij', hermi=1, shls_slice=shls_slice, vhfopt=opt) t1 = logger.timer_debug1(dfobj, 'df-vj pass 2', *t1) logger.timer(dfobj, 'df-vj', *t0) return numpy.asarray(vj).reshape(dm_shape)
[docs] def r_get_jk(dfobj, dms, hermi=1, with_j=True, with_k=True): '''Relativistic density fitting JK''' t0 = (logger.process_clock(), logger.perf_counter()) mol = dfobj.mol c1 = .5 / lib.param.LIGHT_SPEED tao = mol.tmap() ao_loc = mol.ao_loc_2c() n2c = ao_loc[-1] if hermi == 0 and DEBUG: # J matrix is symmetrized in this function which is only true for # density matrix with time reversal symmetry scf.dhf._ensure_time_reversal_symmetry(mol, dms) def fjk(dm): dm = numpy.asarray(dm, dtype=numpy.complex128) fmmm = libri.RIhalfmmm_r_s2_bra_noconj fdrv = _ao2mo.libao2mo.AO2MOr_e2_drv ftrans = libri.RItranse2_r_s2 vj = numpy.zeros_like(dm) vk = numpy.zeros_like(dm) fcopy = libri.RImmm_r_s2_transpose rargs = (ctypes.c_int(n2c), (ctypes.c_int*4)(0, n2c, 0, 0), tao.ctypes.data_as(ctypes.c_void_p), ao_loc.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas)) dmll = numpy.asarray(dm[:n2c,:n2c], order='C') dmls = numpy.asarray(dm[:n2c,n2c:], order='C') * c1 dmsl = numpy.asarray(dm[n2c:,:n2c], order='C') * c1 dmss = numpy.asarray(dm[n2c:,n2c:], order='C') * c1**2 for erill, eriss in dfobj.loop(): naux, nao_pair = erill.shape buf = numpy.empty((naux,n2c,n2c), dtype=numpy.complex128) buf1 = numpy.empty((naux,n2c,n2c), dtype=numpy.complex128) fdrv(ftrans, fmmm, buf.ctypes.data_as(ctypes.c_void_p), erill.ctypes.data_as(ctypes.c_void_p), dmll.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux), *rargs) # buf == (P|LL) rho = numpy.einsum('kii->k', buf) fdrv(ftrans, fcopy, buf1.ctypes.data_as(ctypes.c_void_p), erill.ctypes.data_as(ctypes.c_void_p), dmll.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux), *rargs) # buf1 == (P|LL) vk[:n2c,:n2c] += numpy.dot(buf1.reshape(-1,n2c).T, buf.reshape(-1,n2c)) fdrv(ftrans, fmmm, buf.ctypes.data_as(ctypes.c_void_p), eriss.ctypes.data_as(ctypes.c_void_p), dmls.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux), *rargs) # buf == (P|LS) vk[:n2c,n2c:] += numpy.dot(buf1.reshape(-1,n2c).T, buf.reshape(-1,n2c)) * c1 fdrv(ftrans, fmmm, buf.ctypes.data_as(ctypes.c_void_p), eriss.ctypes.data_as(ctypes.c_void_p), dmss.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux), *rargs) # buf == (P|SS) rho += numpy.einsum('kii->k', buf) vj[:n2c,:n2c] += lib.unpack_tril(numpy.dot(rho, erill), 1) vj[n2c:,n2c:] += lib.unpack_tril(numpy.dot(rho, eriss), 1) * c1**2 fdrv(ftrans, fcopy, buf1.ctypes.data_as(ctypes.c_void_p), eriss.ctypes.data_as(ctypes.c_void_p), dmss.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux), *rargs) # buf == (P|SS) vk[n2c:,n2c:] += numpy.dot(buf1.reshape(-1,n2c).T, buf.reshape(-1,n2c)) * c1**2 if hermi != 1: fdrv(ftrans, fmmm, buf.ctypes.data_as(ctypes.c_void_p), erill.ctypes.data_as(ctypes.c_void_p), dmsl.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(naux), *rargs) # buf == (P|SL) vk[n2c:,:n2c] += numpy.dot(buf1.reshape(-1,n2c).T, buf.reshape(-1,n2c)) * c1 if hermi == 1: vk[n2c:,:n2c] = vk[:n2c,n2c:].T.conj() return vj, vk if isinstance(dms, numpy.ndarray) and dms.ndim == 2: vj, vk = fjk(dms) else: vjk = [fjk(dm) for dm in dms] vj = numpy.array([x[0] for x in vjk]) vk = numpy.array([x[1] for x in vjk]) logger.timer(dfobj, 'vj and vk', *t0) return vj, vk
if __name__ == '__main__': import pyscf.gto import pyscf.scf mol = pyscf.gto.Mole() mol.build( verbose = 0, atom = [["O" , (0. , 0. , 0.)], [1 , (0. , -0.757 , 0.587)], [1 , (0. , 0.757 , 0.587)] ], basis = 'ccpvdz', ) method = density_fit(pyscf.scf.RHF(mol), 'weigend') method.max_memory = 0 energy = method.scf() print(energy, -76.0259362997) method = density_fit(pyscf.scf.DHF(mol), 'weigend') energy = method.scf() print(energy, -76.0807386770) # normal DHF energy is -76.0815679438127 method = density_fit(pyscf.scf.UKS(mol), 'weigend', only_dfj = True) energy = method.scf() print(energy, -75.8547753298) mol.build( verbose = 0, atom = [["O" , (0. , 0. , 0.)], [1 , (0. , -0.757 , 0.587)], [1 , (0. , 0.757 , 0.587)] ], basis = 'ccpvdz', spin = 1, charge = 1, ) method = density_fit(pyscf.scf.UHF(mol), 'weigend') energy = method.scf() print(energy, -75.6310072359) method = density_fit(pyscf.scf.RHF(mol), 'weigend') energy = method.scf() print(energy, -75.6265157244)