Analytical PCM 2nd derivative qv term

gpu4pyscf/gto/int3c1e_ip.py

@@ -103,7 +103,7 @@ def get_int3c1e_ip(mol, grids, charge_exponents, intopt):
                 int3c_angular_slice = cart2sph(int3c_angular_slice, axis=2, ang=lj)
                 int3c_angular_slice = cart2sph(int3c_angular_slice, axis=3, ang=li)
 
-            int3c_grid_slice[:, :, j0:j1, i0:i1] = int3c_angular_slice
+            int3c_grid_slice[:, :, i0:i1, j0:j1] = int3c_angular_slice.transpose(0,1,3,2)
 
         ao_idx = np.argsort(intopt._ao_idx)
         grid_idx = np.arange(ngrids_of_split)
@@ -191,14 +191,73 @@ def get_int3c1e_ip1_charge_contracted(mol, grids, charge_exponents, charges, int
             int1e_angular_slice = cart2sph(int1e_angular_slice, axis=1, ang=lj)
             int1e_angular_slice = cart2sph(int1e_angular_slice, axis=2, ang=li)
 
-        int1e_charge_contracted[:, j0:j1, i0:i1] = int1e_angular_slice
+        int1e_charge_contracted[:, i0:i1, j0:j1] = int1e_angular_slice.transpose(0,2,1)
 
     ao_idx = np.argsort(intopt._ao_idx)
     derivative_idx = np.arange(3)
     int1e_charge_contracted = int1e_charge_contracted[np.ix_(derivative_idx, ao_idx, ao_idx)]
 
     return int1e_charge_contracted
 
+def get_int3c1e_ip1_density_contracted(mol, grids, charge_exponents, dm, intopt):
+    omega = mol.omega
+    assert omega >= 0.0, "Short-range one electron integrals with GPU acceleration is not implemented."
+
+    ngrids = grids.shape[0]
+    grids = cp.asarray(grids, order='C')
+    if charge_exponents is not None:
+        charge_exponents = cp.asarray(charge_exponents, order='C')
+
+    dm = cp.asarray(dm)
+    assert dm.ndim == 2
+    assert dm.shape[0] == dm.shape[1] and dm.shape[0] == mol.nao
+
+    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
+        dm = cart2sph_transformation_matrix @ dm @ cart2sph_transformation_matrix.T
+    dm = dm.flatten(order='F') # Column major order matches (i + j * n_ao) access pattern in the C function
+
+    nao = intopt._sorted_mol.nao
+
+    i_atom_of_each_shell = intopt._sorted_mol._bas[:, ATOM_OF]
+    i_atom_of_each_shell = cp.array(i_atom_of_each_shell, dtype=np.int32)
+
+    ip1_per_atom = cp.zeros([mol.natm, 3, ngrids])
+
+    for cp_ij_id, _ in enumerate(intopt.log_qs):
+        stream = cp.cuda.get_current_stream()
+
+        log_q_ij = intopt.log_qs[cp_ij_id]
+
+        nbins = 1
+        bins_locs_ij = np.array([0, len(log_q_ij)], dtype=np.int32)
+
+        charge_exponents_pointer = c_null_ptr()
+        if charge_exponents is not None:
+            charge_exponents_pointer = charge_exponents.data.ptr
+
+        err = libgint.GINTfill_int3c1e_ip1_density_contracted(
+            ctypes.cast(stream.ptr, ctypes.c_void_p),
+            intopt.bpcache,
+            ctypes.cast(grids.data.ptr, ctypes.c_void_p),
+            ctypes.cast(charge_exponents_pointer, ctypes.c_void_p),
+            ctypes.c_int(ngrids),
+            ctypes.cast(ip1_per_atom.data.ptr, ctypes.c_void_p),
+            bins_locs_ij.ctypes.data_as(ctypes.c_void_p),
+            ctypes.c_int(nbins),
+            ctypes.c_int(cp_ij_id),
+            ctypes.cast(dm.data.ptr, ctypes.c_void_p),
+            ctypes.cast(i_atom_of_each_shell.data.ptr, ctypes.c_void_p),
+            ctypes.c_int(nao),
+            ctypes.c_double(omega))
+
+        if err != 0:
+            raise RuntimeError('GINTfill_int3c1e_charge_contracted failed')
+
+    return ip1_per_atom
+
 def get_int3c1e_ip2_density_contracted(mol, grids, charge_exponents, dm, intopt):
     omega = mol.omega
     assert omega >= 0.0, "Short-range one electron integrals with GPU acceleration is not implemented."
@@ -290,6 +349,82 @@ def get_int3c1e_ip2_density_contracted(mol, grids, charge_exponents, dm, intopt)
 
     return int3c_density_contracted
 
+def get_int3c1e_ip2_charge_contracted(mol, grids, charge_exponents, charges, gridslice, output, intopt):
+    omega = mol.omega
+    assert omega >= 0.0, "Short-range one electron integrals with GPU acceleration is not implemented."
+
+    ngrids = grids.shape[0]
+    grids = cp.asarray(grids, order='C')
+    if charge_exponents is not None:
+        charge_exponents = cp.asarray(charge_exponents, order='C')
+
+    assert charges.ndim == 1 and charges.shape[0] == grids.shape[0]
+    charges = cp.asarray(charges).astype(np.float64)
+
+    charges = charges.reshape([-1, 1], order='C')
+    grids = cp.concatenate([grids, charges], axis=1)
+
+    n_atom = len(gridslice)
+    i_atom_of_each_charge = [[i_atom] * (gridslice[i_atom][1] - gridslice[i_atom][0]) for i_atom in range(n_atom)]
+    i_atom_of_each_charge = sum(i_atom_of_each_charge, [])
+    i_atom_of_each_charge = cp.array(i_atom_of_each_charge, dtype=np.int32)
+
+    assert isinstance(output, cp.ndarray)
+    assert output.shape == (n_atom, 3, mol.nao, mol.nao)
+
+    for cp_ij_id, _ in enumerate(intopt.log_qs):
+        cpi = intopt.cp_idx[cp_ij_id]
+        cpj = intopt.cp_jdx[cp_ij_id]
+        li = intopt.angular[cpi]
+        lj = intopt.angular[cpj]
+
+        stream = cp.cuda.get_current_stream()
+
+        log_q_ij = intopt.log_qs[cp_ij_id]
+
+        nbins = 1
+        bins_locs_ij = np.array([0, len(log_q_ij)], dtype=np.int32)
+
+        i0, i1 = intopt.cart_ao_loc[cpi], intopt.cart_ao_loc[cpi+1]
+        j0, j1 = intopt.cart_ao_loc[cpj], intopt.cart_ao_loc[cpj+1]
+        ni = i1 - i0
+        nj = j1 - j0
+
+        ao_offsets = np.array([i0, j0], dtype=np.int32)
+        strides = np.array([ni, ni*nj], dtype=np.int32)
+
+        charge_exponents_pointer = c_null_ptr()
+        if charge_exponents is not None:
+            charge_exponents_pointer = charge_exponents.data.ptr
+
+        int1e_angular_slice = cp.zeros([n_atom, 3, j1-j0, i1-i0], order='C')
+
+        err = libgint.GINTfill_int3c1e_ip2_charge_contracted(
+            ctypes.cast(stream.ptr, ctypes.c_void_p),
+            intopt.bpcache,
+            ctypes.cast(grids.data.ptr, ctypes.c_void_p),
+            ctypes.cast(charge_exponents_pointer, ctypes.c_void_p),
+            ctypes.c_int(ngrids),
+            ctypes.cast(int1e_angular_slice.data.ptr, ctypes.c_void_p),
+            strides.ctypes.data_as(ctypes.c_void_p),
+            ao_offsets.ctypes.data_as(ctypes.c_void_p),
+            bins_locs_ij.ctypes.data_as(ctypes.c_void_p),
+            ctypes.c_int(nbins),
+            ctypes.c_int(cp_ij_id),
+            ctypes.cast(i_atom_of_each_charge.data.ptr, ctypes.c_void_p),
+            ctypes.c_double(omega))
+
+        if err != 0:
+            raise RuntimeError('GINTfill_int3c1e_charge_contracted failed')
+
+        i0, i1 = intopt.ao_loc[cpi], intopt.ao_loc[cpi+1]
+        j0, j1 = intopt.ao_loc[cpj], intopt.ao_loc[cpj+1]
+        if not mol.cart:
+            int1e_angular_slice = cart2sph(int1e_angular_slice, axis=2, ang=lj)
+            int1e_angular_slice = cart2sph(int1e_angular_slice, axis=3, ang=li)
+
+        output[np.ix_(range(n_atom), range(3), intopt._ao_idx[i0:i1], intopt._ao_idx[j0:j1])] += int1e_angular_slice.transpose(0,1,3,2)
+
 def get_int3c1e_ip1_charge_and_density_contracted(mol, grids, charge_exponents, dm, charges, intopt):
     dm = cp.asarray(dm)
     if dm.ndim == 3:
@@ -302,7 +437,7 @@ def get_int3c1e_ip1_charge_and_density_contracted(mol, grids, charge_exponents,
     assert dm.shape[0] == dm.shape[1] and dm.shape[0] == mol.nao
 
     int3c_ip1 = get_int3c1e_ip1_charge_contracted(mol, grids, charge_exponents, charges, intopt)
-    int3c_ip1 = cp.einsum('xji,ij->xi', int3c_ip1, dm)
+    int3c_ip1 = cp.einsum('xij,ij->xi', int3c_ip1, dm)
     return int3c_ip1
 
 def get_int3c1e_ip2_charge_and_density_contracted(mol, grids, charge_exponents, dm, charges, intopt):
@@ -319,13 +454,18 @@ def int1e_grids_ip1(mol, grids, charge_exponents=None, dm=None, charges=None, di
     $$\left(\frac{\partial}{\partial \vec{A}} \mu \middle| \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
     where $\mu(\vec{r})$ centers at $\vec{A}$ and $\nu(\vec{r})$ centers at $\vec{B}$.
 
-    If charges is not None, the function computes the following contraction:
+    If charges is not None and density is None, the function computes the following contraction:
     $$\sum_{C}^{n_{charge}} q_C \left(\frac{\partial}{\partial \vec{A}} \mu \middle| \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
     where $q_C$ is the charge centered at $\vec{C}$.
 
     If charges is not None and dm is not None, the function computes the following contraction:
     $$\sum_\nu^{n_{ao}} D_{\mu\nu} \sum_{C}^{n_{charge}} q_C
         \left(\frac{\partial}{\partial \vec{A}} \mu \middle| \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
+
+    If dm is not None and charges is None, the function computes the following contraction:
+    $$\sum_{\mu \in \{\text{AO of atom A}\}} \sum_\nu^{n_{ao}} D_{\mu\nu}
+        \left(\frac{\partial}{\partial \vec{A}} \mu \middle| \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
+    The output dimension is $(n_{atom}, 3, n_{charge})$.
     '''
     assert grids is not None
 
@@ -340,25 +480,31 @@ def int1e_grids_ip1(mol, grids, charge_exponents=None, dm=None, charges=None, di
 
     if dm is None and charges is None:
         return get_int3c1e_ip(mol, grids, charge_exponents, intopt)[0]
-    else:
-        assert charges is not None
+    elif charges is not None:
         if dm is not None:
             return get_int3c1e_ip1_charge_and_density_contracted(mol, grids, charge_exponents, dm, charges, intopt)
         else:
             return get_int3c1e_ip1_charge_contracted(mol, grids, charge_exponents, charges, intopt)
+    else:
+        assert dm is not None
+        return get_int3c1e_ip1_density_contracted(mol, grids, charge_exponents, dm, intopt)
 
 def int1e_grids_ip2(mol, grids, charge_exponents=None, dm=None, charges=None, direct_scf_tol=1e-13, intopt=None):
     r'''
     This function computes
     $$\left(\mu \middle| \frac{\partial}{\partial \vec{C}} \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
     where $\mu(\vec{r})$ centers at $\vec{A}$ and $\nu(\vec{r})$ centers at $\vec{B}$.
 
-    If dm is not None, the function computes the following contraction:
+    If dm is not None and charges is None, the function computes the following contraction:
     $$\sum_{\mu, \nu}^{n_{ao}} D_{\mu\nu} \left(\mu \middle| \frac{\partial}{\partial \vec{C}} \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
 
     If dm is not None and charges is not None, the function computes the following contraction:
     $$q_C \sum_{\mu, \nu}^{n_{ao}} D_{\mu\nu} \left(\mu \middle| \frac{\partial}{\partial \vec{C}} \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
     where $q_C$ is the charge centered at $\vec{C}$.
+
+    If charges is not None and dm is None, the function computes the following contraction:
+    $$\sum_{C}^{n_{charge}} q_C \left(\mu \middle| \frac{\partial}{\partial \vec{C}} \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
+    Notice that this summation should not be performed if the charges originates from different atomic centers.
     '''
     assert grids is not None
 
@@ -373,9 +519,36 @@ def int1e_grids_ip2(mol, grids, charge_exponents=None, dm=None, charges=None, di
 
     if dm is None and charges is None:
         return get_int3c1e_ip(mol, grids, charge_exponents, intopt)[1]
-    else:
-        assert dm is not None
+    elif dm is not None:
         if charges is not None:
             return get_int3c1e_ip2_charge_and_density_contracted(mol, grids, charge_exponents, dm, charges, intopt)
         else:
             return get_int3c1e_ip2_density_contracted(mol, grids, charge_exponents, dm, intopt)
+    else:
+        assert charges is not None
+        output = cp.zeros([1, 3, mol.nao, mol.nao])
+        get_int3c1e_ip2_charge_contracted(mol, grids, charge_exponents, charges, [[0, grids.shape[0]]], output, intopt)
+        return output.reshape([3, mol.nao, mol.nao])
+
+def int1e_grids_ip2_charge_contracted(mol, grids, charges, gridslice, output, charge_exponents=None, direct_scf_tol=1e-13, intopt=None):
+    r'''
+    This function computes the following contraction:
+    $$\sum_{C \in \{\text{grid attached to atom A}\}} q_C
+        \left(\mu \middle| \frac{\partial}{\partial \vec{C}} \frac{1}{|\vec{r} - \vec{C}|} \middle| \nu\right)$$
+    where $q_C$ is the charge centered at $\vec{C}$. The output dimension is $(n_{atom}, 3, n_{ao}, n_{ao})$.
+    '''
+    assert grids is not None
+    assert charges is not None
+    assert gridslice is not None
+    assert output is not None
+
+    if intopt is None:
+        intopt = VHFOpt(mol)
+        intopt.build(direct_scf_tol, aosym=False)
+    else:
+        assert isinstance(intopt, VHFOpt), \
+            f"Please make sure intopt is a {VHFOpt.__module__}.{VHFOpt.__name__} object."
+        assert hasattr(intopt, "density_offset"), "Please call build() function for VHFOpt object first."
+        assert not intopt.aosym
+
+    return get_int3c1e_ip2_charge_contracted(mol, grids, charge_exponents, charges, gridslice, output, intopt)