Implementing OccRI

examples/44-occri.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python
+# Copyright 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.
+
+"""
+Using multi-grid algorithm for DFT calculation
+"""
+
+import numpy as np
+import pyscf
+
+
+from pyscf.pbc import gto
+
+cell = gto.Cell(
+    a=np.eye(3) * 3.5668,
+    atom="""C     0.      0.      0.    
+              C     0.8917  0.8917  0.8917
+              C     1.7834  1.7834  0.    
+              C     2.6751  2.6751  0.8917
+              C     1.7834  0.      1.7834
+              C     2.6751  0.8917  2.6751
+              C     0.      1.7834  1.7834
+              C     0.8917  2.6751  2.6751""",
+    basis='gth-dzvp',
+    pseudo='gth-hf-rev',
+    verbose=5,
+    exp_to_discard=0.1,
+    ke_cutoff=60,
+)
+
+kmesh = [2, 2, 2]
+kpts = cell.make_kpts(kmesh)
+exxdiv = 'ewald'
+
+from gpu4pyscf.pbc.scf import KRHF, KUHF
+from gpu4pyscf.pbc.df.fft import FFTDF, OccRI
+
+rhf = KRHF(cell, kpts)
+
+rhf.with_df = OccRI(cell, kpts)
+rhf.exxdiv = exxdiv
+rhf.conv_tol = 1e-6
+rhf.max_cycle = 50
+e_tot = rhf.kernel()
+print('-> KRHF energy with OccRI: %12.8f' % e_tot)
+
+dm0 = rhf.make_rdm1()
+
+rhf.with_df = FFTDF(cell, kpts)
+e_tot = rhf.kernel()
+print('-> KRHF energy with FFTDF: %12.8f' % e_tot)
+
+uhf = KUHF(cell, kpts)
+uhf.with_df = OccRI(cell, kpts)
+uhf.exxdiv = exxdiv
+uhf.conv_tol = 1e-6
+uhf.max_cycle = 50
+e_tot = uhf.kernel()
+print('-> KUHF energy with OccRI: %12.8f' % e_tot)
+
+dm0 = uhf.make_rdm1()
+
+uhf.with_df = FFTDF(cell, kpts)
+e_tot = uhf.kernel()
+print('-> KUHF energy with FFTDF: %12.8f' % e_tot)

gpu4pyscf/pbc/df/__init__.py

@@ -15,6 +15,6 @@
 from . import fft
 from . import aft
 from . import df
-from .fft import FFTDF
+from .fft import FFTDF, OccRI
 from .df import GDF
 from .aft import AFTDF

gpu4pyscf/pbc/df/fft.py

@@ -14,9 +14,8 @@
 
 '''GPW method'''
 
-__all__ = [
-    'get_nuc', 'get_pp', 'get_SI', 'FFTDF'
-]
+__all__ = ['get_nuc', 'get_pp', 'get_SI', 'FFTDF', 'OccRI']
+
 
 import numpy as np
 import cupy as cp
@@ -29,6 +28,7 @@
 from gpu4pyscf.lib import logger, utils
 from gpu4pyscf.lib.cupy_helper import contract
 from gpu4pyscf.pbc import tools
+from gpu4pyscf.pbc.gto import int1e
 from gpu4pyscf.pbc.df import fft_jk, aft
 from gpu4pyscf.pbc.df.aft import _check_kpts
 from gpu4pyscf.pbc.df.ft_ao import ft_ao
@@ -223,6 +223,7 @@ class FFTDF(lib.StreamObject):
     '''
 
     blockdim = 240
+    blksize = 32
 
     _keys = fft_cpu.FFTDF._keys
 
@@ -326,3 +327,59 @@ def to_cpu(self):
         out = FFTDF(self.cell, kpts=self.kpts)
         out.mesh = self.mesh
         return out
+
+class OccRI(FFTDF):
+    def __init__(self, cell, kpts):
+        super().__init__(cell, kpts)
+        self._ovlp_kpts = None
+
+    def get_jk(self, dm, hermi=1, kpts=None, kpts_band=None, with_j=True, with_k=True, omega=None, exxdiv=None):
+        assert omega is None, 'omega is not supported for OccRI'
+        if self._ovlp_kpts is None:
+            self._ovlp_kpts = int1e.int1e_ovlp(self.cell, self.kpts)
+
+        kpts, is_single_kpt = _check_kpts(kpts, dm)
+        vj = vk = None
+        if is_single_kpt:
+            if with_j:
+                vj = fft_jk.get_j_kpts(self, dm, hermi, kpts[0], kpts_band)
+
+            if with_k:
+                vk = fft_jk.get_k_occri_kpts(self, dm, hermi, kpts[0], kpts_band, exxdiv)
+
+        else:
+            if with_j:
+                vj = fft_jk.get_j_kpts(self, dm, hermi, kpts, kpts_band)
+
+            if with_k:
+                vk = fft_jk.get_k_occri_kpts(self, dm, hermi, kpts, kpts_band, exxdiv)
+
+        return vj, vk
+
+    def dump_flags(self):
+        return super().dump_flags() + f' exxdiv = {self.exxdiv}'
+
+    def _get_vR_dm(self, mo1T, mo2T, coulg, mesh):
+        nmo1 = mo1T.shape[0]
+        nmo2 = mo2T.shape[0]
+        ngrids = cp.prod(mesh)
+
+        mo1T = mo1T.reshape(nmo1, 1, ngrids)
+        mo2T = mo2T.reshape(1, nmo2, ngrids)
+
+        blksize = self.blksize
+        out = cp.zeros((nmo1, ngrids), dtype=np.complex128)
+        for i0, i1 in lib.prange(0, nmo1, blksize):
+            rhoR = mo1T[i0:i1].conj() * mo2T
+            rhoR = rhoR.reshape(-1, *mesh)
+
+            rhoG = tools.fft(rhoR, mesh)
+            vG = rhoG * coulg
+            rhoR = rhoG = None
+
+            vR = tools.ifft(vG, mesh).reshape(i1 - i0, nmo2, ngrids)
+            out[i0:i1] += contract('ijg,jg->ig', vR, mo2T[0].conj())
+            vR = vG = None
+
+        return out
+

gpu4pyscf/pbc/df/fft_jk.py

@@ -18,7 +18,7 @@
 
 __all__ = [
     'get_j_kpts', 'get_k_kpts', 'get_jk', 'get_j', 'get_k',
-    'get_j_e1_kpts', 'get_k_e1_kpts'
+    'get_j_e1_kpts', 'get_k_e1_kpts', 'get_k_occri_kpts'
 ]
 
 import numpy as np
@@ -94,7 +94,7 @@ def get_j_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None):
 
 def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None,
                exxdiv=None):
-    '''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.
+    '''Get the exchange (K) AO matrices at sampled k-points.
 
     Args:
         dm_kpts : (nkpts, nao, nao) ndarray
@@ -160,7 +160,7 @@ def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None,
         mo2_kpts = None
 
     vR_dm = cp.empty((nset,nao,ngrids), dtype=vk_kpts.dtype)
-    blksize = 32
+    blksize = mydf.blksize
 
     for k2, ao2 in enumerate(ao2_kpts):
         ao2T = ao2.T
@@ -199,6 +199,115 @@ def get_k_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None,
 
     return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
 
+
+def get_full_k(s, v, c):
+    sc = cp.dot(s, c)
+    ccs = cp.dot(c, sc.T.conj())
+    scv = cp.dot(sc, v)
+    return scv + scv.T.conj() - cp.dot(scv, ccs)
+
+
+def get_k_occri_kpts(mydf, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None, exxdiv=None):
+    '''Get the exchange (K) AO matrices at sampled k-points,
+    using OccRI.
+
+    Args:
+        dm_kpts : (nkpts, nao, nao) ndarray
+            Density matrix at each k-point
+        kpts : (nkpts, 3) ndarray
+
+    Kwargs:
+        hermi : int
+            Whether K matrix is hermitian
+
+            | 0 : not hermitian and not symmetric
+            | 1 : hermitian
+
+        kpts_band : (3,) ndarray or (*,3) ndarray
+            A list of arbitrary "band" k-points at which to evalute the matrix.
+
+    Returns:
+        vk : (nkpts, nao, nao) ndarray
+        or list of vk if the input dm_kpts is a list of DMs
+    '''
+    cell = mydf.cell
+    mesh = mydf.mesh
+    assert cell.low_dim_ft_type != 'inf_vacuum'
+    assert cell.dimension > 1
+    coords = mydf.grids.coords
+    ngrids = coords.shape[0]
+
+    mo_coeff = getattr(dm_kpts, 'mo_coeff', None)
+    mo_occ = getattr(dm_kpts, 'mo_occ', None)
+
+    if (kpts_band is not None) or (mo_coeff is None):
+        return get_k_kpts(mydf, dm_kpts, hermi, kpts, kpts_band, exxdiv)
+
+    kpts = np.asarray(kpts)
+    dm_kpts = cp.asarray(dm_kpts, order='C')
+    dms = _format_dms(dm_kpts, kpts)
+    nset, nkpts, nao = dms.shape[:3]
+
+    weight = cell.vol / ngrids / nkpts
+
+    kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
+    nband = len(kpts_band)
+
+    dtype = dms.dtype if (is_zero(kpts_band) and is_zero(kpts)) else np.complex128
+    vk_kpts = cp.zeros((nset, nband, nao, nao), dtype=dtype)
+
+    mo_coeff = cp.asarray(mo_coeff).reshape(nset, nkpts, nao, -1)
+    mo_occ = cp.asarray(mo_occ).reshape(nset, nkpts, -1)
+
+    from gpu4pyscf.pbc.dft.numint import eval_ao
+
+    for s in range(nset):
+        cocc_kpts = []
+        nocc_kpts = []
+        for k in range(nkpts):
+            mask = mo_occ[s, k] > 0
+            cocc_kpts.append(mo_coeff[s, k][:, mask])
+            nocc_kpts.append(mo_occ[s, k][mask] ** 0.5)
+
+        for k1 in range(nkpts):
+            kpt1 = np.array(kpts[k1])
+            ao1 = eval_ao(cell, coords, kpt=kpt1)
+
+            cocc1 = cocc_kpts[k1]
+            mo1 = cp.dot(ao1, cocc1)
+            nmo1 = mo1.shape[1]
+            mo1T = mo1.T
+
+            vk_k1 = cp.zeros((nmo1, nao), dtype=dtype)
+            for k2 in range(nkpts):
+                kpt2 = np.array(kpts[k2])
+                ao2 = eval_ao(cell, coords, kpt=kpt2)
+
+                k21 = kpt2 - kpt1
+                coulg = tools.get_coulG(cell, k21, exxdiv, mesh, kpts=kpts)
+
+                if not is_zero(k21):
+                    k21 = cp.asarray(k21)
+                    theta = cp.dot(coords, k21)
+                    phase = cp.exp(-1j * theta)
+                    ao2 = ao2 * phase.reshape(-1, 1)
+
+                cocc2 = cocc_kpts[k2]
+                mo2 = cp.dot(ao2, cocc2 * nocc_kpts[k2])
+                mo2T = mo2.T
+
+                vR_dm = mydf._get_vR_dm(mo1T, mo2T, coulg, mesh)
+                if vk_kpts.dtype == np.double:
+                    vR_dm = vR_dm.real
+
+                vk_k1 += cp.dot(vR_dm, ao1) * weight
+                vR_dm = None
+
+            ovlp1 = mydf._ovlp_kpts[k1]
+            vk_kpts[s, k1] = get_full_k(ovlp1, vk_k1, cocc1)
+    return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
+
+
 def get_jk(mydf, dm, hermi=1, kpt=np.zeros(3), kpts_band=None,
            with_j=True, with_k=True, exxdiv=None):
     '''Get the Coulomb (J) and exchange (K) AO matrices for the given density matrix.