Support multiple density matrices and charge arrays in an inefficient…

… way.

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__}")

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
 

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()

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)

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):