From baf2709fc9f04904f2303492c1b47e3938bb9ede Mon Sep 17 00:00:00 2001 From: Henry Wang Date: Thu, 19 Dec 2024 08:02:51 +0800 Subject: [PATCH] Support multiple density matrices and charge arrays in an inefficient way. --- gpu4pyscf/gto/int3c1e.py | 32 ++++++++++++++++++------- gpu4pyscf/gto/int3c1e_ip.py | 6 ----- gpu4pyscf/gto/tests/test_int1e_grids.py | 24 +++++++++++++++++++ gpu4pyscf/qmmm/chelpg.py | 3 +++ gpu4pyscf/solvent/pcm.py | 12 ++-------- 5 files changed, 53 insertions(+), 24 deletions(-) diff --git a/gpu4pyscf/gto/int3c1e.py b/gpu4pyscf/gto/int3c1e.py index 5d6369f6..bdc4516c 100644 --- a/gpu4pyscf/gto/int3c1e.py +++ b/gpu4pyscf/gto/int3c1e.py @@ -366,12 +366,6 @@ def get_int3c1e_density_contracted(mol, grids, charge_exponents, dm, intopt): assert omega >= 0.0, "Short-range one electron integrals with GPU acceleration is not implemented." dm = cp.asarray(dm) - if dm.ndim == 3: - if dm.shape[0] > 2: - print("Warning: more than two density matrices are found for int3c1e kernel. " - "They will be summed up to one density matrix.") - dm = cp.einsum("ijk->jk", dm) - assert dm.ndim == 2 assert dm.shape[0] == dm.shape[1] and dm.shape[0] == mol.nao @@ -488,8 +482,30 @@ def int1e_grids(mol, grids, charge_exponents=None, dm=None, charges=None, direct if dm is None and charges is None: return get_int3c1e(mol, grids, charge_exponents, intopt) elif dm is not None: - return get_int3c1e_density_contracted(mol, grids, charge_exponents, dm, intopt) + if dm.ndim == 2: + return get_int3c1e_density_contracted(mol, grids, charge_exponents, dm, intopt) + else: + assert dm.ndim == 3 + n_dm = dm.shape[0] + ngrids = grids.shape[0] + if n_dm == 1: + return get_int3c1e_density_contracted(mol, grids, charge_exponents, dm[0], intopt).reshape(1, ngrids) + int3c_density_contracted = cp.empty((n_dm, ngrids)) + for i_dm in range(n_dm): + int3c_density_contracted[i_dm] = get_int3c1e_density_contracted(mol, grids, charge_exponents, dm[i_dm], intopt) + return int3c_density_contracted elif charges is not None: - return get_int3c1e_charge_contracted(mol, grids, charge_exponents, charges, intopt) + if charges.ndim == 1: + return get_int3c1e_charge_contracted(mol, grids, charge_exponents, charges, intopt) + else: + assert charges.ndim == 2 + n_charges = charges.shape[0] + nao = mol.nao + if n_charges == 1: + return get_int3c1e_charge_contracted(mol, grids, charge_exponents, charges[0], intopt).reshape(1, nao, nao) + int3c_charge_contracted = cp.empty((n_charges, nao, nao)) + for i_charge in range(n_charges): + int3c_charge_contracted[i_charge] = get_int3c1e_charge_contracted(mol, grids, charge_exponents, charges[i_charge], intopt) + return int3c_charge_contracted else: raise ValueError(f"Logic error in {__file__} {__name__}") diff --git a/gpu4pyscf/gto/int3c1e_ip.py b/gpu4pyscf/gto/int3c1e_ip.py index 939096e0..63717d09 100644 --- a/gpu4pyscf/gto/int3c1e_ip.py +++ b/gpu4pyscf/gto/int3c1e_ip.py @@ -212,12 +212,6 @@ def get_int3c1e_ip2_density_contracted(mol, grids, charge_exponents, dm, intopt) charge_exponents = cp.asarray(charge_exponents, order='C') dm = cp.asarray(dm) - if dm.ndim == 3: - if dm.shape[0] > 2: - print("Warning: more than two density matrices are found for int3c1e kernel. " - "They will be summed up to one density matrix.") - dm = cp.einsum("ijk->jk", dm) - assert dm.ndim == 2 assert dm.shape[0] == dm.shape[1] and dm.shape[0] == mol.nao diff --git a/gpu4pyscf/gto/tests/test_int1e_grids.py b/gpu4pyscf/gto/tests/test_int1e_grids.py index 2ab28bf1..5401b216 100644 --- a/gpu4pyscf/gto/tests/test_int1e_grids.py +++ b/gpu4pyscf/gto/tests/test_int1e_grids.py @@ -261,6 +261,30 @@ def test_int1e_grids_charge_contracted_guassian_charge_omega(self): assert isinstance(test_int1e_dot_q, cp.ndarray) cp.testing.assert_allclose(ref_int1e_dot_q, test_int1e_dot_q, atol = charge_contraction_threshold) + # Multiple inputs + + def test_int1e_grids_density_contracted_sph_asymmetric(self): + np.random.seed(12348) + n_dm = 4 + dm = np.random.uniform(-2.0, 2.0, (n_dm, mol_sph.nao, mol_sph.nao)) + ref_int1e_dot_D = np.zeros((n_dm, grid_points.shape[0])) + for i_dm in range(n_dm): + ref_int1e_dot_D[i_dm] = np.einsum('pij,ij->p', mol_sph.intor('int1e_grids', grids=grid_points), dm[i_dm]) + test_int1e_dot_D = int1e_grids(mol_sph, grid_points, dm = dm) + assert isinstance(test_int1e_dot_D, cp.ndarray) + cp.testing.assert_allclose(ref_int1e_dot_D, test_int1e_dot_D, atol = density_contraction_threshold) + + def test_int1e_grids_charge_contracted_sph(self): + np.random.seed(12348) + n_charges = 4 + charges = np.random.uniform(-2.0, 2.0, (n_charges, grid_points.shape[0])) + ref_int1e_dot_q = np.zeros((n_charges, mol_sph.nao, mol_sph.nao)) + for i_charge in range(n_charges): + ref_int1e_dot_q[i_charge] = np.einsum('pij,p->ij', mol_sph.intor('int1e_grids', grids=grid_points), charges[i_charge]) + test_int1e_dot_q = int1e_grids(mol_sph, grid_points, charges = charges) + assert isinstance(test_int1e_dot_q, cp.ndarray) + cp.testing.assert_allclose(ref_int1e_dot_q, test_int1e_dot_q, atol = charge_contraction_threshold) + if __name__ == "__main__": print("Full Tests for One Electron Coulomb Integrals") unittest.main() diff --git a/gpu4pyscf/qmmm/chelpg.py b/gpu4pyscf/qmmm/chelpg.py index 2e658b21..c0d44af7 100644 --- a/gpu4pyscf/qmmm/chelpg.py +++ b/gpu4pyscf/qmmm/chelpg.py @@ -104,6 +104,9 @@ def tau_f(R, Rcut, Roff): r_pX_potential = 1/r_pX potential_real = cupy.dot(cupy.array(mf.mol.atom_charges()), r_pX_potential) + if dm.ndim == 3: # Unrestricted + assert dm.shape[0] == 2 + dm = dm[0] + dm[1] potential_real -= int1e_grids(mf.mol, gridcoords, dm=dm, direct_scf_tol=1e-14) w = cupy.array(w) diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index c244bb0d..a0981877 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -393,24 +393,16 @@ def _get_v(self, dms): ''' return electrostatic potential on surface ''' - nset = dms.shape[0] charge_exp = self.surface['charge_exp'] grid_coords = self.surface['grid_coords'] - ngrids = grid_coords.shape[0] - v_grids_e = cupy.empty([nset, ngrids]) - for i in range(nset): - v_grids_e[i] = int1e_grids(self.mol, grid_coords, dm = dms[i], charge_exponents = charge_exp**2, intopt = self.intopt) + v_grids_e = int1e_grids(self.mol, grid_coords, dm = dms, charge_exponents = charge_exp**2, intopt = self.intopt) return v_grids_e def _get_vmat(self, q): assert q.ndim == 2 - nset = q.shape[0] - nao = self.mol.nao charge_exp = self.surface['charge_exp'] grid_coords = self.surface['grid_coords'] - vmat = cupy.empty([nset, nao, nao]) - for i in range(nset): - vmat[i] = -int1e_grids(self.mol, grid_coords, charges = q[i], charge_exponents = charge_exp**2, intopt = self.intopt) + vmat = -int1e_grids(self.mol, grid_coords, charges = q, charge_exponents = charge_exp**2, intopt = self.intopt) return vmat def nuc_grad_method(self, grad_method):