#!/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>
#
'''
JK with discrete Fourier transformation
'''
import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc import tools
from pyscf.pbc.df.df_jk import _format_dms, _format_kpts_band, _format_jks
from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0
from pyscf.pbc.lib.kpts_helper import is_zero
[docs]
def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None):
'''Get the Coulomb (J) AO matrix at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray
Density matrix at each k-point. If a list of k-point DMs, eg,
UHF alpha and beta DM, the alpha and beta DMs are contracted
separately.
kpts : (nkpts, 3) ndarray
Kwargs:
kpts_band : (3,) ndarray or (*,3) ndarray
A list of arbitrary "band" k-points at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
or list of vj if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
mesh = mydf.mesh
assert cell.low_dim_ft_type != 'inf_vacuum'
assert cell.dimension > 1
ni = mydf._numint
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
make_rho, nset, nao = ni._gen_rho_evaluator(cell, dms, hermi)
coulG = tools.get_coulG(cell, mesh=mesh)
ngrids = len(coulG)
if hermi == 1 or is_zero(kpts):
vR = rhoR = np.zeros((nset,ngrids))
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
rhoR[i,p0:p1] += make_rho(i, ao_ks, mask, 'LDA').real
ao = ao_ks = None
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh).real
else: # vR may be complex if the underlying density is complex
vR = rhoR = np.zeros((nset,ngrids), dtype=np.complex128)
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
for k, ao in enumerate(ao_ks):
ao_dm = lib.dot(ao, dms[i,k])
rhoR[i,p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj())
rhoR *= 1./nkpts
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
weight = cell.vol / ngrids
vR *= weight
if is_zero(kpts_band):
vj_kpts = np.zeros((nset,nband,nao,nao))
else:
vj_kpts = np.zeros((nset,nband,nao,nao), dtype=np.complex128)
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
# ni.eval_mat can handle real vR only
# vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA')
for k, ao in enumerate(ao_ks):
aow = np.einsum('xi,x->xi', ao, vR[i,p0:p1])
vj_kpts[i,k] += lib.dot(ao.conj().T, aow)
return _format_jks(vj_kpts, dm_kpts, input_band, kpts)
[docs]
def get_j_e1_kpts(mydf, dm_kpts, kpts=np.zeros((1,3)), kpts_band=None):
'''Derivatives of Coulomb (J) AO matrix at sampled k-points.
'''
cell = mydf.cell
mesh = mydf.mesh
assert cell.low_dim_ft_type != 'inf_vacuum'
assert cell.dimension > 1
ni = mydf._numint
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
make_rho, nset, nao = ni._gen_rho_evaluator(cell, dms, hermi=1)
coulG = tools.get_coulG(cell, mesh=mesh)
ngrids = len(coulG)
if is_zero(kpts):
vR = rhoR = np.zeros((nset,ngrids))
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
rhoR[i,p0:p1] += make_rho(i, ao_ks, mask, 'LDA').real
ao = ao_ks = None
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh).real
else: # vR may be complex if the underlying density is complex
vR = rhoR = np.zeros((nset,ngrids), dtype=np.complex128)
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
for k, ao in enumerate(ao_ks):
ao_dm = lib.dot(ao, dms[i,k])
rhoR[i,p0:p1] += np.einsum('xi,xi->x', ao_dm, ao.conj())
rhoR *= 1./nkpts
for i in range(nset):
rhoG = tools.fft(rhoR[i], mesh)
vG = coulG * rhoG
vR[i] = tools.ifft(vG, mesh)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
weight = cell.vol / ngrids
vR *= weight
if is_zero(kpts_band):
vj_kpts = np.zeros((3,nset,nband,nao,nao))
else:
vj_kpts = np.zeros((3,nset,nband,nao,nao), dtype=np.complex128)
for ao_ks_etc, p0, p1 in mydf.aoR_loop(mydf.grids, kpts_band, deriv=1):
ao_ks, mask = ao_ks_etc[0], ao_ks_etc[2]
for i in range(nset):
# ni.eval_mat can handle real vR only
# vj_kpts[i] += ni.eval_mat(cell, ao_ks, 1., None, vR[i,p0:p1], mask, 'LDA')
for k, ao in enumerate(ao_ks):
aow = np.einsum('xi,x->xi', ao[0], vR[i,p0:p1])
vj_kpts[:,i,k] -= lib.einsum('axi,xj->aij', ao[1:].conj(), aow)
vj_kpts = np.asarray([_format_jks(vj, dm_kpts, input_band, kpts) for vj in vj_kpts])
return vj_kpts
[docs]
def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None,
exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray
Density matrix at each k-point
kpts : (nkpts, 3) ndarray
Kwargs:
hermi : int
Whether K matrix is hermitian
| 0 : not hermitian and not symmetric
| 1 : hermitian
kpts_band : (3,) ndarray or (*,3) ndarray
A list of arbitrary "band" k-points at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
vk : (nkpts, nao, nao) ndarray
or list of vj and vk if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
mesh = mydf.mesh
assert cell.low_dim_ft_type != 'inf_vacuum'
assert cell.dimension > 1
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if getattr(dm_kpts, 'mo_coeff', None) is not None:
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1./nkpts * (cell.vol/ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if is_zero(kpts_band) and is_zero(kpts):
vk_kpts = np.zeros((nset,nband,nao,nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((nset,nband,nao,nao), dtype=np.complex128)
coords = mydf.grids.coords
ao2_kpts = [np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)]
if input_band is None:
ao1_kpts = ao2_kpts
else:
ao1_kpts = [np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band)]
if mo_coeff is not None and nset == 1:
mo_coeff = [mo_coeff[k][:,occ>0] * np.sqrt(occ[occ>0])
for k, occ in enumerate(mo_occ)]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(min(nao, max(1, (max_memory-mem_now)*1e6/16/4/ngrids/nao)))
logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d',
max_memory, blksize)
#ao1_dtype = np.result_type(*ao1_kpts)
#ao2_dtype = np.result_type(*ao2_kpts)
vR_dm = np.empty((nset,nao,ngrids), dtype=vk_kpts.dtype)
t1 = (logger.process_clock(), logger.perf_counter())
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
naoj = ao2T.shape[0]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i,k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2-kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2-kpt1, exxdiv, mydf, mesh)
if is_zero(kpt1-kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2-kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('ig,jg->ijg', ao1T[p0:p1].conj()*expmikr, ao2T)
vG = tools.fft(rho1.reshape(-1,ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(p1-p0,naoj,ngrids)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i,p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[i,k1] += weight * lib.dot(vR_dm[i], ao1T.T)
t1 = logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)'%k2, *t1)
# Function _ewald_exxdiv_for_G0 to add back in the G=0 component to vk_kpts
# Note in the _ewald_exxdiv_for_G0 implementation, the G=0 treatments are
# different for 1D/2D and 3D systems. The special treatments for 1D and 2D
# can only be used with AFTDF/GDF/MDF method. In the FFTDF method, 1D, 2D
# and 3D should use the ewald probe charge correction.
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
[docs]
def get_k_e1_kpts(mydf, dm_kpts, kpts=np.zeros((1,3)), kpts_band=None,
exxdiv=None):
'''Derivatives of exchange (K) AO matrix at sampled k-points.
'''
cell = mydf.cell
mesh = mydf.mesh
assert cell.low_dim_ft_type != 'inf_vacuum'
assert cell.dimension > 1
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if getattr(dm_kpts, 'mo_coeff', None) is not None:
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1./nkpts * (cell.vol/ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if is_zero(kpts_band) and is_zero(kpts):
vk_kpts = np.zeros((3,nset,nband,nao,nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((3,nset,nband,nao,nao), dtype=np.complex128)
coords = mydf.grids.coords
if input_band is None:
ao2_kpts = [np.asarray(ao.transpose(0,2,1), order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts, deriv=1)]
ao1_kpts = ao2_kpts
ao2_kpts = [ao2_kpt[0] for ao2_kpt in ao2_kpts]
else:
ao2_kpts = [np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)]
ao1_kpts = [np.asarray(ao.transpose(0,2,1), order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band, deriv=1)]
if mo_coeff is not None and nset == 1:
mo_coeff = [mo_coeff[k][:,occ>0] * np.sqrt(occ[occ>0])
for k, occ in enumerate(mo_occ)]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(min(nao, max(1, (max_memory-mem_now)*1e6/16/4/3/ngrids/nao)))
logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d',
max_memory, blksize)
vR_dm = np.empty((3,nset,nao,ngrids), dtype=vk_kpts.dtype)
t1 = (logger.process_clock(), logger.perf_counter())
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
naoj = ao2T.shape[0]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i,k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2-kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2-kpt1, exxdiv, mydf, mesh)
if is_zero(kpt1-kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2-kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('aig,jg->aijg', ao1T[1:,p0:p1].conj()*expmikr, ao2T)
vG = tools.fft(rho1.reshape(-1,ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(3,p1-p0,naoj,ngrids)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('aijg,jg->aig', vR, ao_dms[i], out=vR_dm[:,i,p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[:,i,k1] -= weight * np.einsum('aig,jg->aij', vR_dm[:,i], ao1T[0])
t1 = logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)'%k2, *t1)
# Ewald correction has no contribution to nuclear gradient unless range separated Coulomb is used
# The gradient correction part is not added in the vk matrix
if exxdiv == 'ewald' and cell.omega!=0:
raise NotImplementedError("Range Separated Coulomb")
# when cell.omega !=0: madelung constant will have a non-zero derivative
vk_kpts = np.asarray([_format_jks(vk, dm_kpts, input_band, kpts) for vk in vk_kpts])
return vk_kpts
[docs]
def get_jk(mydf, dm, hermi=1, kpt=np.zeros(3), kpts_band=None,
with_j=True, with_k=True, exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices for the given density matrix.
Args:
dm : ndarray or list of ndarrays
A density matrix or a list of density matrices
Kwargs:
hermi : int
Whether J, K matrix is hermitian
| 0 : no hermitian or symmetric
| 1 : hermitian
| 2 : anti-hermitian
kpt : (3,) ndarray
The "inner" dummy k-point at which the DM was evaluated (or
sampled).
kpts_band : (3,) ndarray or (*,3) ndarray
The "outer" primary k-point at which J and K are evaluated.
Returns:
The function returns one J and one K matrix, corresponding to the input
density matrix (both order and shape).
'''
dm = np.asarray(dm, order='C')
vj = vk = None
if with_j:
vj = get_j(mydf, dm, hermi, kpt, kpts_band)
if with_k:
vk = get_k(mydf, dm, hermi, kpt, kpts_band, exxdiv)
return vj, vk
[docs]
def get_j(mydf, dm, hermi=1, kpt=np.zeros(3), kpts_band=None):
'''Get the Coulomb (J) AO matrix for the given density matrix.
Args:
dm : ndarray or list of ndarrays
A density matrix or a list of density matrices
Kwargs:
hermi : int
Whether J, K matrix is hermitian
| 0 : no hermitian or symmetric
| 1 : hermitian
| 2 : anti-hermitian
kpt : (3,) ndarray
The "inner" dummy k-point at which the DM was evaluated (or
sampled).
kpts_band : (3,) ndarray or (*,3) ndarray
The "outer" primary k-point at which J and K are evaluated.
Returns:
The function returns one J matrix, corresponding to the input
density matrix (both order and shape).
'''
dm = np.asarray(dm, order='C')
nao = dm.shape[-1]
dm_kpts = dm.reshape(-1,1,nao,nao)
vj = get_j_kpts(mydf, dm_kpts, hermi, kpt.reshape(1,3), kpts_band)
if kpts_band is None:
vj = vj[:,0,:,:]
if dm.ndim == 2:
vj = vj[0]
return vj
[docs]
def get_k(mydf, dm, hermi=1, kpt=np.zeros(3), kpts_band=None, exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices for the given density matrix.
Args:
dm : ndarray or list of ndarrays
A density matrix or a list of density matrices
Kwargs:
hermi : int
Whether J, K matrix is hermitian
| 0 : no hermitian or symmetric
| 1 : hermitian
| 2 : anti-hermitian
kpt : (3,) ndarray
The "inner" dummy k-point at which the DM was evaluated (or
sampled).
kpts_band : (3,) ndarray or (*,3) ndarray
The "outer" primary k-point at which J and K are evaluated.
Returns:
The function returns one J and one K matrix, corresponding to the input
density matrix (both order and shape).
'''
dm = np.asarray(dm, order='C')
nao = dm.shape[-1]
dm_kpts = dm.reshape(-1,1,nao,nao)
vk = get_k_kpts(mydf, dm_kpts, hermi, kpt.reshape(1,3), kpts_band, exxdiv)
if kpts_band is None:
vk = vk[:,0,:,:]
if dm.ndim == 2:
vk = vk[0]
return vk
