Skip to content

Commit

Permalink
Debug UKS gradient
Browse files Browse the repository at this point in the history
  • Loading branch information
@puzhichen
puzhichen committed Feb 19, 2025
1 parent 2942d41 commit b8bb071
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 26 deletions.
47 changes: 25 additions & 22 deletions gpu4pyscf/dft/xc_deriv.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -194,11 +194,12 @@ def transform_kxc(rho, fxc, kxc, xctype, spin=0):
vp = cupy.empty((2,nvar, 2,nvar, 2,nvar, ngrids)).transpose(1,3,5,0,2,4,6)
vp[0,0,0] = _stack_frrr(frrr)
i3 = np.arange(3)
#:qggg = _stack_fggg(fggg)
#:qggg = np.einsum('abcdefg,axg->xbcdefg', qggg, rho[:,1:4])
#:qggg = np.einsum('xbcdefg,cyg->xybdefg', qggg, rho[:,1:4])
#:qggg = np.einsum('xybdefg,ezg->xyzbdfg', qggg, rho[:,1:4])
qggg = _stack_fggg(fggg, rho=rho).transpose(1,3,5,0,2,4,6)
qggg = _stack_fggg(fggg)
qggg = cupy.einsum('abcdefg,axg->xbcdefg', qggg, rho[:,1:4])
qggg = cupy.einsum('xbcdefg,cyg->xybdefg', qggg, rho[:,1:4])
qggg = cupy.einsum('xybdefg,ezg->xyzbdfg', qggg, rho[:,1:4])
# qggg = _stack_fggg(fggg, rho=rho).transpose(1,3,5,0,2,4,6)
# qggg = cupy.asarray(qggg)
qgg = _stack_fgg(fgg)
qgg = cupy.einsum('abcdg,axg->xbcdg', qgg, rho[:,1:4])
for i in range(3):
Expand All @@ -208,31 +209,33 @@ def transform_kxc(rho, fxc, kxc, xctype, spin=0):
vp[1:4,1:4,1:4] = qggg

frgg = frgg.reshape(2,6,ngrids)
#:qrgg = _stack_fgg(frgg, axis=1)
#:qrgg = np.einsum('rabcdg,axg->xrbcdg', qrgg, rho[:,1:4])
#:qrgg = np.einsum('xrbcdg,cyg->xyrbdg', qrgg, rho[:,1:4])
qrgg = _stack_fgg(frgg, axis=1, rho=rho).transpose(2,4,0,1,3,5)
qrgg = _stack_fgg(frgg, axis=1)
qrgg = cupy.einsum('rabcdg,axg->xrbcdg', qrgg, rho[:,1:4])
qrgg = cupy.einsum('xrbcdg,cyg->xyrbdg', qrgg, rho[:,1:4])
# qrgg = _stack_fgg(frgg.get(), axis=1, rho=rho.get()).transpose(2,4,0,1,3,5)
qrg = _stack_fg(frg.reshape(2,3,ngrids), axis=1)
# qrgg = cupy.asarray(qrgg)
qrgg[i3,i3] += qrg
vp[0,1:4,1:4] = qrgg
vp[1:4,0,1:4] = qrgg.transpose(0,1,3,2,4,5)
vp[1:4,1:4,0] = qrgg.transpose(0,1,3,4,2,5)

frrg = frrg.reshape(3,3,ngrids)
qrrg = _stack_frr(frrg, axis=0)
#:qrrg = _stack_fg(qrrg, axis=2)
#:qrrg = np.einsum('rsabg,axg->rsxbg', qrrg, rho[:,1:4])
qrrg = _stack_fg(qrrg, axis=2, rho=rho).transpose(3,0,1,2,4)
qrrg = _stack_fg(qrrg, axis=2)
qrrg = cupy.einsum('rsabg,axg->rsxbg', qrrg, rho[:,1:4]).transpose(2,0,1,3,4)
# qrrg = _stack_fg(qrrg.get(), axis=2, rho=rho.get()).transpose(3,0,1,2,4)
# qrrg = cupy.asarray(qrrg)
vp[0,0,1:4] = qrrg
vp[0,1:4,0] = qrrg.transpose(0,1,3,2,4)
vp[1:4,0,0] = qrrg.transpose(0,3,1,2,4)

if order > 1:
fggt = fggt.reshape(6,2,ngrids)
#:qggt = _stack_fgg(fggt, axis=0)
#:qggt = np.einsum('abcdrg,axg->xbcdrg', qggt, rho[:,1:4])
#:qggt = np.einsum('xbcdrg,cyg->xybdrg', qggt, rho[:,1:4])
qggt = _stack_fgg(fggt, axis=0, rho=rho).transpose(1,3,0,2,4,5)
qggt = _stack_fgg(fggt, axis=0)
qggt = cupy.einsum('abcdrg,axg->xbcdrg', qggt, rho[:,1:4])
qggt = cupy.einsum('xbcdrg,cyg->xybdrg', qggt, rho[:,1:4])
# qggt = _stack_fgg(fggt, axis=0, rho=rho).transpose(1,3,0,2,4,5)
qgt = _stack_fg(fgt.reshape(3,2,ngrids), axis=0)
i3 = np.arange(3)
qggt[i3,i3] += qgt
Expand All @@ -241,17 +244,17 @@ def transform_kxc(rho, fxc, kxc, xctype, spin=0):
vp[4,1:4,1:4] = qggt.transpose(0,1,4,2,3,5)

qgtt = _stack_frr(fgtt.reshape(3,3,ngrids), axis=1)
#:qgtt = _stack_fg(qgtt, axis=0)
#:qgtt = np.einsum('abrsg,axg->xbrsg', qgtt, rho[:,1:4])
qgtt = _stack_fg(qgtt, axis=0, rho=rho).transpose(1,0,2,3,4)
qgtt = _stack_fg(qgtt, axis=0)
qgtt = cupy.einsum('abrsg,axg->xbrsg', qgtt, rho[:,1:4])
# qgtt = _stack_fg(qgtt, axis=0, rho=rho).transpose(1,0,2,3,4)
vp[1:4,4,4] = qgtt
vp[4,1:4,4] = qgtt.transpose(0,2,1,3,4)
vp[4,4,1:4] = qgtt.transpose(0,2,3,1,4)

frgt = frgt.reshape(2,3,2,ngrids)
#:qrgt = _stack_fg(frgt, axis=1)
#:qrgt = np.einsum('rabsg,axg->xrbsg', qrgt, rho[:,1:4])
qrgt = _stack_fg(frgt, axis=1, rho=rho).transpose(2,0,1,3,4)
qrgt = _stack_fg(frgt, axis=1)
qrgt = cupy.einsum('rabsg,axg->xrbsg', qrgt, rho[:,1:4])
# qrgt = _stack_fg(frgt, axis=1, rho=rho).transpose(2,0,1,3,4)
vp[0,1:4,4] = qrgt
vp[0,4,1:4] = qrgt.transpose(0,1,3,2,4)
vp[1:4,0,4] = qrgt.transpose(0,2,1,3,4)
Expand Down
9 changes: 5 additions & 4 deletions gpu4pyscf/tdscf/rhf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
]


def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None):
def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None, singlet=True):
r'''A and B matrices for TDDFT response function.
A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ai||jb)
Expand All @@ -43,7 +43,8 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None):
Ref: Chem Phys Lett, 256, 454
'''
from pyscf import ao2mo

if not singlet:
raise NotImplementedError('Only singlet is implemented')
if mo_energy is None: mo_energy = mf.mo_energy
if mo_coeff is None: mo_coeff = mf.mo_coeff
if mo_occ is None: mo_occ = mf.mo_occ
Expand Down Expand Up @@ -237,9 +238,9 @@ class TDBase(lib.StreamObject):
_finalize = tdhf_cpu.TDBase._finalize

gen_vind = NotImplemented
def get_ab(self, mf=None):
def get_ab(self, mf=None, singlet=True):
if mf is None: mf = self._scf
return get_ab(mf)
return get_ab(mf, singlet=singlet)

def get_precond(self, hdiag):
threshold_t=1.0e-4
Expand Down

0 comments on commit b8bb071

Please sign in to comment.