diff --git a/examples/36-amlo_eda.py b/examples/36-almo_eda.py similarity index 100% rename from examples/36-amlo_eda.py rename to examples/36-almo_eda.py diff --git a/examples/44-occri.py b/examples/44-occri.py new file mode 100644 index 000000000..5a5efb66d --- /dev/null +++ b/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) diff --git a/gpu4pyscf/pbc/df/__init__.py b/gpu4pyscf/pbc/df/__init__.py index f99f76b53..3ebe8b12b 100644 --- a/gpu4pyscf/pbc/df/__init__.py +++ b/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 diff --git a/gpu4pyscf/pbc/df/fft.py b/gpu4pyscf/pbc/df/fft.py index ddde5fc9c..6e1a79313 100644 --- a/gpu4pyscf/pbc/df/fft.py +++ b/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 + diff --git a/gpu4pyscf/pbc/df/fft_jk.py b/gpu4pyscf/pbc/df/fft_jk.py index e2a428bca..3aa8589a2 100644 --- a/gpu4pyscf/pbc/df/fft_jk.py +++ b/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.