Skip to content

Commit

Permalink
merge master
Browse files Browse the repository at this point in the history
  • Loading branch information
@henryw7
henryw7 committed Dec 4, 2024
2 parents e1ae69e + 28323a5 commit b0ff950
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 10 deletions.
23 changes: 22 additions & 1 deletion gpu4pyscf/gto/moleintor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,23 @@ def get_n_hermite_density_of_angular_pair(l):
else:
self.ao_loc = self.sph_ao_loc

def sort_orbitals(self, mat, axis=[]):
''' Transform given axis of a matrix into sorted AO,
and transform given auxiliary axis of a matrix into sorted auxiliary AO
'''
idx = self._ao_idx
shape_ones = (1,) * mat.ndim
fancy_index = []
for dim, n in enumerate(mat.shape):
if dim in axis:
assert n == len(idx)
indices = idx
else:
indices = np.arange(n)
idx_shape = shape_ones[:dim] + (-1,) + shape_ones[dim+1:]
fancy_index.append(indices.reshape(idx_shape))
return mat[tuple(fancy_index)]

@property
def bpcache(self):
device_id = cp.cuda.Device().id
Expand Down Expand Up @@ -349,6 +366,10 @@ def get_int3c1e_density_contracted(mol, grids, charge_exponents, dm, intopt):

dm = cp.asarray(dm)
if dm.ndim == 3:
if dm.shape[0] > 2:
print("Warning: There are more than two density matrices to contract with one electron integrals, "
"it's not from an unrestricted calculation, and we're unsure about your purpose. "
"We sum the density matrices up, please check if that's expected.")
dm = cp.einsum("ijk->jk", dm)

assert dm.ndim == 2
Expand All @@ -357,7 +378,7 @@ def get_int3c1e_density_contracted(mol, grids, charge_exponents, dm, intopt):
nao_cart = intopt._sorted_mol.nao
ngrids = grids.shape[0]

dm = dm[np.ix_(intopt._ao_idx, intopt._ao_idx)] # intopt.ao_idx is in spherical basis
dm = intopt.sort_orbitals(dm, [0,1])
if not mol.cart:
cart2sph_transformation_matrix = intopt.cart2sph
# TODO: This part is inefficient (O(N^3)), should be changed to the O(N^2) algorithm
Expand Down
6 changes: 3 additions & 3 deletions gpu4pyscf/hessian/rks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -374,8 +374,8 @@ def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
vmat_dm[ia][:,:,mask] += contract('yjg,xjg->xyj', ao_mask[1:4], aow)
ao_dm0 = aow = None
t1 = log.timer_debug2('integration', *t1)
vmat_dm = opt.unsort_orbitals(vmat_dm, axis=[3])
for ia in range(_sorted_mol.natm):
vmat_dm[ia][:,:,opt._ao_idx] = vmat_dm[ia]
p0, p1 = aoslices[ia][2:]
vmat_dm[ia] += contract('xypq,pq->xyp', ipip[:,:,:,p0:p1], dm0[:,p0:p1])
elif xctype == 'GGA':
Expand Down Expand Up @@ -411,8 +411,8 @@ def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
vmat_dm[ia][:,:,mask] += vmat_dm_tmp
ao_dm0 = aow = None
t1 = log.timer_debug2('integration', *t1)
vmat_dm = opt.unsort_orbitals(vmat_dm, axis=[3])
for ia in range(_sorted_mol.natm):
vmat_dm[ia][:,:,opt._ao_idx] = vmat_dm[ia]
p0, p1 = aoslices[ia][2:]
vmat_dm[ia] += contract('xypq,pq->xyp', ipip[:,:,:,p0:p1], dm0[:,p0:p1])
vmat_dm[ia] += contract('yxqp,pq->xyp', ipip[:,:,p0:p1], dm0[:,p0:p1])
Expand Down Expand Up @@ -478,8 +478,8 @@ def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):

vmat_dm[ia][:,:,mask] += vmat_dm_tmp
t1 = log.timer_debug2('integration', *t1)
vmat_dm = opt.unsort_orbitals(vmat_dm, axis=[3])
for ia in range(_sorted_mol.natm):
vmat_dm[ia][:,:,opt._ao_idx] = vmat_dm[ia]
p0, p1 = aoslices[ia][2:]
vmat_dm[ia] += contract('xypq,pq->xyp', ipip[:,:,:,p0:p1], dm0[:,p0:p1])
vmat_dm[ia] += contract('yxqp,pq->xyp', ipip[:,:,p0:p1], dm0[:,p0:p1])
Expand Down
12 changes: 6 additions & 6 deletions gpu4pyscf/hessian/uks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -446,10 +446,10 @@ def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
vmatb_dm[ia][:,:,mask] += contract('yjg,xjg->xyj', ao_mask[1:4], aow)
ao_dm0a = ao_dm0b = aow = None
t1 = log.timer_debug2('integration', *t1)
vmata_dm = opt.unsort_orbitals(vmata_dm, axis=[3])
vmatb_dm = opt.unsort_orbitals(vmatb_dm, axis=[3])
for ia in range(_sorted_mol.natm):
p0, p1 = aoslices[ia][2:]
vmata_dm[ia][:,:,opt._ao_idx] = vmata_dm[ia]
vmatb_dm[ia][:,:,opt._ao_idx] = vmatb_dm[ia]
vmata_dm[ia] += contract('xypq,pq->xyp', ipipa[:,:,:,p0:p1], dm0a[:,p0:p1])
vmatb_dm[ia] += contract('xypq,pq->xyp', ipipb[:,:,:,p0:p1], dm0b[:,p0:p1])
elif xctype == 'GGA':
Expand Down Expand Up @@ -503,9 +503,9 @@ def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
vmatb_dm[ia][:,:,mask] += vmatb_dm_tmp
ao_dm0a = ao_dm0b = aow = None
t1 = log.timer_debug2('integration', *t1)
vmata_dm = opt.unsort_orbitals(vmata_dm, axis=[3])
vmatb_dm = opt.unsort_orbitals(vmatb_dm, axis=[3])
for ia in range(_sorted_mol.natm):
vmata_dm[ia][:,:,opt._ao_idx] = vmata_dm[ia]
vmatb_dm[ia][:,:,opt._ao_idx] = vmatb_dm[ia]
p0, p1 = aoslices[ia][2:]
vmata_dm[ia] += contract('xypq,pq->xyp', ipipa[:,:,:,p0:p1], dm0a[:,p0:p1])
vmata_dm[ia] += contract('yxqp,pq->xyp', ipipa[:,:,p0:p1], dm0a[:,p0:p1])
Expand Down Expand Up @@ -618,9 +618,9 @@ def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
vmata_dm[ia][:,:,mask] += vmata_dm_tmp
vmatb_dm[ia][:,:,mask] += vmatb_dm_tmp
t1 = log.timer_debug2('integration', *t1)
vmata_dm = opt.unsort_orbitals(vmata_dm, axis=[3])
vmatb_dm = opt.unsort_orbitals(vmatb_dm, axis=[3])
for ia in range(_sorted_mol.natm):
vmata_dm[ia][:,:,opt._ao_idx] = vmata_dm[ia]
vmatb_dm[ia][:,:,opt._ao_idx] = vmatb_dm[ia]
p0, p1 = aoslices[ia][2:]
vmata_dm[ia] += contract('xypq,pq->xyp', ipipa[:,:,:,p0:p1], dm0a[:,p0:p1])
vmata_dm[ia] += contract('yxqp,pq->xyp', ipipa[:,:,p0:p1], dm0a[:,p0:p1])
Expand Down

0 comments on commit b0ff950

Please sign in to comment.