From c1a678f3d5f01e0bf5a4e41a31236e7c275eaf7b Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Fri, 18 Sep 2026 11:51:50 +0800 Subject: [PATCH 1/9] First round of GDF Hcore gradient, supporting all electron only (no pp), and disabling many other cases --- gpu4pyscf/pbc/df/grad/krhf.py | 235 +++++++++++++++++++++++++++++++++- gpu4pyscf/pbc/grad/krhf.py | 24 ++-- gpu4pyscf/pbc/grad/rhf.py | 28 ++-- 3 files changed, 254 insertions(+), 33 deletions(-) diff --git a/gpu4pyscf/pbc/df/grad/krhf.py b/gpu4pyscf/pbc/df/grad/krhf.py index 137871c1f..72e8e8c09 100644 --- a/gpu4pyscf/pbc/df/grad/krhf.py +++ b/gpu4pyscf/pbc/df/grad/krhf.py @@ -23,11 +23,12 @@ from pyscf import lib from pyscf.pbc.tools.k2gamma import double_translation_indices from pyscf.pbc.lib.kpts_helper import is_zero +from pyscf.pbc.df import aft as aft_cpu from gpu4pyscf.lib import logger from gpu4pyscf.lib.cupy_helper import ( contract, asarray, ndarray, transpose_sum, get_avail_mem, empty_aligned) from gpu4pyscf.__config__ import props as gpu_specs -from gpu4pyscf.pbc.df.int3c2e import libpbc, POOL_SIZE, MAX_IMGS_PER_TASK +from gpu4pyscf.pbc.df.int3c2e import libpbc, POOL_SIZE, MAX_IMGS_PER_TASK, SRInt3c2eOpt from gpu4pyscf.pbc.df.rsdf_builder import LINEAR_DEP_THR from gpu4pyscf.pbc.tools.pbc import madelung, _Gv_wrap_around from gpu4pyscf.pbc.df import ft_ao, aft_jk @@ -41,6 +42,8 @@ from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.gto.cell import get_Gv_weights from gpu4pyscf.pbc.lib.kpts_helper import fft_matrix, kk_adapted_iter +from gpu4pyscf.pbc.tools.k2gamma import kpts_to_kmesh +from gpu4pyscf.pbc.df.rsdf_builder import estimate_ke_cutoff_for_omega, estimate_omega_for_ke_cutoff, _weighted_coulG_LR def _get_ejk_derivatives(int3c2e_opt, dm, kpts=None, hermi=0, j_factor=1., k_factor=1., @@ -830,3 +833,233 @@ def _j_energy_per_atom(int3c2e_opt, dm, kpts=None, hermi=0, omega=None, '''Compatibility wrapper returning only the atomic Coulomb derivatives.''' return _get_ej_derivatives( int3c2e_opt, dm, kpts, hermi, omega, verbose, linear_dep_threshold)[:-3] + +def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbose=None): + log = logger.new_logger(cell, verbose) + t0 = log.init_timer() + + is_single_kpt = kpts is not None and kpts.ndim == 1 + is_gamma_point = kpts is None or is_zero(kpts) + if is_single_kpt: + kpts = kpts.reshape(1, 3) + if is_gamma_point: + bvk_kmesh = np.ones(3, dtype=int) + else: + bvk_kmesh = kpts_to_kmesh(cell, kpts, bound_by_supmol=True) + + # Guess range-separation parameter based on system size + omega = 0.4 + ke_cutoff = estimate_ke_cutoff_for_omega(cell, omega) + mesh = cell.cutoff_to_mesh(ke_cutoff) + nGv = np.prod(mesh) + ke_cutoff *= (3e3/nGv)**(2./3) + omega = estimate_omega_for_ke_cutoff(cell, ke_cutoff) + ke_cutoff = estimate_ke_cutoff_for_omega(cell, omega) + mesh = cell.cutoff_to_mesh(ke_cutoff) + mesh = cell.symmetrize_mesh(mesh) + log.debug('get_pp_loc_part1_grad: omega = %g Ecut = %s mesh = %s', + omega, ke_cutoff, mesh) + + fakenuc = aft_cpu._fake_nuc(cell, with_pseudo=with_pseudo) + int3c2e_opt = SRInt3c2eOpt(cell, fakenuc, omega=-omega, bvk_kmesh=bvk_kmesh).build() + charges = -cp.asarray(cell.atom_charges(), dtype=np.float64) + nuc = int3c2e_opt.contract_auxvec(charges, kpts) + + cell = int3c2e_opt.cell + auxcell = int3c2e_opt.auxcell + + dm = cp.asarray(dm) + dm = cell.apply_C_mat_CT(dm) + if hermi != 1: + dm = transpose_sum(dm, inplace=True) + dm[:] *= .5 + + if kpts is None or is_zero(kpts): + dm = cp.asarray(dm.real, order='C') + nkpts = 1 + else: + assert len(int3c2e_opt.bvkmesh_Ls) == len(kpts) + nkpts = len(kpts) + #:expLk = cp.exp(1j*asarray(int3c2e_opt.bvkmesh_Ls).dot(asarray(kpts).T)) + expLk = fft_matrix(int3c2e_opt.bvk_kmesh) + dm = contract('Lk,kpq->Lpq', expLk, dm) + dm = cp.asarray(dm.real, order='C') + dm *= 1./nkpts + + Gv, _, kws = cell.get_Gv_weights(mesh) + ngrids = len(Gv) + if with_pseudo: + raise NotImplementedError("") + # #TODO: call multigrid.eval_vpplocG after removing its part2 contribution + # ZG = ft_ao.ft_ao(fakenuc, Gv).conj() + # ZG = ZG.dot(charges) + # ZG *= _weighted_coulG_LR(cell, Gv, omega, kws) + # if ((cell.dimension == 3 or + # (cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum'))): + # exps = cp.asarray(np.hstack(fakenuc.bas_exps())) + # ZG[0] -= charges.dot(np.pi/exps) / cell.vol + else: + # ZG = _get_ZSI(cell, mesh).conj() + # ZG *= _weighted_coulG_LR(cell, Gv, omega, kws) + + wcoulG_LR = _weighted_coulG_LR(cell, Gv, omega, kws) + ft_opt = ft_ao.FTOpt.from_intopt(int3c2e_opt) + + bvk_ncells = len(int3c2e_opt.bvkmesh_Ls) + aux_loc = auxcell.ao_loc + nao = dm.shape[-1] + naux = int(aux_loc[-1]) + + eval_ft = ft_opt.ft_evaluator( + compressing=True, cart=True, original_ao_order=False)[0] + pair_addresses, diag_idx = ft_opt.pair_and_diag_indices( + cart=True, original_ao_order=False) + i_addr, j_addr = divmod(pair_addresses, bvk_ncells * nao) + dm_tril = dm.reshape(bvk_ncells*nao, nao).real[j_addr, i_addr] + dm_tril[diag_idx] *= .5 + dm_tril *= 2 + + mem_avail = get_avail_mem(exclude_memory_pool=True) + nao_pair = len(dm_tril) + Gblksize = int(mem_avail*.8//((nao_pair+naux*2)*16))//32*32 + Gblksize = min(Gblksize, ngrids) + assert Gblksize > 0 + log.debug1('%.3f GB free memory. blksize=%d for LR part', + mem_avail*1e-9, Gblksize) + + rhoG = cp.empty(ngrids, dtype=np.complex128) + buf = cp.empty(max(nao_pair,naux)*Gblksize, dtype=np.complex128) + for p0, p1 in lib.prange(0, ngrids, Gblksize): + nGv = p1 - p0 + # conj((r|G)^{[0]}) (ij|G)^{[0]} + pqG = eval_ft(Gv[p0:p1], out=buf) + rhoGz = cp.einsum('pG,p->G', pqG.view(np.float64), dm_tril) + rhoG[p0:p1] = rhoGz.view(np.complex128) + + aft_envs = ft_opt.aft_envs + shm_size = aft_jk._estimate_max_shm_size(cell, (1, 0)) + mem_avail = get_avail_mem(exclude_memory_pool=True) + Gblksize = int(mem_avail*.8//(naux*2*16))//32*32 + Gblksize = min(Gblksize, ngrids) + rho_nucG = cp.empty(ngrids, dtype=np.complex128) + buf = cp.empty(naux*Gblksize, dtype=np.complex128) + for p0, p1 in lib.prange(0, ngrids, Gblksize): + auxG = ft_ao.ft_ao(auxcell, Gv[p0:p1], out=buf).T + rho_nucG[p0:p1] = charges.dot( + auxG.view(np.float64)).view(np.complex128) + + vG = rhoG * wcoulG_LR + GvT = cp.asarray(Gv.T.ravel()) + ej_sigma_aux = cp.zeros([cell.natm+3, 3]) + aux_ft_envs = RysIntEnvVars.new( + auxcell.natm, auxcell.nbas, auxcell._atm, auxcell._bas, + _scale_sp_ctr_coeff(auxcell), auxcell.ao_loc) + err = libpbc.PBC_ft_ao_deriv( + ctypes.cast(ej_sigma_aux[:-3].data.ptr, ctypes.c_void_p), + ctypes.cast(ej_sigma_aux[-3:].data.ptr, ctypes.c_void_p), + ctypes.cast(charges.data.ptr, ctypes.c_void_p), + ctypes.cast(vG.data.ptr, ctypes.c_void_p), + ctypes.cast(GvT.data.ptr, ctypes.c_void_p), + ctypes.byref(aux_ft_envs), ctypes.c_int(ngrids)) + if err != 0: + raise RuntimeError('ft_ao_deriv failed') + + ej_sigma_lr = cp.zeros([cell.natm+3, 3]) + vG_conj = rho_nucG.conj() * wcoulG_LR + bas_ij_idx, bas_ij_img_idx, shl_pair_offsets = aft_jk._generate_shl_pairs(ft_opt) + nbatches_shl_pair = len(shl_pair_offsets) - 1 + err = libpbc.PBC_ft_aopair_ej_deriv( + ctypes.cast(ej_sigma_lr[:-3].data.ptr, ctypes.c_void_p), + ctypes.cast(ej_sigma_lr[-3:].data.ptr, ctypes.c_void_p), + ctypes.cast(dm.data.ptr, ctypes.c_void_p), + ctypes.cast(vG_conj.data.ptr, ctypes.c_void_p), + ctypes.cast(GvT.data.ptr, ctypes.c_void_p), + ctypes.byref(aft_envs), + ctypes.c_int(nbatches_shl_pair), + ctypes.c_int(ngrids), + ctypes.c_int(shm_size), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(bas_ij_img_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(shl_pair_offsets.data.ptr, ctypes.c_void_p), + ctypes.c_int(ft_opt.permutation_symmetry)) + if err != 0: + raise RuntimeError('PBC_ft_aopair_ej_deriv failed') + + ej_sigma_lr *= 2 + ej_sigma_lr += ej_sigma_aux + + ej_sigma = ej_sigma_lr + t0 = log.timer_debug1('lr_int3c2e_deriv via aft', *t0) + ft_opt = None + + nsp_per_block, gout_stride, shm_size = int3c2e_scheme( + gout_width=54, deriv=(1,0,0)) + lmax = cell.uniq_l_ctr[:,0].max() + laux = auxcell.uniq_l_ctr[:,0].max() + shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + + l_ctr_aux_offsets = np.append(0, np.cumsum(auxcell.l_ctr_counts)) + l_ctr_aux_offsets, uniq_l_ctr_aux = _split_l_ctr_pattern( + l_ctr_aux_offsets, auxcell.uniq_l_ctr, POOL_SIZE) + ksh_offsets_cpu = l_ctr_aux_offsets + ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu, dtype=np.int32) + + nksh_per_batch = ksh_offsets_cpu[1:] - ksh_offsets_cpu[:-1] + shl_pair_batch_size = rhf._get_shl_pair_batch_size( + nksh_per_batch, bvk_ncells) + bas_ij_idx, shl_pair_offsets = cell.aggregate_shl_pairs( + int3c2e_opt.bas_ij_cache, nsp_per_block=shl_pair_batch_size) + + diffuse_exps = cp.asarray(int3c2e_opt.diffuse_exps) + diffuse_coefs = cp.asarray(int3c2e_opt.diffuse_coefs) + log_cutoff = math.log(int3c2e_opt.cutoff) + + ej_sigma_sr = cp.zeros([cell.natm+3, 3]) + workers = gpu_specs['multiProcessorCount'] + pool = cp.empty(workers * POOL_SIZE*(MAX_IMGS_PER_TASK+2) + 1, dtype=np.uint32) + head = pool[-1:] + task_pool = empty_aligned((workers, POOL_SIZE*16), np.int32, alignment=128) + int3c2e_envs = int3c2e_opt.int3c2e_envs + kern = libpbc.PBCsr_ejk_int3c2e_deriv + err = kern( + ctypes.cast(ej_sigma_sr[:-3].data.ptr, ctypes.c_void_p), + ctypes.cast(ej_sigma_sr[-3:].data.ptr, ctypes.c_void_p), + ctypes.cast(dm.data.ptr, ctypes.c_void_p), + ctypes.cast(charges.data.ptr, ctypes.c_void_p), + ctypes.c_double(-int3c2e_opt.omega), + ctypes.byref(int3c2e_envs), + ctypes.cast(pool.data.ptr, ctypes.c_void_p), + ctypes.cast(task_pool.data.ptr, ctypes.c_void_p), + ctypes.cast(head.data.ptr, ctypes.c_void_p), + ctypes.c_int(shm_size_max), + ctypes.c_int(len(shl_pair_offsets) - 1), + ctypes.c_int(len(ksh_offsets_gpu) - 1), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(shl_pair_offsets.data.ptr, ctypes.c_void_p), + ctypes.cast(ksh_offsets_gpu.data.ptr, ctypes.c_void_p), + ctypes.cast(int3c2e_opt.img_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(int3c2e_opt.img_offsets.data.ptr, ctypes.c_void_p), + ctypes.cast(gout_stride.data.ptr, ctypes.c_void_p), + lib.c_null_ptr(), + ctypes.c_int(0), + ctypes.c_int(auxcell.nbas), + ctypes.c_int(naux), + ctypes.cast(diffuse_exps.data.ptr, ctypes.c_void_p), + ctypes.cast(diffuse_coefs.data.ptr, ctypes.c_void_p), + ctypes.c_float(log_cutoff)) + if err != 0: + raise RuntimeError('PBCsr_ejk_int3c2e_deriv failed') + ej_sigma += ej_sigma_sr * 2 + + t0 = log.timer_debug1('contract int3c2e_ejk_deriv', *t0) + return ej_sigma.get() + +def get_nuc(cell, dm, kpts=None, hermi=1): + log = logger.new_logger(cell) + t0 = log.init_timer() + nuc = get_pp_loc_part1_grad(cell, dm, kpts, hermi, with_pseudo=False, verbose=log) + log.timer('get_nuc gradient', *t0) + return nuc + +def get_pp(cell, dm, kpts=None, hermi=1): + raise NotImplementedError("") diff --git a/gpu4pyscf/pbc/grad/krhf.py b/gpu4pyscf/pbc/grad/krhf.py index ad3c61a71..69abef14c 100644 --- a/gpu4pyscf/pbc/grad/krhf.py +++ b/gpu4pyscf/pbc/grad/krhf.py @@ -28,7 +28,7 @@ from gpu4pyscf.pbc.dft import numint as pbc_numint from gpu4pyscf.pbc.dft.numint import eval_ao_kpts, _GTOvalOpt from gpu4pyscf.pbc.dft import UniformGrids, BeckeGrids -from gpu4pyscf.pbc.dft import multigrid_v3 +from gpu4pyscf.pbc.dft import multigrid, multigrid_v3 from gpu4pyscf.pbc.df import ft_ao, GDF from gpu4pyscf.pbc.df.aft import get_SI, _get_ZSI from gpu4pyscf.pbc.gto import int1e @@ -38,6 +38,8 @@ vppnl_nuc_grad, _get_pp_nonloc_strain_derivatives) from gpu4pyscf.pbc.grad.rhf import contract_h1e_dm, _get_ejk_derivatives from gpu4pyscf.pbc.grad import rhf as pbchf_grad +from gpu4pyscf.pbc.scf import hf as pbchf +from gpu4pyscf.pbc.df.grad.krhf import get_nuc __all__ = ['Gradients'] @@ -354,21 +356,13 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None): if isinstance(ni, multigrid_v3.MultiGridNumInt): # Vne or pploc contribution is evaluated in energy_ee grad_sigma += int1e.kin_derivatives(cell, dm0, kpts) + elif isinstance(ni, multigrid.MultiGridNumIntBase): + raise NotImplementedError("") + elif np.prod(cell.mesh) < pbchf.ALLOWED_FFT_MESH_SIZE: + raise NotImplementedError("") else: - hcore_deriv = self.hcore_generator(cell, kpts) - dh1e = cp.empty([cell.natm, 3]) - for ia in range(cell.natm): - h1ao = hcore_deriv(ia) - dh1e[ia] = cp.einsum('kxij,kji->x', h1ao, dm0).real - grad_sigma[:-3] += dh1e.get() / nkpts - if isinstance(self.grids or getattr(mf, 'grids', None), BeckeGrids): - grad_sigma[-3:] = np.nan - else: - # hcore_generator includes kinetic gradients, but not kinetic strain. - grad_sigma[-3:] += int1e.kin_derivatives(cell, dm0, kpts)[-3:] - ni = multigrid_v3.MultiGridNumInt(cell) - grad_sigma[-3:] += ni.energy_strain_gradient( - 'HF', dm0, kpts, spin=0, with_j=False, with_nuc=True) + grad_sigma += get_nuc(cell, dm0, kpts) + grad_sigma += int1e.kin_derivatives(cell, dm0, kpts) if cell._pseudo: grad_sigma[:-3] += vppnl_nuc_grad(cell, dm0, kpts=kpts) / nkpts diff --git a/gpu4pyscf/pbc/grad/rhf.py b/gpu4pyscf/pbc/grad/rhf.py index a2424ab5f..3b97b5fed 100644 --- a/gpu4pyscf/pbc/grad/rhf.py +++ b/gpu4pyscf/pbc/grad/rhf.py @@ -25,13 +25,15 @@ from gpu4pyscf.lib import logger import gpu4pyscf.grad.rhf as mol_rhf from gpu4pyscf.pbc.tools.k2gamma import kpts_to_kmesh -from gpu4pyscf.pbc.dft import multigrid_v3 +from gpu4pyscf.pbc.dft import multigrid, multigrid_v3 from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt from gpu4pyscf.pbc.df import aft_jk, AFTDF, GDF from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.dft import KohnShamDFT, BeckeGrids from gpu4pyscf.pbc.grad.pp import ( vppnl_nuc_grad, _get_pp_nonloc_strain_derivatives) +from gpu4pyscf.pbc.scf import hf as pbchf +from gpu4pyscf.pbc.df.grad.krhf import get_nuc from gpu4pyscf.gto.mole import groupby __all__ = ['Gradients'] @@ -184,22 +186,13 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None): if isinstance(ni, multigrid_v3.MultiGridNumInt): # Vne or pploc contribution is evaluated in energy_ee grad_sigma += int1e.kin_derivatives(cell, dm0) + elif isinstance(ni, multigrid.MultiGridNumIntBase): + raise NotImplementedError("") + elif np.prod(cell.mesh) < pbchf.ALLOWED_FFT_MESH_SIZE: + raise NotImplementedError("") else: - from gpu4pyscf.pbc.grad.krhf import hcore_generator - hcore_deriv = hcore_generator(self, cell, np.zeros((1, 3))) - dh1e = cp.empty([cell.natm, 3]) - for ia in range(cell.natm): - h1ao = hcore_deriv(ia) - dh1e[ia] = cp.einsum('xij,ji->x', h1ao[0], dm0).real - grad_sigma[:-3] += dh1e.get() - if isinstance(self.grids or getattr(mf, 'grids', None), BeckeGrids): - grad_sigma[-3:] = np.nan - else: - # hcore_generator includes kinetic gradients, but not kinetic strain. - grad_sigma[-3:] += int1e.kin_derivatives(cell, dm0)[-3:] - ni = multigrid_v3.MultiGridNumInt(cell) - grad_sigma[-3:] += ni.energy_strain_gradient( - 'HF', dm0, spin=0, with_j=False, with_nuc=True) + grad_sigma += get_nuc(cell, dm0) + grad_sigma += int1e.kin_derivatives(cell, dm0) if cell._pseudo: grad_sigma[:-3] += vppnl_nuc_grad(cell, dm0) @@ -375,8 +368,9 @@ def strain_tensor_dispalcement(x, y, disp): return E_strain def _finite_diff_cells(cell, x, y, disp=1e-4, precision=None): + cell = cell.copy() + cell.verbose = 0 if precision is not None: - cell = cell.copy() cell.precision = precision a = cell.lattice_vectors() r = cell.atom_coords() From 307b3e6b8af40635f3778c8d9b4db2e14656907b Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Fri, 18 Sep 2026 17:08:06 +0800 Subject: [PATCH 2/9] Support stress tensor, add tests for all restricted cases (unrestricted not supported yet) --- gpu4pyscf/pbc/df/grad/krhf.py | 21 +- .../df/tests/test_pbc_hcore_derivatives.py | 354 ++++++++++++++++++ gpu4pyscf/pbc/grad/krhf.py | 86 ----- 3 files changed, 368 insertions(+), 93 deletions(-) create mode 100644 gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py diff --git a/gpu4pyscf/pbc/df/grad/krhf.py b/gpu4pyscf/pbc/df/grad/krhf.py index 72e8e8c09..e978f4e18 100644 --- a/gpu4pyscf/pbc/df/grad/krhf.py +++ b/gpu4pyscf/pbc/df/grad/krhf.py @@ -43,7 +43,7 @@ from gpu4pyscf.pbc.gto.cell import get_Gv_weights from gpu4pyscf.pbc.lib.kpts_helper import fft_matrix, kk_adapted_iter from gpu4pyscf.pbc.tools.k2gamma import kpts_to_kmesh -from gpu4pyscf.pbc.df.rsdf_builder import estimate_ke_cutoff_for_omega, estimate_omega_for_ke_cutoff, _weighted_coulG_LR +from gpu4pyscf.pbc.df.rsdf_builder import estimate_ke_cutoff_for_omega, estimate_omega_for_ke_cutoff def _get_ejk_derivatives(int3c2e_opt, dm, kpts=None, hermi=0, j_factor=1., k_factor=1., @@ -886,8 +886,15 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos dm = cp.asarray(dm.real, order='C') dm *= 1./nkpts - Gv, _, kws = cell.get_Gv_weights(mesh) + Gv, _, kws = get_Gv_weights(cell, mesh) ngrids = len(Gv) + wcoulG_LR0, wcoulG_LR1 = get_wcoulG( + cell, Gv, int3c2e_opt.omega) + wcoulG_SR_at_G0 = np.pi / int3c2e_opt.omega**2 * kws + wcoulG_LR0[0] -= wcoulG_SR_at_G0 + wcoulG_LR1[:,:,0] += wcoulG_SR_at_G0 * cp.eye(3) + ft_opt = ft_ao.FTOpt.from_intopt(int3c2e_opt) + if with_pseudo: raise NotImplementedError("") # #TODO: call multigrid.eval_vpplocG after removing its part2 contribution @@ -899,12 +906,10 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos # exps = cp.asarray(np.hstack(fakenuc.bas_exps())) # ZG[0] -= charges.dot(np.pi/exps) / cell.vol else: + pass # ZG = _get_ZSI(cell, mesh).conj() # ZG *= _weighted_coulG_LR(cell, Gv, omega, kws) - wcoulG_LR = _weighted_coulG_LR(cell, Gv, omega, kws) - ft_opt = ft_ao.FTOpt.from_intopt(int3c2e_opt) - bvk_ncells = len(int3c2e_opt.bvkmesh_Ls) aux_loc = auxcell.ao_loc nao = dm.shape[-1] @@ -948,7 +953,7 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos rho_nucG[p0:p1] = charges.dot( auxG.view(np.float64)).view(np.complex128) - vG = rhoG * wcoulG_LR + vG = rhoG * wcoulG_LR0 GvT = cp.asarray(Gv.T.ravel()) ej_sigma_aux = cp.zeros([cell.natm+3, 3]) aux_ft_envs = RysIntEnvVars.new( @@ -965,7 +970,7 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos raise RuntimeError('ft_ao_deriv failed') ej_sigma_lr = cp.zeros([cell.natm+3, 3]) - vG_conj = rho_nucG.conj() * wcoulG_LR + vG_conj = rho_nucG.conj() * wcoulG_LR0 bas_ij_idx, bas_ij_img_idx, shl_pair_offsets = aft_jk._generate_shl_pairs(ft_opt) nbatches_shl_pair = len(shl_pair_offsets) - 1 err = libpbc.PBC_ft_aopair_ej_deriv( @@ -987,6 +992,8 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos ej_sigma_lr *= 2 ej_sigma_lr += ej_sigma_aux + ej_sigma_lr[-3:] += cp.einsum( + 'g,g,xyg->xy', rho_nucG, rhoG.conj(), wcoulG_LR1).real ej_sigma = ej_sigma_lr t0 = log.timer_debug1('lr_int3c2e_deriv via aft', *t0) diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py new file mode 100644 index 000000000..ca8ddb155 --- /dev/null +++ b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py @@ -0,0 +1,354 @@ +# Copyright 2024-2025 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. + +import unittest +import numpy as np +import cupy as cp + +import pyscf +from gpu4pyscf.pbc.dft.rks import RKS +from gpu4pyscf.pbc.dft.krks import KRKS +from gpu4pyscf.pbc.scf.hf import RHF +from gpu4pyscf.pbc.scf.khf import KRHF +from gpu4pyscf.pbc.dft.gen_grid import BeckeGrids +from gpu4pyscf.pbc.df.grad.krhf import get_nuc +from gpu4pyscf.pbc.gto.int1e import kin_derivatives +from gpu4pyscf.lib.multi_gpu import num_devices + +def numerical_hcore_gradient_and_stresstensor(cell, dm, kmesh): + if dm.ndim == 2: + dm = dm[None,:,:] + + def get_energy(cell): + mf = KRHF(cell).density_fit(auxbasis="def2-universal-jkfit") + kpts = cell.make_kpts(kmesh) + hcore = mf.get_hcore(kpts = kpts) + e = cp.einsum("kij,kji->", hcore, dm) + e = e.real / kpts.shape[0] + return e + + numerical_gradient = np.zeros((cell.natm + 3, 3)) + dx = 1e-4 + cell_copy = cell.copy() + for i_atom in range(cell.natm): + for i_xyz in range(3): + xyz_p = cell.atom_coords() + xyz_p[i_atom, i_xyz] += dx + cell_copy.set_geom_(xyz_p, unit='Bohr') + cell_copy.build() + E_p = get_energy(cell_copy) + + xyz_m = cell.atom_coords() + xyz_m[i_atom, i_xyz] -= dx + cell_copy.set_geom_(xyz_m, unit='Bohr') + cell_copy.build() + E_m = get_energy(cell_copy) + + numerical_gradient[i_atom, i_xyz] = (E_p - E_m) / (2 * dx) + + translation_invariance = np.sum(numerical_gradient, axis=0) + assert np.max(np.abs(translation_invariance)) < 1e-8, "Bad numerical gradient" + + from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells + for i_xyz in range(3): + for j_xyz in range(3): + cell_p, cell_m = _finite_diff_cells(cell, i_xyz, j_xyz, disp = dx) + E_p = get_energy(cell_p) + E_m = get_energy(cell_m) + + numerical_gradient[cell.natm + i_xyz, j_xyz] = (E_p - E_m) / (2 * dx) + + # np.set_printoptions(precision=16, suppress=True, linewidth=np.inf) + # print(repr(numerical_gradient)) + return numerical_gradient + +def numerical_gradient_and_stresstensor(cell, get_energy): + assert callable(get_energy) + + numerical_gradient = np.zeros((cell.natm + 3, 3)) + dx = 1e-4 + cell_copy = cell.copy() + for i_atom in range(cell.natm): + for i_xyz in range(3): + xyz_p = cell.atom_coords() + xyz_p[i_atom, i_xyz] += dx + cell_copy.set_geom_(xyz_p, unit='Bohr') + cell_copy.build() + E_p = get_energy(cell_copy) + if isinstance(E_p, tuple): + E_p = E_p[0] + + xyz_m = cell.atom_coords() + xyz_m[i_atom, i_xyz] -= dx + cell_copy.set_geom_(xyz_m, unit='Bohr') + cell_copy.build() + E_m = get_energy(cell_copy) + if isinstance(E_m, tuple): + E_m = E_m[0] + + numerical_gradient[i_atom, i_xyz] = (E_p - E_m) / (2 * dx) + + translation_invariance = np.sum(numerical_gradient, axis=0) + assert np.max(np.abs(translation_invariance)) < 1e-7, "Bad numerical gradient" + + from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells + for i_xyz in range(3): + for j_xyz in range(3): + cell_p, cell_m = _finite_diff_cells(cell, i_xyz, j_xyz, disp = dx) + E_p = get_energy(cell_p) + E_m = get_energy(cell_m) + if isinstance(E_p, tuple): + E_p = E_p[0] + if isinstance(E_m, tuple): + E_m = E_m[0] + + numerical_gradient[cell.natm + i_xyz, j_xyz] = (E_p - E_m) / (2 * dx) + + np.set_printoptions(precision=16, suppress=True, linewidth=np.inf) + print(f"ref_gradient = np.{repr(numerical_gradient)}") + return numerical_gradient + +class KnownValues(unittest.TestCase): + def test_gdf_hcore_derivatives_rhf(self): + cell = pyscf.M( + atom = """ + O 15.43509000 9.59549000 8.94968000 + H 15.05724000 9.21878000 9.73314000 + H 0.51550474 9.33856000 9.01857000 + He 2.51550474 9.33856000 9.01857000 + """, + a = np.eye(3) * (15.9069652593 / 3), + unit = "Angstrom", + basis = "def2-SVP", + verbose = 0, + ) + + def get_energy(cell): + mf = RHF(cell).density_fit(auxbasis="def2-universal-jkfit") + mf.conv_tol = 1e-12 + e = mf.kernel() + assert mf.converged + return e, mf + test_energy, mf = get_energy(cell) + + gobj = mf.Gradients() + test_gradient = gobj.kernel() + test_derivatives = np.vstack((test_gradient, gobj.stress)) + + # Energy check is consistency check + ref_energy = -78.84270992380236 + # ref_derivatives = numerical_gradient_and_stresstensor(cell, get_energy) + ref_derivatives = np.array([ + [-0.063013447402227 , 0.0161632252826394, -0.0025874167874917], + [-0.000704673723817 , 0.001038964825284 , -0.0022375882480219], + [ 0.0599051951155616, -0.0166196268480689, 0.0044088567108247], + [ 0.0038129256552111, -0.0005825626203659, 0.0004161493194488], + [ 0.1479508021162701, -0.0302303583055163, 0.0074059591526066], + [-0.0302303594423847, 0.0338498824703493, -0.0000233173835795], + [ 0.0074059608579091, -0.0000233153940599, 0.0259459928031447], + ]) + + assert np.abs(test_energy - ref_energy) < 1e-9 + assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 5e-8 + # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 5e-7 # TODO: Something actually wrong here + + dm = mf.make_rdm1() + kmesh = np.array([1,1,1]) + kpts = cell.make_kpts(kmesh) + + test_hcore_derivatives = get_nuc(cell, dm, kpts) + test_hcore_derivatives += kin_derivatives(cell, dm, kpts) + assert test_hcore_derivatives.shape == (cell.natm + 3, 3) + + ref_hcore_derivatives = numerical_hcore_gradient_and_stresstensor(cell, dm, kmesh) + + assert np.max(np.abs(test_hcore_derivatives[:-3, :] - ref_hcore_derivatives[:-3, :])) < 3e-8 + assert np.max(np.abs(test_hcore_derivatives[-3:, :] - ref_hcore_derivatives[-3:, :])) < 3e-7 + + def test_gdf_hcore_derivatives_rks_large_cell(self): + cell = pyscf.M( + atom = """ + O 15.43509000 9.59549000 8.94968000 + H 15.05724000 9.21878000 9.73314000 + H 0.51550474 9.33856000 9.01857000 + He 2.51550474 9.33856000 9.01857000 + """, + a = np.eye(3) * 15.9069652593, + unit = "Angstrom", + basis = "6-31g", + verbose = 0, + ) + + def get_energy(cell): + mf = RKS(cell, xc="r2scan").density_fit(auxbasis="def2-universal-jkfit") + mf.grids = BeckeGrids(cell) + mf.grids.atom_grid = (50, 194) + mf.conv_tol = 1e-12 + e = mf.kernel() + assert mf.converged + return e, mf + test_energy, mf = get_energy(cell) + + gobj = mf.Gradients() + gobj.grid_response = True + test_gradient = gobj.kernel() + test_derivatives = np.vstack((test_gradient, gobj.stress)) + + # Energy check is consistency check + ref_energy = -79.23783762162086 + # ref_derivatives = numerical_gradient_and_stresstensor(cell, get_energy) + ref_derivatives = np.array([ + [-0.052554448615183 , 0.0019232464154584, 0.0164716913531038], + [ 0.0195730997631927, 0.0102995896611446, -0.0248131673430407], + [ 0.0348620867640648, -0.0122673990432531, 0.0083400308170667], + [-0.0018807354251749, 0.0000445629666501, 0.0000014450307617], + [ 0.0405035486750194, -0.0299275689030765, 0.0332470406760876], + [-0.0299700908357181, -0.001340544457662 , 0.0136869616795821], + [ 0.0332548719939041, 0.0136434722008971, -0.0355085206393824], + ]) + + assert np.abs(test_energy - ref_energy) < 1e-9 + assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 3e-8 + # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 3e-7 # TODO: Support Becke grid stress tensor + + dm = mf.make_rdm1() + kmesh = np.array([1,1,1]) + kpts = cell.make_kpts(kmesh) + + test_hcore_derivatives = get_nuc(cell, dm, kpts) + test_hcore_derivatives += kin_derivatives(cell, dm, kpts) + assert test_hcore_derivatives.shape == (cell.natm + 3, 3) + + ref_hcore_derivatives = numerical_hcore_gradient_and_stresstensor(cell, dm, kmesh) + + assert np.max(np.abs(test_hcore_derivatives[:-3, :] - ref_hcore_derivatives[:-3, :])) < 3e-8 + assert np.max(np.abs(test_hcore_derivatives[-3:, :] - ref_hcore_derivatives[-3:, :])) < 3e-7 + + def test_gdf_hcore_derivatives_krhf(self): + cell = pyscf.M( + atom = """ + O 15.43509000 9.59549000 8.94968000 + H 15.05724000 9.21878000 9.73314000 + H 0.51550474 9.33856000 9.01857000 + """, + a = np.eye(3) * (15.9069652593 / 3), + unit = "Angstrom", + basis = "sto-6g", + verbose = 0, + ) + + kmesh = np.array([3,1,1]) + + def get_energy(cell): + kpts = cell.make_kpts(kmesh) + mf = KRHF(cell, kpts=kpts).density_fit(auxbasis="def2-universal-jkfit") + mf.conv_tol = 1e-11 + e = mf.kernel() + assert mf.converged + return e, mf + test_energy, mf = get_energy(cell) + + gobj = mf.Gradients() + test_gradient = gobj.kernel() + test_derivatives = np.vstack((test_gradient, gobj.stress)) + + # Energy check is consistency check + ref_energy = -75.68294188844864 + # ref_derivatives = numerical_gradient_and_stresstensor(cell, get_energy) + ref_derivatives = np.array([ + [-0.0590151794455096, -0.0182923543690094, 0.0523572944644002], + [ 0.0116237738012614, 0.0251060370715095, -0.0479677144227253], + [ 0.047391405857411 , -0.0068136822761744, -0.004389579544295 ], + [ 0.0891122930113397, -0.0303158925873959, 0.0254392091392219], + [-0.0303158941505899, -0.0126309679870928, 0.0368723210897315], + [ 0.0254392105603074, 0.0368723237187396, -0.068196215678995 ], + ]) + + assert np.abs(test_energy - ref_energy) < 1e-9 + assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 5e-7 + # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 5e-6 # TODO: Something actually wrong here + + dm = mf.make_rdm1() + kpts = cell.make_kpts(kmesh) + + test_hcore_derivatives = get_nuc(cell, dm, kpts) + test_hcore_derivatives += kin_derivatives(cell, dm, kpts) + assert test_hcore_derivatives.shape == (cell.natm + 3, 3) + + ref_hcore_derivatives = numerical_hcore_gradient_and_stresstensor(cell, dm, kmesh) + + assert np.max(np.abs(test_hcore_derivatives[:-3, :] - ref_hcore_derivatives[:-3, :])) < 3e-8 + assert np.max(np.abs(test_hcore_derivatives[-3:, :] - ref_hcore_derivatives[-3:, :])) < 3e-7 + + @unittest.skipIf(num_devices > 1, '') + def test_gdf_hcore_derivatives_krks(self): + cell = pyscf.M( + a = '''0. 1.7834 1.7834 + 1.7834 0. 1.7834 + 1.7834 1.7834 0. ''', + atom = 'C 0., 0., 0.; C 0.8917, 0.8917, 0.8917', + basis = 'def2-svp', + mesh = [10086] * 3, + verbose = 4, + ) + + kmesh = np.array([3,1,1]) + + def get_energy(cell): + kpts = cell.make_kpts(kmesh) + mf = KRKS(cell, xc="wB97X", kpts=kpts).density_fit(auxbasis="def2-universal-jkfit") + mf.grids = BeckeGrids(cell) + mf.grids.atom_grid = (50, 194) + mf.conv_tol = 1e-12 + e = mf.kernel() + assert mf.converged + return e, mf + test_energy, mf = get_energy(cell) + + gobj = mf.Gradients() + gobj.grid_response = True + test_gradient = gobj.kernel() + test_derivatives = np.vstack((test_gradient, gobj.stress)) + + # Energy check is consistency check + ref_energy = -75.65250113719456 + # ref_derivatives = numerical_gradient_and_stresstensor(cell, get_energy) + ref_derivatives = np.array([ + [ 0.0207234415228186, -0.0205836620637001, -0.0205837245204066], + [-0.0207234041482707, 0.0205836266786719, 0.0205836975908369], + [-0.4214593518980791, -0.1513553901588693, -0.1513553918641719], + [-0.1513088777471694, -0.4214285198855805, 0.1508046670295471], + [-0.1513088915316985, 0.1508046269549368, -0.4214284877690488], + ]) + + assert np.abs(test_energy - ref_energy) < 1e-9 + assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 5e-7 + # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 5e-6 # TODO: Support Becke grid stress tensor + + dm = mf.make_rdm1() + kpts = cell.make_kpts(kmesh) + + test_hcore_derivatives = get_nuc(cell, dm, kpts) + test_hcore_derivatives += kin_derivatives(cell, dm, kpts) + assert test_hcore_derivatives.shape == (cell.natm + 3, 3) + + ref_hcore_derivatives = numerical_hcore_gradient_and_stresstensor(cell, dm, kmesh) + + assert np.max(np.abs(test_hcore_derivatives[:-3, :] - ref_hcore_derivatives[:-3, :])) < 3e-8 + assert np.max(np.abs(test_hcore_derivatives[-3:, :] - ref_hcore_derivatives[-3:, :])) < 3e-7 + + +if __name__ == '__main__': + print("Full Tests for PBC GDF Hcore gradient and stress tensor") + unittest.main() diff --git a/gpu4pyscf/pbc/grad/krhf.py b/gpu4pyscf/pbc/grad/krhf.py index 69abef14c..1f45bfb2d 100644 --- a/gpu4pyscf/pbc/grad/krhf.py +++ b/gpu4pyscf/pbc/grad/krhf.py @@ -149,92 +149,6 @@ def hcore_deriv(atm_id): return hcore return hcore_deriv -def get_nuc_strain_deriv(mf_grad, cell, dm, kpts): - '''Strain derivatives for nuclear attraction or pp-local with k-points sampling - - This function is deprecated. - ''' - from gpu4pyscf.lib.cupy_helper import sandwich_dot - from gpu4pyscf.pbc.grad.krks_stress import ( - _eval_ao_strain_derivatives, _get_vpplocG_strain_derivatives, - _get_coulG_strain_derivatives, ALIGNED) - assert cell.low_dim_ft_type != 'inf_vacuum' - assert cell.dimension != 1 - assert kpts.ndim == 2 - assert dm.ndim == 3 - if not cell.cart: - c2s = asarray(cell.cart2sph_coeff()) - dm = sandwich_dot(dm, c2s.T) - # Ensure all AOs are evaluated in the Cartesian GTOs as ao_ks strain - # derivatives currently supports Cartesian format only - cell = cell.copy() - cell.cart = True - nkpts, nao = dm.shape[:2] - assert nkpts == len(kpts) - - grids = UniformGrids(cell) - grids_idx = grids.argsort(tile=8) - grids_coords = grids.coords[grids_idx] - ngrids = len(grids_coords) - mesh = grids.mesh - - def partial_dot(bra, ket): - '''conj(ig),ig->g''' - rho = cp.einsum('ig,ig->g', bra.real, ket.real) - rho += cp.einsum('ig,ig->g', bra.imag, ket.imag) - return rho - - eval_gto_opt = _GTOvalOpt(cell, kpts, deriv=1) - max_memory = 4e9 - blksize = int((max_memory/16/(nkpts*10*nao))/ ALIGNED) * ALIGNED - - rho0 = cp.zeros(ngrids) - rho1 = cp.zeros((3,3, ngrids)) - - for p0, p1 in lib.prange(0, ngrids, blksize): - coords = cp.asarray(grids_coords[p0:p1].T, order='C').T - ao_ks = eval_ao_kpts(cell, coords, kpts, deriv=1, opt=eval_gto_opt) - ao_ks_strain = _eval_ao_strain_derivatives( - cell, coords, kpts, deriv=0, opt=eval_gto_opt) - coordsT = coords.T - for k, dm_k in enumerate(dm): - ao = ao_ks[k].transpose(0,2,1) - ao_strain = ao_ks_strain[k] - ao1 = ao_strain[:,:,0] - # Adding the response of the grids - ao1 += contract('xig,yg->xyig', ao[1:4], coordsT) - c0 = dm_k.T.dot(ao[0]) - rho0[p0:p1] += partial_dot(ao[0], c0).real - rho1[:,:,p0:p1] += contract('xyig,ig->xyg', ao1, c0.conj()).real - - rho0 *= 1./nkpts - # *2 for rho1 because the derivatives were applied to the bra only - rho1 *= 2./nkpts - - rho0_fft_order = cp.empty_like(rho0) - rho1_fft_order = cp.empty_like(rho1) - rho0_fft_order[grids_idx] = rho0 - rho1_fft_order[:,:,grids_idx] = rho1 - rho0, rho1 = rho0_fft_order, rho1_fft_order - rhoG = pbctools.fft(rho0, mesh) - - if cell._pseudo: - vpplocG_0, vpplocG_1 = _get_vpplocG_strain_derivatives(cell, mesh) - vpplocR = pbctools.ifft(vpplocG_0, mesh).real - Ene = contract('xyg,g->xy', rho1, vpplocR).real.get() - Ene += contract('g,xyg->xy', rhoG.conj(), vpplocG_1).real.get() * (1./ngrids) - else: - Gv = cell.get_Gv(mesh) - coulG_0, coulG_1 = _get_coulG_strain_derivatives(cell, Gv) - # SI corresponds to Fourier components of the fractional atomic - # positions within the cell. It does not respond to the strain - # transformation - ZG = _get_ZSI(cell, mesh) - vR = pbctools.ifft(ZG * coulG_0, mesh).real - Ene = contract('xyg,g->xy', rho1, vR).real.get() - Ene += contract('xyg,g->xy', coulG_1, rhoG.conj()*ZG).real.get() * (1./ngrids) - return Ene - class GradientsBase(pbchf_grad.GradientsBase): ''' Basic nuclear gradient functions for non-relativistic methods From 23f7f2ba21032c19e927473409d4e3dcdd527f8d Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Fri, 18 Sep 2026 17:09:06 +0800 Subject: [PATCH 3/9] Fix linter --- gpu4pyscf/pbc/df/grad/krhf.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gpu4pyscf/pbc/df/grad/krhf.py b/gpu4pyscf/pbc/df/grad/krhf.py index e978f4e18..b5c733787 100644 --- a/gpu4pyscf/pbc/df/grad/krhf.py +++ b/gpu4pyscf/pbc/df/grad/krhf.py @@ -863,7 +863,6 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos fakenuc = aft_cpu._fake_nuc(cell, with_pseudo=with_pseudo) int3c2e_opt = SRInt3c2eOpt(cell, fakenuc, omega=-omega, bvk_kmesh=bvk_kmesh).build() charges = -cp.asarray(cell.atom_charges(), dtype=np.float64) - nuc = int3c2e_opt.contract_auxvec(charges, kpts) cell = int3c2e_opt.cell auxcell = int3c2e_opt.auxcell From b9c65a742189b3ba0ca63504bd6a4fe979e4a224 Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Sun, 20 Sep 2026 11:04:19 +0800 Subject: [PATCH 4/9] Fix stress tests, add UHF/KUKS tests --- .../df/tests/test_pbc_hcore_derivatives.py | 129 ++++++++++++++++-- 1 file changed, 119 insertions(+), 10 deletions(-) diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py index ca8ddb155..0e32b8140 100644 --- a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py +++ b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py @@ -21,10 +21,13 @@ from gpu4pyscf.pbc.dft.krks import KRKS from gpu4pyscf.pbc.scf.hf import RHF from gpu4pyscf.pbc.scf.khf import KRHF +from gpu4pyscf.pbc.scf.uhf import UHF +from gpu4pyscf.pbc.dft.kuks import KUKS from gpu4pyscf.pbc.dft.gen_grid import BeckeGrids from gpu4pyscf.pbc.df.grad.krhf import get_nuc from gpu4pyscf.pbc.gto.int1e import kin_derivatives from gpu4pyscf.lib.multi_gpu import num_devices +from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells def numerical_hcore_gradient_and_stresstensor(cell, dm, kmesh): if dm.ndim == 2: @@ -60,7 +63,6 @@ def get_energy(cell): translation_invariance = np.sum(numerical_gradient, axis=0) assert np.max(np.abs(translation_invariance)) < 1e-8, "Bad numerical gradient" - from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells for i_xyz in range(3): for j_xyz in range(3): cell_p, cell_m = _finite_diff_cells(cell, i_xyz, j_xyz, disp = dx) @@ -102,7 +104,6 @@ def numerical_gradient_and_stresstensor(cell, get_energy): translation_invariance = np.sum(numerical_gradient, axis=0) assert np.max(np.abs(translation_invariance)) < 1e-7, "Bad numerical gradient" - from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells for i_xyz in range(3): for j_xyz in range(3): cell_p, cell_m = _finite_diff_cells(cell, i_xyz, j_xyz, disp = dx) @@ -144,7 +145,7 @@ def get_energy(cell): gobj = mf.Gradients() test_gradient = gobj.kernel() - test_derivatives = np.vstack((test_gradient, gobj.stress)) + test_derivatives = np.vstack((test_gradient, gobj.stress * cell.vol)) # Energy check is consistency check ref_energy = -78.84270992380236 @@ -161,7 +162,7 @@ def get_energy(cell): assert np.abs(test_energy - ref_energy) < 1e-9 assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 5e-8 - # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 5e-7 # TODO: Something actually wrong here + assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-7 dm = mf.make_rdm1() kmesh = np.array([1,1,1]) @@ -176,7 +177,7 @@ def get_energy(cell): assert np.max(np.abs(test_hcore_derivatives[:-3, :] - ref_hcore_derivatives[:-3, :])) < 3e-8 assert np.max(np.abs(test_hcore_derivatives[-3:, :] - ref_hcore_derivatives[-3:, :])) < 3e-7 - def test_gdf_hcore_derivatives_rks_large_cell(self): + def test_gdf_hcore_derivatives_rks(self): cell = pyscf.M( atom = """ O 15.43509000 9.59549000 8.94968000 @@ -203,7 +204,7 @@ def get_energy(cell): gobj = mf.Gradients() gobj.grid_response = True test_gradient = gobj.kernel() - test_derivatives = np.vstack((test_gradient, gobj.stress)) + test_derivatives = np.vstack((test_gradient, gobj.stress * cell.vol)) # Energy check is consistency check ref_energy = -79.23783762162086 @@ -261,7 +262,7 @@ def get_energy(cell): gobj = mf.Gradients() test_gradient = gobj.kernel() - test_derivatives = np.vstack((test_gradient, gobj.stress)) + test_derivatives = np.vstack((test_gradient, gobj.stress * cell.vol)) # Energy check is consistency check ref_energy = -75.68294188844864 @@ -277,7 +278,7 @@ def get_energy(cell): assert np.abs(test_energy - ref_energy) < 1e-9 assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 5e-7 - # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 5e-6 # TODO: Something actually wrong here + assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 5e-7 dm = mf.make_rdm1() kpts = cell.make_kpts(kmesh) @@ -300,7 +301,7 @@ def test_gdf_hcore_derivatives_krks(self): atom = 'C 0., 0., 0.; C 0.8917, 0.8917, 0.8917', basis = 'def2-svp', mesh = [10086] * 3, - verbose = 4, + verbose = 0, ) kmesh = np.array([3,1,1]) @@ -319,7 +320,7 @@ def get_energy(cell): gobj = mf.Gradients() gobj.grid_response = True test_gradient = gobj.kernel() - test_derivatives = np.vstack((test_gradient, gobj.stress)) + test_derivatives = np.vstack((test_gradient, gobj.stress * cell.vol)) # Energy check is consistency check ref_energy = -75.65250113719456 @@ -348,6 +349,114 @@ def get_energy(cell): assert np.max(np.abs(test_hcore_derivatives[:-3, :] - ref_hcore_derivatives[:-3, :])) < 3e-8 assert np.max(np.abs(test_hcore_derivatives[-3:, :] - ref_hcore_derivatives[-3:, :])) < 3e-7 + def test_gdf_hcore_derivatives_uhf(self): + cell = pyscf.M( + atom = """ + O 15.43509000 9.59549000 8.94968000 + H 15.05724000 9.21878000 9.73314000 + H 0.51550474 9.33856000 9.01857000 + He 2.51550474 9.33856000 9.01857000 + """, + a = np.eye(3) * 15.9069652593, + unit = "Angstrom", + charge = 1, + spin = 1, + basis = "6-31g", + verbose = 0, + ) + + def get_energy(cell): + mf = UHF(cell).density_fit(auxbasis="def2-universal-jkfit") + mf.conv_tol = 1e-11 + e = mf.kernel() + assert mf.converged + return e, mf + test_energy, mf = get_energy(cell) + + gobj = mf.Gradients() + gobj.grid_response = True + test_gradient = gobj.kernel() + test_derivatives = np.vstack((test_gradient, gobj.stress * cell.vol)) + + # Energy check is consistency check + ref_energy = -78.48457315211476 + # ref_derivatives = numerical_gradient_and_stresstensor(cell, get_energy) + ref_derivatives = np.array([ + [-0.0555001189184168, 0.00711174671153 , 0.0089337752484653], + [ 0.0414229687351053, 0.0090893369275591, -0.0309807412435248], + [ 0.0139675427135444, -0.0163045759649094, 0.0220698445474454], + [ 0.0001096094592867, 0.0001034923968746, -0.0000228787655487], + [ 0.0143442309763486, -0.0362892617999933, 0.0630823853953189], + [-0.0362892622263189, 0.0190900627927704, 0.0113875727691948], + [ 0.0630823866032415, 0.0113875747587144, -0.0252774496800612], + ]) + + assert np.abs(test_energy - ref_energy) < 1e-9 + assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 1e-7 + assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-7 + + def test_gdf_hcore_derivatives_kuks(self): + cell = pyscf.M( + a = '''0. 1.7834 1.8834 + 1.7834 0. 1.7834 + 1.7834 1.7834 0. ''', + atom = 'C 0., 0., 0.; C 0.8917, 0.9017, 0.8917', + basis = """ + BASIS "ao basis" SPHERICAL PRINT + #BASIS SET: (12s,6p) -> [2s,1p] + C S + 0.7427370491E+03 0.9163596281E-02 + 0.1361800249E+03 0.4936149294E-01 + 0.3809826352E+02 0.1685383049E+00 + 0.1308778177E+02 0.3705627997E+00 + 0.5082368648E+01 0.4164915298E+00 + 0.2093200076E+01 0.1303340841E+00 + C SP + 0.3049723950E+02 -0.1325278809E-01 0.3759696623E-02 + 0.6036199601E+01 -0.4699171014E-01 0.3767936984E-01 + 0.1876046337E+01 -0.3378537151E-01 0.1738967435E+00 + 0.7217826470E+00 0.2502417861E+00 0.4180364347E+00 + # 0.3134706954E+00 0.5951172526E+00 0.4258595477E+00 + # 0.1436865550E+00 0.2407061763E+00 0.1017082955E+00 + END + """, # Modified sto-6g + mesh = [10086] * 3, + precision = 1e-9, + verbose = 0, + ) + + kmesh = np.array([3,1,1]) + + def get_energy(cell): + kpts = cell.make_kpts(kmesh) + mf = KUKS(cell, xc="PBE0", kpts=kpts).density_fit(auxbasis="def2-universal-jkfit") + mf.grids = BeckeGrids(cell) + mf.grids.atom_grid = (99, 590) + mf.conv_tol = 1e-11 + e = mf.kernel() + assert mf.converged + return e, mf + test_energy, mf = get_energy(cell) + + gobj = mf.Gradients() + gobj.grid_response = True + test_gradient = gobj.kernel() + test_derivatives = np.vstack((test_gradient, gobj.stress * cell.vol)) + + # Energy check is consistency check + ref_energy = -73.96736967398559 + # ref_derivatives = numerical_gradient_and_stresstensor(cell, get_energy) + ref_derivatives = np.array([ + [-0.0432393890292815, 0.0152970496714033, 0.0907656186655004], + [ 0.0432393193250391, -0.0152969768407729, -0.0907655656590123], + [ 1.2659054821284599, -0.2414305584608201, -0.2869095677482392], + [-0.2414146290874442, 1.2612306482395752, 0.305671657088169 ], + [-0.2868897047392238, 0.3056643988941232, 1.4221600833508319], + ]) + + assert np.abs(test_energy - ref_energy) < 1e-9 + assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 1e-6 + # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-6 # TODO: Support Becke grid stress tensor if __name__ == '__main__': print("Full Tests for PBC GDF Hcore gradient and stress tensor") From abdd2978e976b571eb7762223d4302f68437c5e5 Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Mon, 21 Sep 2026 11:47:50 +0800 Subject: [PATCH 5/9] Recover FFTDF derivative for other multigrid cases, fix the tests --- gpu4pyscf/pbc/grad/krhf.py | 24 +++++++++++++++++++++--- gpu4pyscf/pbc/grad/rhf.py | 5 +++-- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/gpu4pyscf/pbc/grad/krhf.py b/gpu4pyscf/pbc/grad/krhf.py index 7da6cc626..77f637f43 100644 --- a/gpu4pyscf/pbc/grad/krhf.py +++ b/gpu4pyscf/pbc/grad/krhf.py @@ -148,6 +148,25 @@ def hcore_deriv(atm_id): return hcore return hcore_deriv +def get_nuc_fftdf(mf_grad, cell, dm0, kpts): + nkpts = len(kpts) + if nkpts == 1: + if dm0.ndim == 2: + dm0 = dm0[None, :, :] + grad_sigma = np.zeros((cell.natm + 3, 3)) + hcore_deriv = hcore_generator(mf_grad, cell, kpts) + dh1e = cp.empty([cell.natm, 3]) + for ia in range(cell.natm): + h1ao = hcore_deriv(ia) + dh1e[ia] = cp.einsum('kxij,kji->x', h1ao, dm0).real + grad_sigma[:-3] += dh1e.get() / nkpts + # hcore_generator includes kinetic gradients, but not kinetic strain. + grad_sigma[-3:] += int1e.kin_derivatives(cell, dm0, kpts)[-3:] + ni = multigrid_v3.MultiGridNumInt(cell) + grad_sigma[-3:] += ni.energy_strain_gradient( + 'HF', dm0, kpts, spin=0, with_j=False, with_nuc=True) + return grad_sigma + class GradientsBase(pbchf_grad.GradientsBase): ''' Basic nuclear gradient functions for non-relativistic methods @@ -246,7 +265,6 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None): else: is_uhf = mf.istype('UHF') kpts = mf.kpt - nkpts = len(kpts) if getattr(mf, 'disp', None): raise NotImplementedError('dispersion correction') @@ -270,9 +288,9 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None): # Vne or pploc contribution is evaluated in energy_ee grad_sigma += int1e.kin_derivatives(cell, dm0, kpts) elif isinstance(ni, multigrid.MultiGridNumIntBase): - raise NotImplementedError("") + grad_sigma += get_nuc_fftdf(self, cell, dm0, kpts) elif np.prod(cell.mesh) < pbchf.ALLOWED_FFT_MESH_SIZE: - raise NotImplementedError("") + grad_sigma += get_nuc_fftdf(self, cell, dm0, kpts) else: grad_sigma += get_nuc(cell, dm0, kpts) grad_sigma += int1e.kin_derivatives(cell, dm0, kpts) diff --git a/gpu4pyscf/pbc/grad/rhf.py b/gpu4pyscf/pbc/grad/rhf.py index 18c3100e8..eb7d8013f 100644 --- a/gpu4pyscf/pbc/grad/rhf.py +++ b/gpu4pyscf/pbc/grad/rhf.py @@ -190,14 +190,15 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None): if is_uhf: dm0 = dm0[0] + dm0[1] + from gpu4pyscf.pbc.grad.krhf import get_nuc_fftdf ni = mf._numint if isinstance(ni, multigrid_v3.MultiGridNumInt): # Vne or pploc contribution is evaluated in energy_ee grad_sigma += int1e.kin_derivatives(cell, dm0) elif isinstance(ni, multigrid.MultiGridNumIntBase): - raise NotImplementedError("") + grad_sigma += get_nuc_fftdf(self, cell, dm0, np.zeros((1,3))) elif np.prod(cell.mesh) < pbchf.ALLOWED_FFT_MESH_SIZE: - raise NotImplementedError("") + grad_sigma += get_nuc_fftdf(self, cell, dm0, np.zeros((1,3))) else: grad_sigma += get_nuc(cell, dm0) grad_sigma += int1e.kin_derivatives(cell, dm0) From 2ef9040c6e15f57bf72cfb9a3cffbd5527759095 Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Mon, 21 Sep 2026 14:59:18 +0800 Subject: [PATCH 6/9] Safeguard for pseudopotential gradient under GDF mode --- gpu4pyscf/pbc/df/grad/krhf.py | 22 +++++++++------------- gpu4pyscf/pbc/grad/krhf.py | 7 +++++-- gpu4pyscf/pbc/grad/rhf.py | 7 +++++-- 3 files changed, 19 insertions(+), 17 deletions(-) diff --git a/gpu4pyscf/pbc/df/grad/krhf.py b/gpu4pyscf/pbc/df/grad/krhf.py index dc4839d5e..6cd62207d 100644 --- a/gpu4pyscf/pbc/df/grad/krhf.py +++ b/gpu4pyscf/pbc/df/grad/krhf.py @@ -934,19 +934,9 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos ft_opt = ft_ao.FTOpt.from_intopt(int3c2e_opt) if with_pseudo: - raise NotImplementedError("") - # #TODO: call multigrid.eval_vpplocG after removing its part2 contribution - # ZG = ft_ao.ft_ao(fakenuc, Gv).conj() - # ZG = ZG.dot(charges) - # ZG *= _weighted_coulG_LR(cell, Gv, omega, kws) - # if ((cell.dimension == 3 or - # (cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum'))): - # exps = cp.asarray(np.hstack(fakenuc.bas_exps())) - # ZG[0] -= charges.dot(np.pi/exps) / cell.vol + raise NotImplementedError("get_pp_loc_part1_grad(with_pseudo = True) not implemented yet") else: pass - # ZG = _get_ZSI(cell, mesh).conj() - # ZG *= _weighted_coulG_LR(cell, Gv, omega, kws) bvk_ncells = len(int3c2e_opt.bvkmesh_Ls) aux_loc = auxcell.ao_loc @@ -1106,5 +1096,11 @@ def get_nuc(cell, dm, kpts=None, hermi=1): log.timer('get_nuc gradient', *t0) return nuc -def get_pp(cell, dm, kpts=None, hermi=1): - raise NotImplementedError("") +def get_pp_loc(cell, dm, kpts=None, hermi=1): + log = logger.new_logger(cell) + t0 = log.init_timer() + depp = get_pp_loc_part1_grad(cell, dm, kpts, hermi, with_pseudo=True, verbose=log) + raise NotImplementedError("get_pp_loc_part2_grad not implemented yet") + # depp += get_pp_loc_part2_grad(cell, dm, kpts, hermi) + log.timer('get_pp_loc gradient', *t0) + return depp diff --git a/gpu4pyscf/pbc/grad/krhf.py b/gpu4pyscf/pbc/grad/krhf.py index 77f637f43..b69404043 100644 --- a/gpu4pyscf/pbc/grad/krhf.py +++ b/gpu4pyscf/pbc/grad/krhf.py @@ -38,7 +38,7 @@ from gpu4pyscf.pbc.grad.rhf import contract_h1e_dm, _get_ejk_derivatives from gpu4pyscf.pbc.grad import rhf as pbchf_grad from gpu4pyscf.pbc.scf import hf as pbchf -from gpu4pyscf.pbc.df.grad.krhf import get_nuc +from gpu4pyscf.pbc.df.grad.krhf import get_nuc, get_pp_loc __all__ = ['Gradients'] @@ -292,7 +292,10 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None): elif np.prod(cell.mesh) < pbchf.ALLOWED_FFT_MESH_SIZE: grad_sigma += get_nuc_fftdf(self, cell, dm0, kpts) else: - grad_sigma += get_nuc(cell, dm0, kpts) + if cell._pseudo: + grad_sigma += get_pp_loc(cell, dm0, kpts) + else: + grad_sigma += get_nuc(cell, dm0, kpts) grad_sigma += int1e.kin_derivatives(cell, dm0, kpts) if cell._pseudo: diff --git a/gpu4pyscf/pbc/grad/rhf.py b/gpu4pyscf/pbc/grad/rhf.py index eb7d8013f..eed252ff7 100644 --- a/gpu4pyscf/pbc/grad/rhf.py +++ b/gpu4pyscf/pbc/grad/rhf.py @@ -35,7 +35,7 @@ from gpu4pyscf.pbc.grad.pp import ppnl_derivatives from gpu4pyscf.gto.mole import groupby from gpu4pyscf.pbc.scf import hf as pbchf -from gpu4pyscf.pbc.df.grad.krhf import get_nuc +from gpu4pyscf.pbc.df.grad.krhf import get_nuc, get_pp_loc __all__ = ['Gradients'] @@ -200,7 +200,10 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None): elif np.prod(cell.mesh) < pbchf.ALLOWED_FFT_MESH_SIZE: grad_sigma += get_nuc_fftdf(self, cell, dm0, np.zeros((1,3))) else: - grad_sigma += get_nuc(cell, dm0) + if cell._pseudo: + grad_sigma += get_pp_loc(cell, dm0) + else: + grad_sigma += get_nuc(cell, dm0) grad_sigma += int1e.kin_derivatives(cell, dm0) if cell._pseudo: From 13abb1e4c585255a96450b003e4f5ba5f6ddab31 Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Mon, 21 Sep 2026 16:10:27 +0800 Subject: [PATCH 7/9] Support pseudopotential for get_pp_loc_part1_grad. No tests because get_pp_loc_part2_grad not implemented --- gpu4pyscf/pbc/df/grad/krhf.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/gpu4pyscf/pbc/df/grad/krhf.py b/gpu4pyscf/pbc/df/grad/krhf.py index 6cd62207d..937c6e073 100644 --- a/gpu4pyscf/pbc/df/grad/krhf.py +++ b/gpu4pyscf/pbc/df/grad/krhf.py @@ -934,9 +934,12 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos ft_opt = ft_ao.FTOpt.from_intopt(int3c2e_opt) if with_pseudo: - raise NotImplementedError("get_pp_loc_part1_grad(with_pseudo = True) not implemented yet") + assert (cell.dimension == 3 or + (cell.dimension == 2 and cell.low_dim_ft_type != 'inf_vacuum')) + exps = cp.asarray(np.hstack(fakenuc.bas_exps())) + pp_G0_term = -charges.dot(np.pi/exps) * kws else: - pass + pp_G0_term = 0 bvk_ncells = len(int3c2e_opt.bvkmesh_Ls) aux_loc = auxcell.ao_loc @@ -999,6 +1002,7 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos ej_sigma_lr = cp.zeros([cell.natm+3, 3]) vG_conj = rho_nucG.conj() * wcoulG_LR0 + vG_conj[0] += pp_G0_term bas_ij_idx, bas_ij_img_idx, shl_pair_offsets = aft_jk._generate_shl_pairs(ft_opt) nbatches_shl_pair = len(shl_pair_offsets) - 1 err = libpbc.PBC_ft_aopair_ej_deriv( @@ -1022,6 +1026,7 @@ def get_pp_loc_part1_grad(cell, dm, kpts=None, hermi=0, with_pseudo=True, verbos ej_sigma_lr += ej_sigma_aux ej_sigma_lr[-3:] += cp.einsum( 'g,g,xyg->xy', rho_nucG, rhoG.conj(), wcoulG_LR1).real + ej_sigma_lr[-3:] -= cp.eye(3) * (rhoG[0] * pp_G0_term).real ej_sigma = ej_sigma_lr t0 = log.timer_debug1('lr_int3c2e_deriv via aft', *t0) From 12056ca1ea076a750ba10ec7eecbf428f8299c6b Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Tue, 22 Sep 2026 09:14:34 +0800 Subject: [PATCH 8/9] Give two tests initial guesses, to make them more stable --- .../df/tests/test_pbc_hcore_derivatives.py | 74 ++++++++++++++++++- 1 file changed, 72 insertions(+), 2 deletions(-) diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py index 0e32b8140..1bbb865f2 100644 --- a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py +++ b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py @@ -135,10 +135,43 @@ def test_gdf_hcore_derivatives_rhf(self): verbose = 0, ) + mocc = cp.array([ + [-0.99041089, 0.27950001, -0.01848335, 0.01333192, -0.0938811 , -0.00035243], + [ 0.03294304, 0.55589312, -0.0346783 , 0.03025806, -0.19235982, 0.00005631], + [-0.01081498, 0.30281873, -0.04396756, 0.04771887, -0.32861852, -0.00554105], + [ 0.00088694, 0.02393521, -0.0185682 , 0.40413549, 0.27122707, 0.11677408], + [-0.00137754, -0.04407428, -0.00008411, 0.06637102, -0.28027797, 0.53598057], + [ 0.00203597, 0.06546318, -0.00390287, -0.26610883, 0.37501345, 0.31357959], + [-0.00056695, 0.00071742, -0.01437576, 0.19770758, 0.1819253 , 0.09090388], + [ 0.00016489, 0.001223 , 0.00189027, 0.03116818, -0.18709963, 0.42450686], + [-0.00009328, -0.00273068, -0.00429873, -0.12778522, 0.25098631, 0.2480971 ], + [ 0.00002125, -0.00035278, 0.00042744, -0.00990683, -0.00766189, 0.00437554], + [-0.00049681, -0.00333475, 0.00037362, 0.00916205, -0.01276258, 0.00661061], + [ 0.00019907, 0.00136052, -0.0006717 , -0.01543899, 0.00443065, 0.01133865], + [-0.00042706, -0.00152516, -0.00000493, 0.0102507 , 0.00818746, 0.00601729], + [ 0.00037305, 0.00190751, 0.00056498, 0.01045524, 0.00176675, 0.00920268], + [-0.0019224 , 0.20269791, -0.00084657, -0.3328666 , 0.1858071 , 0.00004932], + [ 0.00177218, 0.00524372, 0.00783172, -0.09101972, 0.04496627, 0.00100695], + [-0.00098335, 0.01947296, -0.00329651, -0.00477764, 0.02907159, 0.00542444], + [-0.00103002, 0.01279762, -0.0001244 , -0.01765917, -0.00176007, 0.02465532], + [ 0.0021158 , -0.02927535, -0.0004487 , 0.03196566, -0.00783496, 0.01441768], + [-0.00086381, 0.16426041, -0.00680487, 0.29963924, 0.22218589, 0.00143911], + [ 0.00187774, 0.00623761, 0.02225619, 0.10326297, 0.07956277, 0.00637142], + [ 0.00182245, -0.03220508, 0.00388776, -0.03057134, -0.02583837, 0.00471044], + [-0.00046351, 0.0053786 , -0.0001501 , 0.01389799, -0.00413402, 0.02216759], + [ 0.00012458, 0.00283229, -0.00009179, -0.01194289, 0.01685874, 0.01286924], + [ 0.0001066 , 0.01915041, 0.59057298, 0.02030031, 0.00478559, 0.01121834], + [-0.00046685, 0.01444229, 0.5073927 , 0.01595701, -0.0086483 , 0.009216 ], + [ 0.00022654, 0.00074338, -0.00314268, -0.00591431, -0.00518211, -0.00260137], + [ 0.0000296 , 0.00002835, -0.00102071, 0.00007384, -0.001246 , 0.00162482], + [-0.00004518, 0.00024298, 0.00197128, -0.00108762, 0.00143038, 0.00107136], + ]) + dm0 = 2 * mocc @ mocc.conj().T + def get_energy(cell): mf = RHF(cell).density_fit(auxbasis="def2-universal-jkfit") mf.conv_tol = 1e-12 - e = mf.kernel() + e = mf.kernel(dm0 = dm0) assert mf.converged return e, mf test_energy, mf = get_energy(cell) @@ -427,13 +460,50 @@ def test_gdf_hcore_derivatives_kuks(self): kmesh = np.array([3,1,1]) + mocc = cp.array([ + [[-0.70267367+0.j , -0.70280126+0.j , -0.14152445+0.j , 0.0005536 +0.j , -0.00015479+0.j , -0.00617958+0.j ], + [-0.02580572+0.j , -0.02071216+0.j , 0.64238801+0.j , -0.00466737+0.j , 0.00132964+0.j , 0.03662597+0.j ], + [ 0.00052735+0.j , 0.00006304+0.j , -0.01838221+0.j , -0.29081125+0.j , -0.41949067+0.j , 0.43157367+0.j ], + [-0.00048372+0.j , -0.00006852+0.j , 0.01679578+0.j , -0.47540259+0.j , -0.12366065+0.j , -0.45328749+0.j ], + [-0.00049458+0.j , -0.00011871+0.j , 0.01334374+0.j , 0.35959094+0.j , -0.50176531+0.j , -0.24720594+0.j ], + [ 0.70245928+0.j , -0.70301555+0.j , -0.14152444+0.j , 0.00055366+0.j , -0.00015479+0.j , -0.00617981+0.j ], + [ 0.0257994 +0.j , -0.02072003+0.j , 0.64238797+0.j , -0.00466772+0.j , 0.00132965+0.j , 0.03662717+0.j ], + [ 0.00052733+0.j , -0.00006321+0.j , 0.01838167+0.j , 0.2908117 +0.j , 0.41949008+0.j , -0.43157348+0.j ], + [-0.0004837 +0.j , 0.00006867+0.j , -0.01679545+0.j , 0.47540182+0.j , 0.12366143+0.j , 0.45328734+0.j ], + [-0.00049455+0.j , 0.00011887+0.j , -0.01334336+0.j , -0.35959129+0.j , 0.50176606+0.j , 0.24720687+0.j ]], + + [[ 0.69185824+0.12284597j, -0.12244734+0.6920467j , 0.00169982-0.13481846j, -0.07151506+0.01313129j, 0.00212679-0.00055695j, 0.00206022+0.00020588j], + [ 0.02481703+0.00400557j, -0.00410915+0.02102312j, -0.00732432+0.61140176j, 0.32924715-0.06118735j, -0.01014939+0.00384823j, -0.00968474-0.00121962j], + [-0.00100367-0.00054517j, -0.00041984+0.00048512j, 0.09788813-0.06128097j, 0.23987595+0.19971993j, -0.07452514+0.25563021j, 0.10681477-0.45887386j], + [ 0.0009512 +0.0005312j , 0.00041264-0.00047494j, -0.09683854+0.05837405j, -0.2413802 -0.18926148j, -0.12475365+0.5287418j , 0.00445371+0.02468709j], + [ 0.00104854+0.00059397j, 0.00046682-0.0004842j , -0.10133274+0.06101002j, -0.24174223-0.23674304j, 0.07002161-0.23725642j, 0.12332986-0.43857581j], + [-0.69061639-0.12889505j, -0.12864782+0.69106133j, 0.04594042-0.12676118j, 0.06229421+0.03749984j, -0.00179383-0.00127096j, -0.00200049-0.00053379j], + [-0.02463188-0.00499888j, -0.00354373+0.02113079j, -0.20869953+0.57472646j, -0.28653729-0.17333029j, 0.00814067+0.00717918j, 0.00949279+0.0022736j ], + [-0.00113154+0.00015643j, -0.0002225 -0.00060151j, 0.06999785+0.09185642j, 0.29488396-0.1023322j , 0.02040371-0.26548856j, -0.06191258+0.46705671j], + [ 0.00107751-0.00016185j, 0.00021933+0.00058946j, -0.07004041-0.0887663j , -0.29260408+0.09201507j, 0.06970666-0.53876963j, 0.01287501-0.02152976j], + [ 0.00119074-0.00018628j, 0.0002668 +0.00061715j, -0.07331696-0.09281752j, -0.3096849 +0.1363204j , -0.01813876+0.24670676j, -0.03929987+0.45388754j]], + + [[ 0.69185824-0.12284597j, -0.12244734-0.6920467j , 0.00169982+0.13481846j, -0.07151506-0.01313129j, 0.00212679+0.00055695j, 0.00206022-0.00020588j], + [ 0.02481703-0.00400557j, -0.00410915-0.02102312j, -0.00732432-0.61140176j, 0.32924715+0.06118735j, -0.01014939-0.00384823j, -0.00968474+0.00121962j], + [-0.00100367+0.00054517j, -0.00041984-0.00048512j, 0.09788813+0.06128097j, 0.23987595-0.19971993j, -0.07452514-0.25563021j, 0.10681477+0.45887386j], + [ 0.0009512 -0.0005312j , 0.00041264+0.00047494j, -0.09683854-0.05837405j, -0.2413802 +0.18926148j, -0.12475365-0.5287418j , 0.00445371-0.02468709j], + [ 0.00104854-0.00059397j, 0.00046682+0.0004842j , -0.10133274-0.06101002j, -0.24174223+0.23674304j, 0.07002161+0.23725642j, 0.12332986+0.43857581j], + [-0.69061639+0.12889505j, -0.12864782-0.69106133j, 0.04594042+0.12676118j, 0.06229421-0.03749984j, -0.00179383+0.00127096j, -0.00200049+0.00053379j], + [-0.02463188+0.00499888j, -0.00354373-0.02113079j, -0.20869953-0.57472646j, -0.28653729+0.17333029j, 0.00814067-0.00717918j, 0.00949279-0.0022736j ], + [-0.00113154-0.00015643j, -0.0002225 +0.00060151j, 0.06999785-0.09185642j, 0.29488396+0.1023322j , 0.02040371+0.26548856j, -0.06191258-0.46705671j], + [ 0.00107751+0.00016185j, 0.00021933-0.00058946j, -0.07004041+0.0887663j , -0.29260408-0.09201507j, 0.06970666+0.53876963j, 0.01287501+0.02152976j], + [ 0.00119074+0.00018628j, 0.0002668 -0.00061715j, -0.07331696+0.09281752j, -0.3096849 -0.1363204j , -0.01813876-0.24670676j, -0.03929987-0.45388754j]], + ]) + dm0 = cp.zeros([2, 3, cell.nao, cell.nao], dtype=cp.complex128) + dm0[0, :, :, :] = dm0[1, :, :, :] = cp.einsum("kpi,kqi->kpq", mocc, mocc.conj()) + def get_energy(cell): kpts = cell.make_kpts(kmesh) mf = KUKS(cell, xc="PBE0", kpts=kpts).density_fit(auxbasis="def2-universal-jkfit") mf.grids = BeckeGrids(cell) mf.grids.atom_grid = (99, 590) mf.conv_tol = 1e-11 - e = mf.kernel() + e = mf.kernel(dm0 = dm0) assert mf.converged return e, mf test_energy, mf = get_energy(cell) From 5484656416caf1a9c3c24b11f4cbe91fefe3cdd3 Mon Sep 17 00:00:00 2001 From: "henry.wang1" Date: Tue, 22 Sep 2026 16:04:42 +0800 Subject: [PATCH 9/9] Fix the tests for pyscf==2.8.0 --- gpu4pyscf/_patch_pyscf.py | 3 - .../df/tests/test_pbc_hcore_derivatives.py | 76 +------------------ gpu4pyscf/pbc/gto/int1e.py | 4 +- 3 files changed, 6 insertions(+), 77 deletions(-) diff --git a/gpu4pyscf/_patch_pyscf.py b/gpu4pyscf/_patch_pyscf.py index c0f5985a9..4452eb000 100644 --- a/gpu4pyscf/_patch_pyscf.py +++ b/gpu4pyscf/_patch_pyscf.py @@ -61,9 +61,6 @@ def get_lattice_Ls(cell, nimgs=None, rcut=None, dimension=None, discard=True): scaled_atom_coords = cell.get_scaled_atom_coords() atom_boundary_max = scaled_atom_coords[:,:dimension].max(axis=0) atom_boundary_min = scaled_atom_coords[:,:dimension].min(axis=0) - if (np.any(atom_boundary_max > 1) or np.any(atom_boundary_min < -1)): - atom_boundary_max[atom_boundary_max > 1] = 1 - atom_boundary_min[atom_boundary_min <-1] = -1 ovlp_penalty = atom_boundary_max - atom_boundary_min dR = ovlp_penalty.dot(a[:dimension]) dR_basis = np.diag(dR) diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py index 1bbb865f2..8c3c2230f 100644 --- a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py +++ b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py @@ -135,43 +135,10 @@ def test_gdf_hcore_derivatives_rhf(self): verbose = 0, ) - mocc = cp.array([ - [-0.99041089, 0.27950001, -0.01848335, 0.01333192, -0.0938811 , -0.00035243], - [ 0.03294304, 0.55589312, -0.0346783 , 0.03025806, -0.19235982, 0.00005631], - [-0.01081498, 0.30281873, -0.04396756, 0.04771887, -0.32861852, -0.00554105], - [ 0.00088694, 0.02393521, -0.0185682 , 0.40413549, 0.27122707, 0.11677408], - [-0.00137754, -0.04407428, -0.00008411, 0.06637102, -0.28027797, 0.53598057], - [ 0.00203597, 0.06546318, -0.00390287, -0.26610883, 0.37501345, 0.31357959], - [-0.00056695, 0.00071742, -0.01437576, 0.19770758, 0.1819253 , 0.09090388], - [ 0.00016489, 0.001223 , 0.00189027, 0.03116818, -0.18709963, 0.42450686], - [-0.00009328, -0.00273068, -0.00429873, -0.12778522, 0.25098631, 0.2480971 ], - [ 0.00002125, -0.00035278, 0.00042744, -0.00990683, -0.00766189, 0.00437554], - [-0.00049681, -0.00333475, 0.00037362, 0.00916205, -0.01276258, 0.00661061], - [ 0.00019907, 0.00136052, -0.0006717 , -0.01543899, 0.00443065, 0.01133865], - [-0.00042706, -0.00152516, -0.00000493, 0.0102507 , 0.00818746, 0.00601729], - [ 0.00037305, 0.00190751, 0.00056498, 0.01045524, 0.00176675, 0.00920268], - [-0.0019224 , 0.20269791, -0.00084657, -0.3328666 , 0.1858071 , 0.00004932], - [ 0.00177218, 0.00524372, 0.00783172, -0.09101972, 0.04496627, 0.00100695], - [-0.00098335, 0.01947296, -0.00329651, -0.00477764, 0.02907159, 0.00542444], - [-0.00103002, 0.01279762, -0.0001244 , -0.01765917, -0.00176007, 0.02465532], - [ 0.0021158 , -0.02927535, -0.0004487 , 0.03196566, -0.00783496, 0.01441768], - [-0.00086381, 0.16426041, -0.00680487, 0.29963924, 0.22218589, 0.00143911], - [ 0.00187774, 0.00623761, 0.02225619, 0.10326297, 0.07956277, 0.00637142], - [ 0.00182245, -0.03220508, 0.00388776, -0.03057134, -0.02583837, 0.00471044], - [-0.00046351, 0.0053786 , -0.0001501 , 0.01389799, -0.00413402, 0.02216759], - [ 0.00012458, 0.00283229, -0.00009179, -0.01194289, 0.01685874, 0.01286924], - [ 0.0001066 , 0.01915041, 0.59057298, 0.02030031, 0.00478559, 0.01121834], - [-0.00046685, 0.01444229, 0.5073927 , 0.01595701, -0.0086483 , 0.009216 ], - [ 0.00022654, 0.00074338, -0.00314268, -0.00591431, -0.00518211, -0.00260137], - [ 0.0000296 , 0.00002835, -0.00102071, 0.00007384, -0.001246 , 0.00162482], - [-0.00004518, 0.00024298, 0.00197128, -0.00108762, 0.00143038, 0.00107136], - ]) - dm0 = 2 * mocc @ mocc.conj().T - def get_energy(cell): mf = RHF(cell).density_fit(auxbasis="def2-universal-jkfit") mf.conv_tol = 1e-12 - e = mf.kernel(dm0 = dm0) + e = mf.kernel() assert mf.converged return e, mf test_energy, mf = get_energy(cell) @@ -345,6 +312,7 @@ def get_energy(cell): mf.grids = BeckeGrids(cell) mf.grids.atom_grid = (50, 194) mf.conv_tol = 1e-12 + mf.with_df.linear_dep_threshold = 1e-10 # The default in pyscf==2.8.0 is 1e-9, in pyscf==2.14.0 is 1e-10 e = mf.kernel() assert mf.converged return e, mf @@ -460,50 +428,14 @@ def test_gdf_hcore_derivatives_kuks(self): kmesh = np.array([3,1,1]) - mocc = cp.array([ - [[-0.70267367+0.j , -0.70280126+0.j , -0.14152445+0.j , 0.0005536 +0.j , -0.00015479+0.j , -0.00617958+0.j ], - [-0.02580572+0.j , -0.02071216+0.j , 0.64238801+0.j , -0.00466737+0.j , 0.00132964+0.j , 0.03662597+0.j ], - [ 0.00052735+0.j , 0.00006304+0.j , -0.01838221+0.j , -0.29081125+0.j , -0.41949067+0.j , 0.43157367+0.j ], - [-0.00048372+0.j , -0.00006852+0.j , 0.01679578+0.j , -0.47540259+0.j , -0.12366065+0.j , -0.45328749+0.j ], - [-0.00049458+0.j , -0.00011871+0.j , 0.01334374+0.j , 0.35959094+0.j , -0.50176531+0.j , -0.24720594+0.j ], - [ 0.70245928+0.j , -0.70301555+0.j , -0.14152444+0.j , 0.00055366+0.j , -0.00015479+0.j , -0.00617981+0.j ], - [ 0.0257994 +0.j , -0.02072003+0.j , 0.64238797+0.j , -0.00466772+0.j , 0.00132965+0.j , 0.03662717+0.j ], - [ 0.00052733+0.j , -0.00006321+0.j , 0.01838167+0.j , 0.2908117 +0.j , 0.41949008+0.j , -0.43157348+0.j ], - [-0.0004837 +0.j , 0.00006867+0.j , -0.01679545+0.j , 0.47540182+0.j , 0.12366143+0.j , 0.45328734+0.j ], - [-0.00049455+0.j , 0.00011887+0.j , -0.01334336+0.j , -0.35959129+0.j , 0.50176606+0.j , 0.24720687+0.j ]], - - [[ 0.69185824+0.12284597j, -0.12244734+0.6920467j , 0.00169982-0.13481846j, -0.07151506+0.01313129j, 0.00212679-0.00055695j, 0.00206022+0.00020588j], - [ 0.02481703+0.00400557j, -0.00410915+0.02102312j, -0.00732432+0.61140176j, 0.32924715-0.06118735j, -0.01014939+0.00384823j, -0.00968474-0.00121962j], - [-0.00100367-0.00054517j, -0.00041984+0.00048512j, 0.09788813-0.06128097j, 0.23987595+0.19971993j, -0.07452514+0.25563021j, 0.10681477-0.45887386j], - [ 0.0009512 +0.0005312j , 0.00041264-0.00047494j, -0.09683854+0.05837405j, -0.2413802 -0.18926148j, -0.12475365+0.5287418j , 0.00445371+0.02468709j], - [ 0.00104854+0.00059397j, 0.00046682-0.0004842j , -0.10133274+0.06101002j, -0.24174223-0.23674304j, 0.07002161-0.23725642j, 0.12332986-0.43857581j], - [-0.69061639-0.12889505j, -0.12864782+0.69106133j, 0.04594042-0.12676118j, 0.06229421+0.03749984j, -0.00179383-0.00127096j, -0.00200049-0.00053379j], - [-0.02463188-0.00499888j, -0.00354373+0.02113079j, -0.20869953+0.57472646j, -0.28653729-0.17333029j, 0.00814067+0.00717918j, 0.00949279+0.0022736j ], - [-0.00113154+0.00015643j, -0.0002225 -0.00060151j, 0.06999785+0.09185642j, 0.29488396-0.1023322j , 0.02040371-0.26548856j, -0.06191258+0.46705671j], - [ 0.00107751-0.00016185j, 0.00021933+0.00058946j, -0.07004041-0.0887663j , -0.29260408+0.09201507j, 0.06970666-0.53876963j, 0.01287501-0.02152976j], - [ 0.00119074-0.00018628j, 0.0002668 +0.00061715j, -0.07331696-0.09281752j, -0.3096849 +0.1363204j , -0.01813876+0.24670676j, -0.03929987+0.45388754j]], - - [[ 0.69185824-0.12284597j, -0.12244734-0.6920467j , 0.00169982+0.13481846j, -0.07151506-0.01313129j, 0.00212679+0.00055695j, 0.00206022-0.00020588j], - [ 0.02481703-0.00400557j, -0.00410915-0.02102312j, -0.00732432-0.61140176j, 0.32924715+0.06118735j, -0.01014939-0.00384823j, -0.00968474+0.00121962j], - [-0.00100367+0.00054517j, -0.00041984-0.00048512j, 0.09788813+0.06128097j, 0.23987595-0.19971993j, -0.07452514-0.25563021j, 0.10681477+0.45887386j], - [ 0.0009512 -0.0005312j , 0.00041264+0.00047494j, -0.09683854-0.05837405j, -0.2413802 +0.18926148j, -0.12475365-0.5287418j , 0.00445371-0.02468709j], - [ 0.00104854-0.00059397j, 0.00046682+0.0004842j , -0.10133274-0.06101002j, -0.24174223+0.23674304j, 0.07002161+0.23725642j, 0.12332986+0.43857581j], - [-0.69061639+0.12889505j, -0.12864782-0.69106133j, 0.04594042+0.12676118j, 0.06229421-0.03749984j, -0.00179383+0.00127096j, -0.00200049+0.00053379j], - [-0.02463188+0.00499888j, -0.00354373-0.02113079j, -0.20869953-0.57472646j, -0.28653729+0.17333029j, 0.00814067-0.00717918j, 0.00949279-0.0022736j ], - [-0.00113154-0.00015643j, -0.0002225 +0.00060151j, 0.06999785-0.09185642j, 0.29488396+0.1023322j , 0.02040371+0.26548856j, -0.06191258-0.46705671j], - [ 0.00107751+0.00016185j, 0.00021933-0.00058946j, -0.07004041+0.0887663j , -0.29260408-0.09201507j, 0.06970666+0.53876963j, 0.01287501+0.02152976j], - [ 0.00119074+0.00018628j, 0.0002668 -0.00061715j, -0.07331696+0.09281752j, -0.3096849 -0.1363204j , -0.01813876-0.24670676j, -0.03929987-0.45388754j]], - ]) - dm0 = cp.zeros([2, 3, cell.nao, cell.nao], dtype=cp.complex128) - dm0[0, :, :, :] = dm0[1, :, :, :] = cp.einsum("kpi,kqi->kpq", mocc, mocc.conj()) - def get_energy(cell): kpts = cell.make_kpts(kmesh) mf = KUKS(cell, xc="PBE0", kpts=kpts).density_fit(auxbasis="def2-universal-jkfit") mf.grids = BeckeGrids(cell) mf.grids.atom_grid = (99, 590) mf.conv_tol = 1e-11 - e = mf.kernel(dm0 = dm0) + mf.with_df.linear_dep_threshold = 1e-10 # The default in pyscf==2.8.0 is 1e-9, in pyscf==2.14.0 is 1e-10 + e = mf.kernel() assert mf.converged return e, mf test_energy, mf = get_energy(cell) diff --git a/gpu4pyscf/pbc/gto/int1e.py b/gpu4pyscf/pbc/gto/int1e.py index a1d66e971..032e16134 100644 --- a/gpu4pyscf/pbc/gto/int1e.py +++ b/gpu4pyscf/pbc/gto/int1e.py @@ -20,7 +20,7 @@ from pyscf.gto import ATOM_OF, PTR_COORD, Mole from pyscf.pbc.gto import Cell from pyscf.pbc.gto.cell import _estimate_rcut -from pyscf.pbc.tools.pbc import super_cell, _build_supcell_, get_lattice_Ls +from pyscf.pbc.tools.pbc import super_cell from pyscf.pbc.lib.kpts_helper import is_zero from pyscf.pbc.tools.k2gamma import translation_vectors_for_kmesh from gpu4pyscf.gto.mole import extract_pgto_params @@ -511,7 +511,7 @@ def _shell_overlap_mask(cell, hermi=1, precision=1e-14, Ls=None, envs=None, def _bvkcell_lattice_sum_Ls(bvkcell, rcut=None): if rcut is None: rcut = bvkcell.rcut - Ls = get_lattice_Ls(bvkcell, rcut=rcut, discard=False) + Ls = bvkcell.get_lattice_Ls(rcut=rcut, discard=False) if len(Ls) > 1: r = asarray(bvkcell.atom_coords()) dist_max = dist_matrix(r, r).max().get()