#!/usr/bin/env python
# Copyright 2014-2020 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>
#
'''
Numerical integration functions for 4-component or 2-component relativistic
methods with j-adapted AO basis
'''
import numpy
from pyscf import lib
from pyscf.dft import numint
from pyscf.dft.numint import _dot_ao_dm, _dot_ao_ao, _scale_ao, _tau_dot, BLKSIZE
from pyscf.dft.numint2c import _eval_xc_eff, mcfun_eval_xc_adapter
from pyscf.dft import xc_deriv
from pyscf import __config__
[docs]
def eval_ao(mol, coords, deriv=0, with_s=True, shls_slice=None,
non0tab=None, cutoff=None, out=None, verbose=None):
'''Evaluates the value of 2-component or 4-component j-adapted basis on grids.
'''
comp = (deriv+1)*(deriv+2)*(deriv+3)//6
feval = 'GTOval_spinor_deriv%d' % deriv
if not with_s:
# aoLa, aoLb = aoL
ao = aoL = mol.eval_gto(feval, coords, comp, shls_slice, non0tab,
cutoff=cutoff, out=out)
else:
assert (deriv <= 1) # only GTOval_ipsp_spinor
ngrids = coords.shape[0]
nao = mol.nao_2c()
ao = numpy.ndarray((4,comp,nao,ngrids), dtype=numpy.complex128, buffer=out)
aoL = mol.eval_gto(feval, coords, comp, shls_slice, non0tab, out=ao[:2]) # noqa
ao = ao.transpose(0,1,3,2)
aoS = ao[2:]
aoSa, aoSb = aoS
feval_gto = ['GTOval_sp_spinor', 'GTOval_ipsp_spinor']
p1 = 0
for n in range(deriv+1):
comp = (n+1)*(n+2)//2
p0, p1 = p1, p1 + comp
aoSa[p0:p1], aoSb[p0:p1] = mol.eval_gto(
feval_gto[n], coords, comp, shls_slice, non0tab, cutoff=cutoff)
if deriv == 0:
ao = ao[:,0]
return ao
def _dot_spinor_dm(mol, ket, dm, non0tab, shls_slice, ao_loc):
ket_a, ket_b = ket
outa = _dot_ao_dm(mol, ket_a, dm, non0tab, shls_slice, ao_loc)
outb = _dot_ao_dm(mol, ket_b, dm, non0tab, shls_slice, ao_loc)
return outa, outb
def _contract_rho_m(bra, ket, hermi=0, bra_eq_ket=False):
'''
hermi indicates whether the density matrix is hermitian.
bra_eq_ket indicates whether bra and ket basis are the same AOs.
'''
ket_a, ket_b = ket
bra_a, bra_b = bra
ngrids = ket_a.shape[0]
if hermi:
raa = numpy.einsum('pi,pi->p', bra_a.real, ket_a.real)
raa+= numpy.einsum('pi,pi->p', bra_a.imag, ket_a.imag)
rab = numpy.einsum('pi,pi->p', bra_a.conj(), ket_b)
rbb = numpy.einsum('pi,pi->p', bra_b.real, ket_b.real)
rbb+= numpy.einsum('pi,pi->p', bra_b.imag, ket_b.imag)
rho_m = numpy.empty((4, ngrids))
rho_m[0,:] = raa + rbb # rho
rho_m[1,:] = rab.real # mx
rho_m[2,:] = rab.imag # my
rho_m[3,:] = raa - rbb # mz
if bra_eq_ket:
rho_m[1,:] *= 2
rho_m[2,:] *= 2
else:
rba = numpy.einsum('pi,pi->p', bra_b.conj(), ket_a)
rho_m[1,:] += rba.real
rho_m[2,:] -= rba.imag
else:
raa = numpy.einsum('pi,pi->p', bra_a.conj(), ket_a)
rab = numpy.einsum('pi,pi->p', bra_a.conj(), ket_b)
rba = numpy.einsum('pi,pi->p', bra_b.conj(), ket_a)
rbb = numpy.einsum('pi,pi->p', bra_b.conj(), ket_b)
rho_m = numpy.empty((4, ngrids), dtype=numpy.complex128)
rho_m[0,:] = raa + rbb # rho
rho_m[1,:] = rba + rab # mx
rho_m[2,:] = (rba - rab) * 1j # my
rho_m[3,:] = raa - rbb # mz
return rho_m
def _eval_rho_2c(mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, with_lapl=False):
aoa, aob = ao
ngrids, nao = aoa.shape[-2:]
xctype = xctype.upper()
if non0tab is None:
non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
dtype=numpy.uint8)
shls_slice = (0, mol.nbas)
ao_loc = mol.ao_loc_2c()
if xctype == 'LDA' or xctype == 'HF':
c0 = _dot_spinor_dm(mol, ao, dm, non0tab, shls_slice, ao_loc)
rho_m = _contract_rho_m(ao, c0, hermi, True)
elif xctype == 'GGA':
# first 4 ~ (rho, m), second 4 ~ (0th order, dx, dy, dz)
if hermi:
rho_m = numpy.empty((4, 4, ngrids))
else:
rho_m = numpy.empty((4, 4, ngrids), dtype=numpy.complex128)
c0 = _dot_spinor_dm(mol, ao[:,0], dm, non0tab, shls_slice, ao_loc)
rho_m[:,0] = _contract_rho_m(ao[:,0], c0, hermi, True)
for i in range(1, 4):
rho_m[:,i] = _contract_rho_m(ao[:,i], c0, hermi, False)
if hermi:
rho_m[:,1:4] *= 2 # *2 for |ao> dm < dx ao| + |dx ao> dm < ao|
else:
for i in range(1, 4):
c1 = _dot_spinor_dm(mol, ao[:,i], dm, non0tab, shls_slice, ao_loc)
rho_m[:,i] += _contract_rho_m(ao[:,0], c1, hermi, False)
else: # meta-GGA
if hermi:
dtype = numpy.double
else:
dtype = numpy.complex128
if with_lapl:
rho_m = numpy.empty((4, 6, ngrids), dtype=dtype)
tau_idx = 5
else:
rho_m = numpy.empty((4, 5, ngrids), dtype=dtype)
tau_idx = 4
c0 = _dot_spinor_dm(mol, ao[:,0], dm, non0tab, shls_slice, ao_loc)
rho_m[:,0] = _contract_rho_m(ao[:,0], c0, hermi, True)
rho_m[:,tau_idx] = 0
for i in range(1, 4):
c1 = _dot_spinor_dm(mol, ao[:,i], dm, non0tab, shls_slice, ao_loc)
rho_m[:,tau_idx] += _contract_rho_m(ao[:,i], c1, hermi, True)
rho_m[:,i] = _contract_rho_m(ao[:,i], c0, hermi, False)
if hermi:
rho_m[:,i] *= 2
else:
rho_m[:,i] += _contract_rho_m(ao[:,0], c1, hermi, False)
if with_lapl:
raise NotImplementedError
# tau = 1/2 (\nabla f)^2
rho_m[:,tau_idx] *= .5
return rho_m
[docs]
def eval_rho(mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, with_lapl=False,
verbose=None):
r'''Calculate the electron density and magnetization spin density with
j-adapted spinor basis.
''' + numint.eval_rho.__doc__
with_s = len(ao) == 4 # aoLa, aoLb, aoSa, aoSb
if with_s:
n2c = dm.shape[1] // 2
dmLL = dm[:n2c,:n2c].copy()
dmSS = dm[n2c:,n2c:].copy()
c1 = .5 / lib.param.LIGHT_SPEED
rho = _eval_rho_2c(mol, ao[:2], dmLL, non0tab, xctype, hermi, with_lapl)
rhoS = _eval_rho_2c(mol, ao[2:], dmSS, non0tab, xctype, hermi, with_lapl)
rhoS *= c1**2
rho[0] += rhoS[0]
# M = |\beta\Sigma|
rho[1:4] -= rhoS[1:4]
return rho
else:
return _eval_rho_2c(mol, ao, dm, non0tab, xctype, hermi, with_lapl)
def _ncol_lda_vxc_mat(mol, ao, weight, rho_tm, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of non-collinear LDA'''
# NOTE vxc in u/d representation
aoa, aob = ao
r, mx, my, mz = rho_tm
vxc = xc_deriv.ud2ts(vxc)
vr, vs = vxc[:,0]
s = lib.norm(rho_tm[1:4], axis=0)
wv = weight * vr
with numpy.errstate(divide='ignore',invalid='ignore'):
ws = vs * weight / s
ws[s < 1e-20] = 0
if not on_LL:
# flip the sign for small components because
# M = \beta\Sigma = [sigma, 0 ]
# [0 , -sigma]
ws *= -1
# * .5 because of v+v.conj().T in r_vxc
if hermi:
wv *= .5
ws *= .5
# einsum('g,g,xgi,xgj->ij', vr, weight, ao, ao)
# + einsum('xy,g,g,xgi,ygj->ij', sx, vs*m[0]/s, weight, ao, ao)
# + einsum('xy,g,g,xgi,ygj->ij', sy, vs*m[1]/s, weight, ao, ao)
# + einsum('xy,g,g,xgi,ygj->ij', sz, vs*m[2]/s, weight, ao, ao)
aow = None
aow = _scale_ao(aoa, ws*(mx+my*1j), out=aow)
mat = _dot_ao_ao(mol, aob, aow, mask, shls_slice, ao_loc)
if hermi:
mat = mat + mat.conj().T
else:
aow = _scale_ao(aob, ws*(mx-my*1j), out=aow) # Mx
mat+= _dot_ao_ao(mol, aoa, aow, mask, shls_slice, ao_loc)
aow = _scale_ao(aoa, wv+ws*mz, out=aow) # Mz
mat+= _dot_ao_ao(mol, aoa, aow, mask, shls_slice, ao_loc)
aow = _scale_ao(aob, wv-ws*mz, out=aow) # Mz
mat+= _dot_ao_ao(mol, aob, aow, mask, shls_slice, ao_loc)
return mat
def _ncol_gga_vxc_mat(mol, ao, weight, rho, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of non-collinear GGA'''
raise NotImplementedError
def _ncol_mgga_vxc_mat(mol, ao, weight, rho, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of non-collinear MGGA'''
raise NotImplementedError
def _col_lda_vxc_mat(mol, ao, weight, rho, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of collinear LDA'''
# NOTE vxc in u/d representation
aoa, aob = ao
wv = weight * vxc
if hermi:
wv *= .5
if on_LL:
wva, wvb = wv
else: # for SS block
# v_rho = (vxc_a + vxc_b) * .5
# v_mz = (vxc_a - vxc_b) * .5
# For small components, M = \beta\Sigma leads to
# (v_rho - sigma_z*v_mz) = [vxc_b, 0 ]
# [0 , vxc_a]
wvb, wva = wv
mat = _dot_ao_ao(mol, aoa, _scale_ao(aoa, wva[0]), mask, shls_slice, ao_loc)
mat += _dot_ao_ao(mol, aob, _scale_ao(aob, wvb[0]), mask, shls_slice, ao_loc)
return mat
def _col_gga_vxc_mat(mol, ao, weight, rho, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of collinear GGA'''
# NOTE vxc in u/d representation
aoa, aob = ao
wv = weight * vxc
if hermi:
wv[:,0] *= .5
if on_LL:
wva, wvb = wv
else: # for SS block
wvb, wva = wv
aow = None
aow = _scale_ao(aoa[:4], wva[:4], out=aow)
mat = _dot_ao_ao(mol, aoa[0], aow, mask, shls_slice, ao_loc)
aow = _scale_ao(aob[:4], wvb[:4], out=aow)
mat+= _dot_ao_ao(mol, aob[0], aow, mask, shls_slice, ao_loc)
if hermi != 1:
aow = _scale_ao(aoa[1:4], wva[1:4].conj(), out=aow)
mat+= _dot_ao_ao(mol, aow, aoa[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aob[1:4], wvb[1:4].conj(), out=aow)
mat+= _dot_ao_ao(mol, aow, aob[0], mask, shls_slice, ao_loc)
return mat
def _col_mgga_vxc_mat(mol, ao, weight, rho, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of collinear MGGA'''
# NOTE vxc in u/d representation
aoa, aob = ao
wv = weight * vxc
tau_idx = 4
wv[:,tau_idx] *= .5 # *.5 for 1/2 in tau
if hermi:
wv[:,0] *= .5
wv[:,tau_idx] *= .5
if on_LL:
wva, wvb = wv
else:
wvb, wva = wv
aow = None
aow = _scale_ao(aoa[:4], wva[:4], out=aow)
mat = _dot_ao_ao(mol, aoa[0], aow, mask, shls_slice, ao_loc)
aow = _scale_ao(aob[:4], wvb[:4], out=aow)
mat+= _dot_ao_ao(mol, aob[0], aow, mask, shls_slice, ao_loc)
if hermi != 1:
aow = _scale_ao(aoa[1:4], wva[1:4].conj(), out=aow)
mat += _dot_ao_ao(mol, aow, aoa[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aob[1:4], wvb[1:4].conj(), out=aow)
mat += _dot_ao_ao(mol, aow, aob[0], mask, shls_slice, ao_loc)
mat += _tau_dot(mol, aoa, aoa, wva[tau_idx], mask, shls_slice, ao_loc)
mat += _tau_dot(mol, aob, aob, wvb[tau_idx], mask, shls_slice, ao_loc)
return mat
def _mcol_lda_vxc_mat(mol, ao, weight, rho, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of multi-collinear LDA'''
# NOTE vxc in t/m representation
aoa, aob = ao
wv = weight * vxc
if hermi:
wv *= .5 # * .5 because of v+v.conj().T in r_vxc
if not on_LL:
# flip the sign for small components because
# M = \beta\Sigma = [sigma, 0 ]
# [0 , -sigma]
wv[1:4] *= -1
wr, wmx, wmy, wmz = wv
# einsum('g,g,xgi,xgj->ij', vxc, weight, ao, ao)
# + einsum('xy,g,g,xgi,ygj->ij', sx, vxc, weight, ao, ao)
# + einsum('xy,g,g,xgi,ygj->ij', sy, vxc, weight, ao, ao)
# + einsum('xy,g,g,xgi,ygj->ij', sz, vxc, weight, ao, ao)
aow = None
aow = _scale_ao(aoa, wmx[0]+wmy[0]*1j, out=aow)
mat = _dot_ao_ao(mol, aob, aow, mask, shls_slice, ao_loc)
if hermi:
assert vxc.dtype == numpy.double
mat = mat + mat.conj().T
else:
aow = _scale_ao(aob, wmx[0]-wmy[0]*1j, out=aow)
mat+= _dot_ao_ao(mol, aoa, aow, mask, shls_slice, ao_loc)
aow = _scale_ao(aoa, wr[0]+wmz[0], out=aow) # Mz
mat+= _dot_ao_ao(mol, aoa, aow, mask, shls_slice, ao_loc)
aow = _scale_ao(aob, wr[0]-wmz[0], out=aow) # Mz
mat+= _dot_ao_ao(mol, aob, aow, mask, shls_slice, ao_loc)
return mat
def _mcol_gga_vxc_mat(mol, ao, weight, rho, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of multi-collinear LDA'''
# NOTE vxc in t/m representation
aoa, aob = ao
wv = weight * vxc
if hermi:
wv[:,0] *= .5
if not on_LL:
# flip the sign for small components because
# M = \beta\Sigma = [sigma, 0 ]
# [0 , -sigma]
wv[1:4] *= -1
wr, wmx, wmy, wmz = wv
aow = None
aow = _scale_ao(aoa[:4], wmx[:4]+wmy[:4]*1j, out=aow)
mat = _dot_ao_ao(mol, aob[0], aow, mask, shls_slice, ao_loc)
if hermi:
assert vxc.dtype == numpy.double
mat = mat + mat.conj().T
else:
aow = _scale_ao(aob[:4], wmx[:4]-wmy[:4]*1j, out=aow)
mat+= _dot_ao_ao(mol, aoa[0], aow, mask, shls_slice, ao_loc)
aow = _scale_ao(aob[1:4], (wmx[1:4]+wmy[1:4]*1j).conj(), out=aow)
mat+= _dot_ao_ao(mol, aow, aoa[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aoa[1:4], (wmx[1:4]-wmy[1:4]*1j).conj(), out=aow)
mat+= _dot_ao_ao(mol, aow, aob[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aoa[1:4], (wr[1:4]+wmz[1:4]).conj(), out=aow) # Mz
mat+= _dot_ao_ao(mol, aow, aoa[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aob[1:4], (wr[1:4]-wmz[1:4]).conj(), out=aow) # Mz
mat+= _dot_ao_ao(mol, aow, aob[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aoa, wr[:4]+wmz[:4], out=aow) # Mz
mat+= _dot_ao_ao(mol, aoa[0], aow, mask, shls_slice, ao_loc)
aow = _scale_ao(aob, wr[:4]-wmz[:4], out=aow) # Mz
mat+= _dot_ao_ao(mol, aob[0], aow, mask, shls_slice, ao_loc)
return mat
def _mcol_mgga_vxc_mat(mol, ao, weight, rho, vxc, mask, shls_slice, ao_loc,
hermi, on_LL=True):
'''Vxc matrix of multi-collinear MGGA'''
# NOTE vxc in t/m representation
aoa, aob = ao
wv = weight * vxc
tau_idx = 4
wv[:,tau_idx] *= .5 # *.5 for 1/2 in tau
if hermi:
wv[:,0] *= .5
wv[:,tau_idx] *= .5
if not on_LL:
# flip the sign for small components because
# M = \beta\Sigma = [sigma, 0 ]
# [0 , -sigma]
wv[1:4] *= -1
wr, wmx, wmy, wmz = wv
aow = None
aow = _scale_ao(aoa[:4], wmx[:4]+wmy[:4]*1j, out=aow)
mat = _dot_ao_ao(mol, aob[0], aow, mask, shls_slice, ao_loc)
mat += _tau_dot(mol, aob, aoa, wmx[tau_idx]+wmy[tau_idx]*1j, mask,
shls_slice, ao_loc)
if hermi:
assert vxc.dtype == numpy.double
mat = mat + mat.conj().T
else:
aow = _scale_ao(aob[:4], wmx[:4]-wmy[:4]*1j, out=aow)
mat+= _dot_ao_ao(mol, aoa[0], aow, mask, shls_slice, ao_loc)
mat+= _tau_dot(mol, aoa, aob, wmx[tau_idx]-wmy[tau_idx]*1j, mask,
shls_slice, ao_loc)
aow = _scale_ao(aob[1:4], (wmx[1:4]+wmy[1:4]*1j).conj(), out=aow)
mat+= _dot_ao_ao(mol, aow, aoa[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aoa[1:4], (wmx[1:4]-wmy[1:4]*1j).conj(), out=aow)
mat+= _dot_ao_ao(mol, aow, aob[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aoa[1:4], (wr[1:4]+wmz[1:4]).conj(), out=aow) # Mz
mat+= _dot_ao_ao(mol, aow, aoa[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aob[1:4], (wr[1:4]-wmz[1:4]).conj(), out=aow) # Mz
mat+= _dot_ao_ao(mol, aow, aob[0], mask, shls_slice, ao_loc)
aow = _scale_ao(aoa, wr[:4]+wmz[:4], out=aow) # Mz
mat+= _dot_ao_ao(mol, aoa[0], aow, mask, shls_slice, ao_loc)
mat+= _tau_dot(mol, aoa, aoa, wr[tau_idx]+wmz[tau_idx], mask, shls_slice, ao_loc)
aow = _scale_ao(aob, wr[:4]-wmz[:4], out=aow) # Mz
mat+= _dot_ao_ao(mol, aob[0], aow, mask, shls_slice, ao_loc)
mat+= _tau_dot(mol, aob, aob, wr[tau_idx]-wmz[tau_idx], mask, shls_slice, ao_loc)
return mat
[docs]
def r_vxc(ni, mol, grids, xc_code, dms, spin=0, relativity=1, hermi=1,
max_memory=2000, verbose=None):
'''Calculate 2-component or 4-component Vxc matrix in j-adapted basis.
''' + numint.nr_rks.__doc__
xctype = ni._xc_type(xc_code)
shls_slice = (0, mol.nbas)
ao_loc = mol.ao_loc_2c()
n2c = ao_loc[-1]
make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
with_s = (nao == n2c*2) # 4C DM
nelec = numpy.zeros(nset)
excsum = numpy.zeros(nset)
matLL = numpy.zeros((nset,n2c,n2c), dtype=numpy.complex128)
matSS = numpy.zeros_like(matLL)
if xctype in ('LDA', 'GGA', 'MGGA'):
f_eval_mat = {
('LDA' , 'c'): (_col_lda_vxc_mat , 0),
('GGA' , 'c'): (_col_gga_vxc_mat , 1),
('MGGA', 'c'): (_col_mgga_vxc_mat , 1),
('LDA' , 'n'): (_ncol_lda_vxc_mat , 0),
('GGA' , 'n'): (_ncol_gga_vxc_mat , 1),
('MGGA', 'n'): (_ncol_mgga_vxc_mat, 1),
('LDA' , 'm'): (_mcol_lda_vxc_mat , 0),
('GGA' , 'm'): (_mcol_gga_vxc_mat , 1),
('MGGA', 'm'): (_mcol_mgga_vxc_mat, 1),
}
fmat, ao_deriv = f_eval_mat[(xctype, ni.collinear[0])]
if ni.collinear[0] == 'm': # mcol
eval_xc = ni.mcfun_eval_xc_adapter(xc_code)
else:
eval_xc = ni.eval_xc_eff
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
with_s=with_s):
for i in range(nset):
rho = make_rho(i, ao, mask, xctype)
exc, vxc = eval_xc(xc_code, rho, deriv=1, xctype=xctype)[:2]
if xctype == 'LDA':
den = rho[0] * weight
else:
den = rho[0,0] * weight
nelec[i] += den.sum()
excsum[i] += numpy.dot(den, exc)
matLL[i] += fmat(mol, ao[:2], weight, rho, vxc,
mask, shls_slice, ao_loc, hermi, True)
if with_s:
matSS[i] += fmat(mol, ao[2:], weight, rho, vxc,
mask, shls_slice, ao_loc, hermi, False)
if hermi:
matLL = matLL + matLL.conj().transpose(0,2,1)
if with_s:
matSS = matSS + matSS.conj().transpose(0,2,1)
elif xctype == 'HF':
pass
else:
raise NotImplementedError(f'r_vxc for functional {xc_code}')
if with_s:
matSS *= (.5 / lib.param.LIGHT_SPEED)**2
vmat = numpy.zeros((nset,nao,nao), dtype=numpy.complex128)
vmat[:,:n2c,:n2c] = matLL
vmat[:,n2c:,n2c:] = matSS
else:
vmat = matLL
if nset == 1:
nelec = nelec[0]
excsum = excsum[0]
return nelec, excsum, vmat.reshape(dms.shape)
def _col_rho_tm2ud(rho_tm):
return numpy.stack([(rho_tm[0] + rho_tm[3]) * .5,
(rho_tm[0] - rho_tm[3]) * .5,])
def _col_lda_fxc_mat(mol, ao, weight, rho0, rho1, fxc,
mask, shls_slice, ao_loc, hermi, on_LL=True):
'''Kernel matrix of collinear LDA'''
rho1 = _col_rho_tm2ud(rho1)
vxc1 = numpy.einsum('ag,abyg->byg', rho1, fxc[:,0,:])
return _col_lda_vxc_mat(mol, ao, weight, rho0, vxc1, mask, shls_slice,
ao_loc, hermi, on_LL)
def _col_gga_fxc_mat(mol, ao, weight, rho0, rho1, fxc,
mask, shls_slice, ao_loc, hermi, on_LL=True):
'''Kernel matrix of collinear GGA'''
rho1 = _col_rho_tm2ud(rho1)
vxc1 = numpy.einsum('axg,axbyg->byg', rho1, fxc)
return _col_gga_vxc_mat(mol, ao, weight, rho0, vxc1, mask, shls_slice,
ao_loc, hermi, on_LL)
def _col_mgga_fxc_mat(mol, ao, weight, rho0, rho1, fxc,
mask, shls_slice, ao_loc, hermi, on_LL=True):
'''Kernel matrix of collinear MGGA'''
rho1 = _col_rho_tm2ud(rho1)
vxc1 = numpy.einsum('axg,axbyg->byg', rho1, fxc)
return _col_mgga_vxc_mat(mol, ao, weight, rho0, vxc1, mask, shls_slice,
ao_loc, hermi, on_LL)
def _mcol_lda_fxc_mat(mol, ao, weight, rho0, rho1, fxc,
mask, shls_slice, ao_loc, hermi, on_LL=True):
'''Kernel matrix of multi-collinear LDA'''
vxc1 = numpy.einsum('ag,abyg->byg', rho1, fxc[:,0,:])
return _mcol_lda_vxc_mat(mol, ao, weight, rho0, vxc1, mask, shls_slice,
ao_loc, hermi, on_LL)
def _mcol_gga_fxc_mat(mol, ao, weight, rho0, rho1, fxc,
mask, shls_slice, ao_loc, hermi, on_LL=True):
'''Kernel matrix of multi-collinear GGA'''
vxc1 = numpy.einsum('axg,axbyg->byg', rho1, fxc)
return _mcol_gga_vxc_mat(mol, ao, weight, rho0, vxc1, mask, shls_slice,
ao_loc, hermi, on_LL)
def _mcol_mgga_fxc_mat(mol, ao, weight, rho0, rho1, fxc,
mask, shls_slice, ao_loc, hermi, on_LL=True):
'''Kernel matrix of multi-collinear MGGA'''
vxc1 = numpy.einsum('axg,axbyg->byg', rho1, fxc)
return _mcol_mgga_vxc_mat(mol, ao, weight, rho0, vxc1, mask, shls_slice,
ao_loc, hermi, on_LL)
[docs]
def r_fxc(ni, mol, grids, xc_code, dm0, dms, spin=0, relativity=1, hermi=0,
rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
'''Calculate 2-component or 4-component Vxc matrix in j-adapted basis.
''' + numint.nr_rks_fxc.__doc__
xctype = ni._xc_type(xc_code)
shls_slice = (0, mol.nbas)
ao_loc = mol.ao_loc_2c()
n2c = ao_loc[-1]
if ni.collinear[0] not in ('c', 'm'): # col or mcol
raise NotImplementedError('non-collinear fxc')
if fxc is None and xctype in ('LDA', 'GGA', 'MGGA'):
fxc = ni.cache_xc_kernel1(mol, grids, xc_code, dm0,
max_memory=max_memory)[2]
make_rho1, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
with_s = (nao == n2c*2) # 4C DM
matLL = numpy.zeros((nset,n2c,n2c), dtype=numpy.complex128)
matSS = numpy.zeros_like(matLL)
if xctype in ('LDA', 'GGA', 'MGGA'):
f_eval_mat = {
('LDA' , 'c'): (_col_lda_fxc_mat , 0),
('GGA' , 'c'): (_col_gga_fxc_mat , 1),
('MGGA', 'c'): (_col_mgga_fxc_mat , 1),
('LDA' , 'm'): (_mcol_lda_fxc_mat , 0),
('GGA' , 'm'): (_mcol_gga_fxc_mat , 1),
('MGGA', 'm'): (_mcol_mgga_fxc_mat , 1),
}
fmat, ao_deriv = f_eval_mat[(xctype, ni.collinear[0])]
_rho0 = None
p1 = 0
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory,
with_s=with_s):
p0, p1 = p1, p1 + weight.size
_fxc = fxc[:,:,:,:,p0:p1]
for i in range(nset):
rho1 = make_rho1(i, ao, mask, xctype)
matLL[i] += fmat(mol, ao[:2], weight, _rho0, rho1, _fxc,
mask, shls_slice, ao_loc, hermi)
if with_s:
matSS[i] += fmat(mol, ao[2:], weight, _rho0, rho1, _fxc,
mask, shls_slice, ao_loc, hermi, False)
if hermi:
# for (\nabla\mu) \nu + \mu (\nabla\nu)
matLL = matLL + matLL.conj().transpose(0,2,1)
if with_s:
matSS = matSS + matSS.conj().transpose(0,2,1)
elif xctype == 'HF':
pass
else:
raise NotImplementedError(f'r_fxc for functional {xc_code}')
if with_s:
matSS *= (.5 / lib.param.LIGHT_SPEED)**2
vmat = numpy.zeros((nset,nao,nao), dtype=numpy.complex128)
vmat[:,:n2c,:n2c] = matLL
vmat[:,n2c:,n2c:] = matSS
else:
vmat = matLL
if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
vmat = vmat[0]
return vmat
[docs]
def cache_xc_kernel(ni, mol, grids, xc_code, mo_coeff, mo_occ, spin=1,
max_memory=2000):
'''Compute the 0th order density, Vxc and fxc. They can be used in TDDFT,
DFT hessian module etc.
'''
dm = numpy.dot(mo_coeff * mo_occ, mo_coeff.conj().T)
return cache_xc_kernel1(ni, mol, grids, xc_code, dm, spin, max_memory)
[docs]
def cache_xc_kernel1(ni, mol, grids, xc_code, dm, spin=1, max_memory=2000):
xctype = ni._xc_type(xc_code)
if xctype in ('GGA', 'MGGA'):
ao_deriv = 1
else:
ao_deriv = 0
hermi = 1 # rho must be real. We need to assume dm hermitian
make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm, hermi)
n2c = mol.nao_2c()
with_s = (nao == n2c*2) # 4C DM
rho = []
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, ao_deriv, max_memory, with_s=with_s):
rho.append(make_rho(0, ao, mask, xctype))
rho = numpy.concatenate(rho,axis=-1)
if ni.collinear[0] == 'm': # mcol
eval_xc = ni.mcfun_eval_xc_adapter(xc_code)
else:
eval_xc = ni.eval_xc_eff
vxc, fxc = eval_xc(xc_code, rho, deriv=2, xctype=xctype)[1:3]
return rho, vxc, fxc
[docs]
def get_rho(ni, mol, dm, grids, max_memory=2000):
make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm, hermi=1)
n2c = mol.nao_2c()
with_s = (nao == n2c*2) # 4C DM
rho = numpy.empty(grids.weights.size)
p1 = 0
for ao, mask, weight, coords \
in ni.block_loop(mol, grids, nao, 0, max_memory, with_s=with_s):
p0, p1 = p1, p1 + weight.size
rho[p0:p1] = make_rho(0, ao, mask, 'LDA')[0]
return rho
[docs]
class RNumInt(lib.StreamObject, numint.LibXCMixin):
'''NumInt for j-adapted (spinor) basis'''
# collinear schemes:
# 'col' (collinear, by default)
# 'ncol' (non-collinear)
# 'mcol' (multi-collinear)
collinear = getattr(__config__, 'dft_numint_RnumInt_collinear', 'col')
spin_samples = getattr(__config__, 'dft_numint_RnumInt_spin_samples', 770)
collinear_thrd = getattr(__config__, 'dft_numint_RnumInt_collinear_thrd', 0.99)
collinear_samples = getattr(__config__, 'dft_numint_RnumInt_collinear_samples', 200)
get_rho = get_rho
cache_xc_kernel = cache_xc_kernel
cache_xc_kernel1 = cache_xc_kernel1
get_vxc = r_vxc = r_vxc
get_fxc = r_fxc = r_fxc
eval_xc_eff = _eval_xc_eff
mcfun_eval_xc_adapter = mcfun_eval_xc_adapter
eval_ao = staticmethod(eval_ao)
eval_rho = staticmethod(eval_rho)
[docs]
def eval_rho1(self, mol, ao, dm, screen_index=None, xctype='LDA', hermi=0,
with_lapl=True, cutoff=None, ao_cutoff=None, pair_mask=None,
verbose=None):
return self.eval_rho(mol, ao, dm, screen_index, xctype, hermi,
with_lapl, verbose=verbose)
[docs]
def eval_rho2(self, mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
with_lapl=True, verbose=None):
raise NotImplementedError
[docs]
def block_loop(self, mol, grids, nao, deriv=0, max_memory=2000,
non0tab=None, blksize=None, buf=None, with_s=False):
'''Define this macro to loop over grids by blocks.
'''
if grids.coords is None:
grids.build(with_non0tab=True)
ngrids = grids.weights.size
comp = (deriv+1)*(deriv+2)*(deriv+3)//6
# NOTE to index ni.non0tab, the blksize needs to be the integer multiplier of BLKSIZE
if blksize is None:
blksize = min(int(max_memory*1e6/((comp*4+4)*nao*16*BLKSIZE))*BLKSIZE, ngrids)
blksize = max(blksize, BLKSIZE)
if non0tab is None:
non0tab = grids.non0tab
if non0tab is None:
non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
dtype=numpy.uint8)
if buf is None:
buf = numpy.empty((4,comp,blksize,nao), dtype=numpy.complex128)
for ip0 in range(0, ngrids, blksize):
ip1 = min(ngrids, ip0+blksize)
coords = grids.coords[ip0:ip1]
weight = grids.weights[ip0:ip1]
non0 = non0tab[ip0//BLKSIZE:]
ao = self.eval_ao(mol, coords, deriv=deriv, with_s=with_s,
non0tab=non0, cutoff=grids.cutoff, out=buf)
yield ao, non0, weight, coords
def _gen_rho_evaluator(self, mol, dms, hermi=1, with_lapl=False, grids=None):
dms = numpy.asarray(dms, order='C')
if dms.ndim == 2:
dms = dms[numpy.newaxis]
ndms, nao = dms.shape[:2]
def make_rho(idm, ao, non0tab, xctype):
return self.eval_rho(mol, ao, dms[idm], non0tab, xctype, hermi, with_lapl)
return make_rho, ndms, nao
_RNumInt = RNumInt
