diff --git a/examples/post_HF/21-df-casscf.py b/examples/post_HF/21-df-casscf.py new file mode 100644 index 000000000..0b9b43efe --- /dev/null +++ b/examples/post_HF/21-df-casscf.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python +# Copyright 2026 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. + +######################################## +# Example of GPU DF-CASSCF +######################################## + +import pyscf +from pyscf import mcscf as cpu_mcscf + +from gpu4pyscf import mcscf as gpu_mcscf + + +mol = pyscf.M( + atom='N 0 0 -0.7; N 0 0 0.7', + basis='6-31g*', +) +mf = mol.RHF().to_gpu().density_fit().run() + +mc_cpu = cpu_mcscf.DFCASSCF(mf.to_cpu(), 6, 6) +e_cpu = mc_cpu.kernel()[0] + +mc_gpu = gpu_mcscf.DFCASSCF(mf.to_gpu(), 6, 6) +e_gpu = mc_gpu.kernel()[0] + +print(f'CPU DF-CASSCF energy: {e_cpu:.12f}') +print(f'GPU DF-CASSCF energy: {e_gpu:.12f}') +print(f'CPU-GPU difference: {e_gpu - e_cpu:.3e}') diff --git a/gpu4pyscf/df/df.py b/gpu4pyscf/df/df.py index e112d7cf0..7b23e124e 100644 --- a/gpu4pyscf/df/df.py +++ b/gpu4pyscf/df/df.py @@ -20,9 +20,11 @@ import cupy as cp from cupyx.scipy.linalg import solve_triangular from pyscf import lib +from pyscf.ao2mo.incore import iden_coeffs from pyscf.df import addons, incore from gpu4pyscf.lib.cupy_helper import ( - cholesky, get_avail_mem, fill_symmetric, asarray, empty_mapped, ndarray) + cholesky, contract, get_avail_mem, fill_symmetric, asarray, empty_mapped, + ndarray) from gpu4pyscf.lib import logger from gpu4pyscf.lib import utils from gpu4pyscf.lib import multi_gpu @@ -282,6 +284,61 @@ def loop(self, blksize=None, unpack=True): out=work[:,:,:p1-p0]) yield out.transpose(2,0,1), cderi_blk + def ao2mo(self, mo_coeffs, compact=True): + '''Transform density-fitting integrals to MO integrals on GPU.''' + if isinstance(mo_coeffs, (np.ndarray, cupy.ndarray)): + if mo_coeffs.ndim != 2: + raise ValueError('mo_coeffs must be a 2D array or four 2D arrays') + mo_coeffs = (mo_coeffs,) * 4 + elif len(mo_coeffs) != 4: + raise ValueError('mo_coeffs must contain four coefficient matrices') + else: + mo_coeffs = tuple(mo_coeffs) + + if any(mo.ndim != 2 or mo.shape[0] != self.mol.nao + for mo in mo_coeffs): + raise ValueError('MO coefficients must have shape (mol.nao, nmo)') + if any(cupy.iscomplexobj(mo) for mo in mo_coeffs): + raise NotImplementedError('GPU DF AO2MO does not support complex orbitals') + + ij_compact = compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]) + kl_compact = compact and iden_coeffs(mo_coeffs[2], mo_coeffs[3]) + same_pairs = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and + iden_coeffs(mo_coeffs[1], mo_coeffs[3])) + + if self._cderi is None: + self.build() + mo_coeffs = [cupy.asarray(mo, dtype=cupy.float64) + for mo in mo_coeffs] + ni, nj, nk, nl = [mo.shape[1] for mo in mo_coeffs] + nij = ni * (ni + 1) // 2 if ij_compact else ni * nj + nkl = nk * (nk + 1) // 2 if kl_compact else nk * nl + + def transform(): + coeffs = [asarray(mo) for mo in mo_coeffs] + eri = cupy.zeros((nij, nkl)) + + def transform_pair(cderi, left, right, compact): + buf = contract('Luv,up->Lpv', cderi, left) + out = contract('Lpv,vq->Lpq', buf, right) + if compact: + idx = cp.tril_indices(left.shape[1]) + return out[:, idx[0], idx[1]] + return out.reshape(len(cderi), -1) + + for cderi, _ in self.loop(unpack=True): + lij = transform_pair( + cderi, coeffs[0], coeffs[1], ij_compact) + lkl = lij if same_pairs else transform_pair( + cderi, coeffs[2], coeffs[3], kl_compact) + contract('Li,Lj->ij', lij, lkl, beta=1, out=eri) + return eri + + eri = multi_gpu.run(transform, non_blocking=True) + return multi_gpu.array_reduce(eri, inplace=True) + + get_mo_eri = ao2mo + def reset(self, mol=None): '''Reset mol and clean up relevant attributes for scanner mode''' if mol is not None: @@ -336,7 +393,6 @@ def range_coulomb(self, omega): auxmol.omega = auxmol_omega get_ao_eri = get_eri = NotImplemented - get_mo_eri = ao2mo = NotImplemented def cholesky_eri_gpu(intopt, mol, auxmol, cd_low, omega=None, sr_only=False, use_gpu_memory=True): diff --git a/gpu4pyscf/df/tests/test_df_ao2mo.py b/gpu4pyscf/df/tests/test_df_ao2mo.py new file mode 100644 index 000000000..6f37dc17c --- /dev/null +++ b/gpu4pyscf/df/tests/test_df_ao2mo.py @@ -0,0 +1,63 @@ +# Copyright 2026 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. + +import unittest + +import pyscf +from pyscf import df as cpu_df +from pyscf import scf + +from gpu4pyscf import df + + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.mol = pyscf.M( + atom='O 0 0 0; H 0 -0.757 0.587; H 0 0.757 0.587', + basis='cc-pvdz', verbose=0, output='/dev/null') + cls.mo = scf.RHF(cls.mol).run(conv_tol=1e-12).mo_coeff + cls.cpu_df = cpu_df.DF(cls.mol, auxbasis='weigend').build() + cls.gpu_df = df.DF(cls.mol, auxbasis='weigend').build() + + @classmethod + def tearDownClass(cls): + cls.mol.stdout.close() + + def test_compact(self): + ref = self.cpu_df.ao2mo(self.mo) + out = self.gpu_df.ao2mo(self.mo).get() + self.assertEqual(out.shape, ref.shape) + self.assertLess(abs(out - ref).max(), 1e-8) + + def test_general(self): + coeffs = (self.mo[:, :3], self.mo[:, 1:5], + self.mo[:, 2:6], self.mo[:, :2]) + ref = self.cpu_df.ao2mo(coeffs, compact=False) + out = self.gpu_df.ao2mo(coeffs, compact=False).get() + self.assertEqual(out.shape, ref.shape) + self.assertLess(abs(out - ref).max(), 1e-8) + + def test_host_cderi(self): + with_df = df.DF(self.mol, auxbasis='weigend') + with_df.use_gpu_memory = False + with_df.build() + ref = self.cpu_df.ao2mo(self.mo[:, 2:9]) + out = with_df.ao2mo(self.mo[:, 2:9]).get() + self.assertLess(abs(out - ref).max(), 1e-8) + + +if __name__ == '__main__': + print('Full tests for GPU DF AO2MO') + unittest.main() diff --git a/gpu4pyscf/fci/__init__.py b/gpu4pyscf/fci/__init__.py index e69de29bb..ceaef390b 100644 --- a/gpu4pyscf/fci/__init__.py +++ b/gpu4pyscf/fci/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2026 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. + +from gpu4pyscf.fci.direct_spin1 import FCI +from gpu4pyscf.fci.direct_spin1 import FCISolver + + +def solver(mol=None, singlet=False, symm=None): + if singlet: + raise NotImplementedError('GPU spin-adapted FCI is not implemented') + if symm or (symm is None and mol is not None and mol.symmetry): + raise NotImplementedError('GPU symmetry-adapted FCI is not implemented') + return FCISolver(mol) diff --git a/gpu4pyscf/fci/direct_spin1.py b/gpu4pyscf/fci/direct_spin1.py index 3a7896f56..277cb8422 100644 --- a/gpu4pyscf/fci/direct_spin1.py +++ b/gpu4pyscf/fci/direct_spin1.py @@ -1,4 +1,4 @@ -# Copyright 2021-2024 The PySCF Developers. All Rights Reserved. +# Copyright 2026 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. @@ -14,58 +14,82 @@ import numpy as np import cupy as cp + +from pyscf import lib +from pyscf.fci import cistring from pyscf.fci import direct_spin1 +from gpu4pyscf.lib import logger +from gpu4pyscf.lib.cupy_helper import get_avail_mem + + TILE = 32 code = r''' #define TILE 32 extern "C" { __global__ -void _build_t1(double *ci0, double *t1, +void _build_t1(const double *ci0, double *t1, long long strb0, long long na, long long nb, long long nnorb, - unsigned short *addra, unsigned short *addrb, char *signa, char *signb) + const unsigned int *addra, const unsigned int *addrb, + const signed char *signa, const signed char *signb) { + long long vec = blockIdx.z; + ci0 += vec * na * nb; + t1 += vec * nnorb * na * TILE; int tx = threadIdx.x; int ty = threadIdx.y; - int stra0 = blockIdx.y * blockDim.y; - int strb = strb0 + tx; - int stra = stra0 + ty; - - int nab = na * TILE; - int ab_id = stra * TILE + tx; - __shared__ unsigned short _addra[TILE*TILE]; - __shared__ unsigned short _addrb[TILE*TILE]; - __shared__ char _signa[TILE*TILE]; - __shared__ char _signb[TILE*TILE]; - int sign, str1, j0, j; - int dj = TILE; - double val; - - for (j0 = 0; j0 < nnorb; j0+=TILE) { - _addra[ty*TILE+tx] = addra[(j0+ty)*na+stra0+tx]; - _addrb[ty*TILE+tx] = addrb[(j0+ty)*nb+strb0+tx]; - _signa[ty*TILE+tx] = signa[(j0+ty)*na+stra0+tx]; - _signb[ty*TILE+tx] = signb[(j0+ty)*nb+strb0+tx]; - if (j0 + TILE > nnorb) { - dj = nnorb - j0; + long long stra0 = (long long)blockIdx.y * blockDim.y; + long long stra = stra0 + ty; + long long strb = strb0 + tx; + __shared__ unsigned int _addra[TILE*TILE]; + __shared__ unsigned int _addrb[TILE*TILE]; + __shared__ signed char _signa[TILE*TILE]; + __shared__ signed char _signb[TILE*TILE]; + + for (long long j0 = 0; j0 < nnorb; j0 += TILE) { + long long ja = j0 + ty; + long long ia = stra0 + tx; + long long jb = j0 + ty; + long long ib = strb0 + tx; + int tile_idx = ty * TILE + tx; + if (ja < nnorb && ia < na) { + long long link_idx = ja * na + ia; + _addra[tile_idx] = addra[link_idx]; + _signa[tile_idx] = signa[link_idx]; + } else { + _addra[tile_idx] = 0; + _signa[tile_idx] = 0; + } + if (jb < nnorb && ib < nb) { + long long link_idx = jb * nb + ib; + _addrb[tile_idx] = addrb[link_idx]; + _signb[tile_idx] = signb[link_idx]; + } else { + _addrb[tile_idx] = 0; + _signb[tile_idx] = 0; } __syncthreads(); - if (stra < na && strb < nb) { - for (j = 0; j < dj; j++) { - val = 0; - sign = _signa[j*TILE+ty]; - str1 = _addra[j*TILE+ty]; - if (sign != 0) { - val = sign * ci0[str1*nb+strb]; - } - sign = _signb[j*TILE+tx]; - str1 = _addrb[j*TILE+tx]; - if (sign != 0) { - val += sign * ci0[stra*nb+str1]; + if (stra < na) { + int dj = min((long long)TILE, nnorb - j0); + for (int j = 0; j < dj; j++) { + double val = 0.; + if (strb < nb) { + int sign = _signa[j*TILE+ty]; + unsigned int str1 = _addra[j*TILE+ty]; + if (sign != 0) { + val = sign * ci0[(long long)str1 * nb + strb]; + } + + sign = _signb[j*TILE+tx]; + str1 = _addrb[j*TILE+tx]; + if (sign != 0) { + val += sign * ci0[stra * nb + str1]; + } } - t1[(j0+j)*nab + ab_id] = val; + long long t1_idx = ((j0+j) * na + stra) * TILE + tx; + t1[t1_idx] = val; } } __syncthreads(); @@ -73,53 +97,70 @@ } __global__ -void _gather(double *out, double *t1, +void _gather(const double *t1, double *out, long long strb0, long long na, long long nb, long long nnorb, - unsigned short *addra, unsigned short *addrb, char *signa, char *signb) + const unsigned int *addra, const unsigned int *addrb, + const signed char *signa, const signed char *signb) { + long long vec = blockIdx.z; + t1 += vec * nnorb * na * TILE; + out += vec * na * nb; int tx = threadIdx.x; int ty = threadIdx.y; - int stra0 = blockIdx.y * blockDim.y; - int strb = strb0 + tx; - int stra = stra0 + ty; - int nab = na * TILE; - int ab_id = stra * TILE + tx; - __shared__ unsigned short _addra[TILE*TILE]; - __shared__ unsigned short _addrb[TILE*TILE]; - __shared__ char _signa[TILE*TILE]; - __shared__ char _signb[TILE*TILE]; - int sign, str1, j0, j; - int dj = TILE; + long long stra0 = (long long)blockIdx.y * blockDim.y; + long long stra = stra0 + ty; + long long strb = strb0 + tx; + __shared__ unsigned int _addra[TILE*TILE]; + __shared__ unsigned int _addrb[TILE*TILE]; + __shared__ signed char _signa[TILE*TILE]; + __shared__ signed char _signb[TILE*TILE]; double val = 0.; - for (j0 = 0; j0 < nnorb; j0+=TILE) { - _addra[ty*TILE+tx] = addra[(j0+ty)*na+stra0+tx]; - _addrb[ty*TILE+tx] = addrb[(j0+ty)*nb+strb0+tx]; - _signa[ty*TILE+tx] = signa[(j0+ty)*na+stra0+tx]; - _signb[ty*TILE+tx] = signb[(j0+ty)*nb+strb0+tx]; - if (j0 + TILE > nnorb) { - dj = nnorb - j0; + for (long long j0 = 0; j0 < nnorb; j0 += TILE) { + long long ja = j0 + ty; + long long ia = stra0 + tx; + long long jb = j0 + ty; + long long ib = strb0 + tx; + int tile_idx = ty * TILE + tx; + if (ja < nnorb && ia < na) { + long long link_idx = ja * na + ia; + _addra[tile_idx] = addra[link_idx]; + _signa[tile_idx] = signa[link_idx]; + } else { + _addra[tile_idx] = 0; + _signa[tile_idx] = 0; + } + if (jb < nnorb && ib < nb) { + long long link_idx = jb * nb + ib; + _addrb[tile_idx] = addrb[link_idx]; + _signb[tile_idx] = signb[link_idx]; + } else { + _addrb[tile_idx] = 0; + _signb[tile_idx] = 0; } __syncthreads(); + if (stra < na && strb < nb) { - for (j = 0; j < dj; j++) { - sign = _signa[j*TILE+ty]; - str1 = _addra[j*TILE+ty]; + int dj = min((long long)TILE, nnorb - j0); + for (int j = 0; j < dj; j++) { + int sign = _signa[j*TILE+ty]; + unsigned int str1 = _addra[j*TILE+ty]; if (sign != 0) { - val += sign * t1[(j0+j)*nab + (str1*TILE+tx)]; + val += sign * t1[((j0+j) * na + str1) * TILE + tx]; } sign = _signb[j*TILE+tx]; str1 = _addrb[j*TILE+tx]; if (sign != 0) { - out[stra*nb+str1] += sign * t1[(j0+j)*nab + ab_id]; + atomicAdd(out + stra * nb + str1, + sign * t1[((j0+j) * na + stra) * TILE + tx]); } } } __syncthreads(); } if (stra < na && strb < nb) { - out[stra*nb+strb] += val; + atomicAdd(out + stra * nb + strb, val); } } }''' @@ -128,65 +169,413 @@ _build_t1 = _contract_2e_spin1.get_function('_build_t1') _gather = _contract_2e_spin1.get_function('_gather') + def _link_index_to_addrs(link_index, nnorb): - na, nov = link_index.shape[:2] - ia = link_index[:,:,0].T - addr = np.zeros((nnorb, na), dtype=np.uint16) - sign = np.zeros((nnorb, na), dtype=np.int8) - #:for j in range(nov): - #: for a in range(na): - #: addr[ia[a,j],a] = link_index[a,j,2] - #: sign[ia[a,j],a] = link_index[a,j,3] - idx = np.arange(na) - addr[ia,idx] = link_index[:,:,2].T - sign[ia,idx] = link_index[:,:,3].T - # Add paddings to avoid illegal address in kernel - _addr = cp.empty((nnorb+TILE-1, na), dtype=np.uint16)[:nnorb] - _sign = cp.empty((nnorb+TILE-1, na), dtype=np.int8)[:nnorb] - _addr.set(addr) - _sign.set(sign) - return _addr, _sign - -def contract_2e(eri, ci0, norb, nelec, link_index): - ci0 = cp.asarray(ci0) - out = cp.zeros_like(ci0) - na, nb = ci0.shape + nstr = link_index.shape[0] + pair = link_index[:, :, 0].T + addr = np.zeros((nnorb, nstr), dtype=np.uint32) + sign = np.zeros((nnorb, nstr), dtype=np.int8) + idx = np.arange(nstr) + addr[pair, idx] = link_index[:, :, 2].T + sign[pair, idx] = link_index[:, :, 3].T + return cp.asarray(addr), cp.asarray(sign) + + +def _prepare_link_index(link_index, norb): + if len(link_index) == 4 and isinstance(link_index[0], cp.ndarray): + return link_index + nnorb = norb * (norb + 1) // 2 - assert eri.shape == (nnorb, nnorb) - eri = cp.asarray(eri) link_indexa, link_indexb = link_index addra, signa = _link_index_to_addrs(link_indexa, nnorb) if link_indexa is link_indexb: addrb, signb = addra, signa else: addrb, signb = _link_index_to_addrs(link_indexb, nnorb) + return addra, signa, addrb, signb + + +def _prepare_rdm_link_index(link_index, norb, nelec): + if link_index is None: + neleca, nelecb = direct_spin1._unpack_nelec(nelec) + link_indexa = cistring.gen_linkstr_index(range(norb), neleca) + if neleca == nelecb: + link_indexb = link_indexa + else: + link_indexb = cistring.gen_linkstr_index(range(norb), nelecb) + else: + link_indexa, link_indexb = link_index + + nnorb = norb * norb + + def convert(link): + full_link = link.copy() + full_link[:, :, 0] = link[:, :, 1] * norb + link[:, :, 0] + return _link_index_to_addrs(full_link, nnorb) + + addra, signa = convert(link_indexa) + if link_indexa is link_indexb: + addrb, signb = addra, signa + else: + addrb, signb = convert(link_indexb) + return addra, signa, addrb, signb + + +def contract_2e(eri, ci0, norb, nelec, link_index=None): + nelec = direct_spin1._unpack_nelec(nelec) + if link_index is None: + link_index = direct_spin1._unpack(norb, nelec, None) + gpu_link_index = _prepare_link_index(link_index, norb) + addra, signa, addrb, signb = gpu_link_index + + neleca, nelecb = nelec + na = lib.comb(norb, neleca) + nb = lib.comb(norb, nelecb) + ci0 = cp.asarray(ci0, dtype=cp.float64, order='C') + input_shape = ci0.shape + single_vector = ci0.ndim == 1 or ci0.shape == (na, nb) + if single_vector: + nvec = 1 + elif ci0.ndim in (2, 3): + nvec = ci0.shape[0] + else: + raise ValueError(f'Invalid CI shape {ci0.shape}') + if ci0.size != nvec * na * nb: + raise ValueError( + f'Invalid CI size {ci0.size}; expected {nvec * na * nb}') + ci0 = ci0.reshape(nvec, na, nb) + + nnorb = norb * (norb + 1) // 2 + eri = cp.asarray(eri, dtype=cp.float64, order='C') + if eri.shape != (nnorb, nnorb): + raise ValueError(f'Invalid ERI shape {eri.shape}; expected {(nnorb, nnorb)}') + + out = cp.zeros_like(ci0) + t1 = cp.empty((nvec, nnorb, na * TILE), dtype=cp.float64) + gt1 = cp.empty_like(t1) + threads = (TILE, TILE) + blocks = (1, (na + threads[1] - 1) // threads[1], nvec) + rest_args = (na, nb, nnorb, addra, addrb, signa, signb) + for strb0 in range(0, nb, TILE): + _build_t1(blocks, threads, (ci0, t1, strb0) + rest_args) + cp.matmul(eri, t1, out=gt1) + _gather(blocks, threads, (gt1, out, strb0) + rest_args) + return out[0] if single_vector else out.reshape(input_shape) + + +def make_rdm12(fcivec, norb, nelec, link_index=None, reorder=True): + nelec = direct_spin1._unpack_nelec(nelec) + neleca, nelecb = nelec + na = lib.comb(norb, neleca) + nb = lib.comb(norb, nelecb) + fcivec = cp.asarray(fcivec, dtype=cp.float64, order='C') + if fcivec.size != na * nb: + raise ValueError(f'Invalid CI size {fcivec.size}; expected {na * nb}') + fcivec = fcivec.reshape(na, nb) + gpu_link_index = _prepare_rdm_link_index(link_index, norb, nelec) + addra, signa, addrb, signb = gpu_link_index + nnorb = norb * norb + dm1 = cp.zeros(nnorb, dtype=cp.float64) + dm2 = cp.zeros((nnorb, nnorb), dtype=cp.float64) + t1 = cp.empty((nnorb, na * TILE), dtype=cp.float64) + ci_tile = cp.zeros((na, TILE), dtype=cp.float64) threads = (TILE, TILE) - blocks = (1, (na+TILE-1)//TILE) + blocks = (1, (na + threads[1] - 1) // threads[1]) rest_args = (na, nb, nnorb, addra, addrb, signa, signb) - t1 = cp.empty((nnorb, na*TILE)) - gt1 = cp.empty((nnorb, na*TILE)) - if 0: - # Pipeline the three operations. Seems no performance improvement - buf = cp.empty((nnorb, na*TILE)) - sm_t1 = cp.cuda.Stream(non_blocking=True) - sm_gt1 = cp.cuda.Stream(non_blocking=True) - for strb0 in range(0, nb, TILE): - _build_t1(blocks, threads, (ci0, t1, strb0) + rest_args, stream=sm_t1) - sm_t1.synchronize() - with sm_gt1: - eri.dot(t1, out=gt1) - sm_gt1.synchronize() - _gather(blocks, threads, (out, gt1, strb0) + rest_args) - t1, gt1, buf = buf, t1, gt1 + + for strb0 in range(0, nb, TILE): + _build_t1(blocks, threads, (fcivec, t1, strb0) + rest_args) + blen = min(TILE, nb - strb0) + ci_tile.fill(0.) + ci_tile[:, :blen] = fcivec[:, strb0:strb0 + blen] + dm1 += cp.dot(t1, ci_tile.ravel()) + dm2 += cp.dot(t1, t1.T) + + dm1 = dm1.reshape(norb, norb).T + dm2 = dm2.reshape(norb, norb, norb, norb).transpose(1, 0, 2, 3) + if reorder: + for k in range(norb): + dm2[:, k, k, :] -= dm1.T + dm2 = dm2.reshape(nnorb, nnorb) + dm2 = (dm2 + dm2.T) * .5 + dm2 = dm2.reshape(norb, norb, norb, norb) + return dm1, dm2 + + +def _qr(vectors, lindep): + vectors = cp.asarray(vectors).copy() + nvec = 0 + for vector in vectors: + for _ in range(2): + if nvec: + vector -= vectors[:nvec].conj().dot(vector).dot( + vectors[:nvec]) + norm = cp.linalg.norm(vector) + if norm**2 > lindep: + vectors[nvec] = vector / norm + nvec += 1 + return vectors[:nvec] + + +def davidson1(aop, x0, precond, tol=1e-12, max_cycle=50, max_space=12, + lindep=1e-14, max_memory=4000, nroots=1, pick=None, + verbose=logger.WARN, tol_residual=None): + """In-core GPU Davidson solver. Host-memory and disk spill are unsupported.""" + if isinstance(verbose, logger.Logger): + log = verbose else: - for strb0 in range(0, nb, TILE): - _build_t1(blocks, threads, (ci0, t1, strb0) + rest_args) - eri.dot(t1, out=gt1) - _gather(blocks, threads, (out, gt1, strb0) + rest_args) - return out.get() + log = logger.Logger(verbose=verbose) + + if tol_residual is None: + tol_residual = np.sqrt(tol) + + if isinstance(x0, (list, tuple)): + x0 = cp.stack([cp.asarray(x) for x in x0]) + else: + x0 = cp.asarray(x0) + if x0.ndim == 1: + x0 = x0[None] + vector_size = x0.shape[1] + if not 0 < nroots <= vector_size: + raise ValueError(f'Invalid nroots {nroots} for vector size {vector_size}') + + vector_bytes = vector_size * x0.dtype.itemsize + max_space = min(max_space + (nroots - 1) * 4, vector_size) + work_vectors = max(5 * nroots, len(x0)) + required_memory = ((2 * max_space + work_vectors) * vector_bytes + + max_space**2 * x0.dtype.itemsize) + available_memory = get_avail_mem() + memory_limit = min(max_memory * 1e6, available_memory) + if required_memory > memory_limit: + raise MemoryError( + f'GPU Davidson subspace requires ' + f'{required_memory / 1e6:.0f} MB in-core; ' + f'max_memory={max_memory:.0f} MB, ' + f'available GPU memory={available_memory / 1e6:.0f} MB') + log.debug('Davidson max_space %d, in-core memory %.0f MB', max_space, + required_memory / 1e6) + + try: + xs = cp.empty((max_space, vector_size), dtype=x0.dtype) + ax = cp.empty_like(xs) + heff = cp.empty((max_space, max_space), dtype=x0.dtype) + except cp.cuda.memory.OutOfMemoryError as err: + raise MemoryError( + f'Failed to allocate GPU Davidson subspace of size {max_space}') \ + from err + + converged = cp.zeros(nroots, dtype=bool) + energy = previous_energy = None + ritz = None + trial = _qr(x0, lindep) + x0 = None + space = 0 + + for cycle in range(max(1, max_cycle)): + if len(trial) == 0: + break + if space + len(trial) > max_space: + trial = trial[:max_space-space] + + try: + atrial = aop(trial) + except cp.cuda.memory.OutOfMemoryError as err: + raise MemoryError( + 'Insufficient GPU memory for the Davidson Hamiltonian-vector ' + 'product') from err + old_space = space + space += len(trial) + xs[old_space:space] = trial + ax[old_space:space] = atrial + + hsub = trial.conj().dot(ax[:space].T) + heff[old_space:space, :space] = hsub + heff[:old_space, old_space:space] = hsub[:, :old_space].T.conj() + block = hsub[:, old_space:space] + heff[old_space:space, old_space:space] = ( + block + block.T.conj()) * .5 + + w, v = cp.linalg.eigh(heff[:space, :space]) + if pick is not None: + w, v, _ = pick(w, v, nroots, locals()) + if len(w) < nroots: + raise RuntimeError('Not enough eigenvalues') + + previous_energy, energy = energy, w[:nroots] + coeff = v[:, :nroots].T + ritz = coeff.dot(xs[:space]) + aritz = coeff.dot(ax[:space]) + residual = aritz - energy[:, None] * ritz + residual_norm = cp.linalg.norm(residual, axis=1) + if previous_energy is None: + de = cp.full_like(energy, cp.inf) + else: + de = energy - previous_energy + converged = ((cp.abs(de) < tol) & + (residual_norm < tol_residual)) + + max_residual = residual_norm.max() + max_de = cp.abs(de).max() + if bool(cp.all(converged).get()): + log.debug('converged %d %d |r|= %.3g e= %s max|de|= %.3g', + cycle, space, max_residual, energy, max_de) + break + + active = cp.where( + ~converged & (residual_norm**2 > lindep))[0] + if len(active) == 0: + converged = residual_norm < tol_residual + break + + trial = precond(residual[active], energy[active]) + for _ in range(2): + trial -= xs[:space].conj().dot(trial.T).T.dot(xs[:space]) + trial = _qr(trial, lindep) + log.debug('davidson %d %d |r|= %.3g e= %s max|de|= %.3g', + cycle, space, max_residual, energy, max_de) + if len(trial) == 0: + converged = residual_norm < tol_residual + break -class FCI(direct_spin1.FCI): + if space + len(trial) > max_space: + trial = _qr(ritz, lindep) + space = 0 + + if ritz is None: + raise RuntimeError('Davidson failed to build a trial subspace') + return converged, cp.asnumpy(energy), ritz + + +class FCISolver(direct_spin1.FCISolver): from gpu4pyscf.lib.utils import to_cpu, to_gpu, device contract_2e = staticmethod(contract_2e) + + def kernel(self, h1e, eri, norb, nelec, ci0=None, + tol=None, lindep=None, max_cycle=None, max_space=None, + nroots=None, max_memory=None, verbose=None, ecore=None, **kwargs): + if nroots is None: + nroots = self.nroots + if tol is None: + tol = self.conv_tol + if lindep is None: + lindep = self.lindep + if max_cycle is None: + max_cycle = self.max_cycle + if max_space is None: + max_space = self.max_space + if ecore is None: + ecore = 0 + if isinstance(verbose, lib.logger.Logger): + log = logger.Logger(verbose.stdout, verbose.verbose) + else: + log = logger.new_logger(self, verbose) + fci_t0 = log.init_timer() + nelec = direct_spin1._unpack_nelec(nelec, self.spin) + h1e = cp.asnumpy(h1e) if isinstance(h1e, cp.ndarray) else np.asarray(h1e) + eri = cp.asnumpy(eri) if isinstance(eri, cp.ndarray) else np.asarray(eri) + hdiag = self.make_hdiag(h1e, eri, norb, nelec, compress=False).ravel() + link_index = direct_spin1._unpack(norb, nelec, None) + gpu_link_index = _prepare_link_index(link_index, norb) + h2e = cp.asarray(self.absorb_h1e(h1e, eri, norb, nelec, .5)) + hdiag_gpu = cp.asarray(hdiag) + + if ci0 is None: + ci0 = self.get_init_guess(norb, nelec, nroots, hdiag) + elif callable(ci0): + ci0 = ci0() + + hop_calls = 0 + hop_vectors = 0 + + def hop(cis): + nonlocal hop_calls, hop_vectors + t0 = log.init_timer() + out = contract_2e( + h2e, cis, norb, nelec, gpu_link_index).reshape(len(cis), -1) + hop_calls += 1 + hop_vectors += len(cis) + log.timer_debug1( + f'contract_2e for {len(cis)} CI vectors', *t0) + return out + + def precond(residual, energy): + denominator = hdiag_gpu - ( + energy[:, None] - self.level_shift) + denominator = cp.where( + cp.abs(denominator) < 1e-8, 1e-8, denominator) + return residual / denominator + + tol_residual = getattr(self, 'conv_tol_residual', None) + if tol_residual is None: + tol_residual = np.sqrt(tol) + if isinstance(ci0, (list, tuple)): + ci0 = cp.stack([cp.asarray(x, dtype=cp.float64) for x in ci0]) + else: + ci0 = cp.asarray(ci0, dtype=cp.float64) + ci0 = ci0.reshape(-1, hdiag.size) + if max_memory is None: + max_memory = self.max_memory + + setup_t1 = log.timer('FCI setup', *fci_t0) + setup_wall = setup_t1[1] - fci_t0[1] + davidson_t0 = log.init_timer() + converged, energies, ci = davidson1( + hop, ci0, precond, tol=tol, tol_residual=tol_residual, + lindep=lindep, nroots=nroots, max_cycle=max_cycle, + max_space=max_space, + max_memory=max_memory, verbose=log) + davidson_t1 = log.timer('FCI Davidson', *davidson_t0) + davidson_wall = davidson_t1[1] - davidson_t0[1] + total_t1 = log.timer('GPU FCI solver', *fci_t0) + total_wall = total_t1[1] - fci_t0[1] + self.timing = { + 'total_wall': total_wall, + 'setup_wall': setup_wall, + 'davidson_wall': davidson_wall, + 'davidson_iterations': hop_calls, + 'davidson_avg_wall': davidson_wall / max(1, hop_calls), + 'contract_2e_calls': hop_calls, + 'contract_2e_vectors': hop_vectors, + } + log.debug('GPU FCI timing: total %.3f s; setup %.3f s; Davidson ' + '%.3f s in %d iterations (%.3f s/iteration)', + total_wall, setup_wall, davidson_wall, hop_calls, + self.timing['davidson_avg_wall']) + + neleca, nelecb = nelec + na = lib.comb(norb, neleca) + nb = lib.comb(norb, nelecb) + self.norb = norb + self.nelec = nelec + if nroots == 1: + self.converged = bool(converged[0].get()) + self.eci = float(energies[0]) + ecore + self.ci = ci[0].reshape(na, nb) + else: + self.converged = cp.asnumpy(converged) + self.eci = energies + ecore + self.ci = [root.reshape(na, nb) for root in ci[:nroots]] + return self.eci, self.ci + + def energy(self, h1e, eri, fcivec, norb, nelec, link_index=None): + h2e = self.absorb_h1e(h1e, eri, norb, nelec, .5) + ci1 = contract_2e(h2e, fcivec, norb, nelec, link_index) + return cp.vdot(cp.asarray(fcivec).ravel(), ci1.ravel()).real + + def spin_square(self, fcivec, norb, nelec): + return super().spin_square(cp.asnumpy(fcivec), norb, nelec) + + def make_rdm1s(self, fcivec, norb, nelec, link_index=None): + return super().make_rdm1s(cp.asnumpy(fcivec), norb, nelec, link_index) + + def make_rdm1(self, fcivec, norb, nelec, link_index=None): + return super().make_rdm1(cp.asnumpy(fcivec), norb, nelec, link_index) + + def make_rdm12(self, fcivec, norb, nelec, link_index=None, reorder=True): + dm1, dm2 = make_rdm12(fcivec, norb, nelec, link_index, reorder) + return cp.asnumpy(dm1), cp.asnumpy(dm2) + + +FCI = FCISolver diff --git a/gpu4pyscf/fci/tests/test_direct_spin1.py b/gpu4pyscf/fci/tests/test_direct_spin1.py index fbf9a74a7..e72805a9a 100644 --- a/gpu4pyscf/fci/tests/test_direct_spin1.py +++ b/gpu4pyscf/fci/tests/test_direct_spin1.py @@ -1,24 +1,60 @@ -import numpy as np +# Copyright 2026 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. + +import unittest + import cupy as cp +import pyscf + +from gpu4pyscf.fci.direct_spin1 import FCISolver from gpu4pyscf.fci.direct_spin1 import contract_2e -from pyscf.fci import cistring, direct_spin1 - -def test_contract_2e(): - norb = 12 - nelec = 12 - npair = norb * (norb + 1) // 2 - np.random.seed(np.asarray(12, np.uint64)) - g2e = np.random.rand(npair,npair) - g2e = g2e + g2e.T - link = cistring.gen_linkstr_index(range(norb), nelec//2, tril=True) - na = link.shape[0] - cp.random.seed(np.asarray(11, np.uint64)) - ci0 = cp.random.rand(na) - ci0 = cp.einsum('i,j->ij', ci0, ci0) - ci0 *= 1/cp.linalg.norm(ci0) - - ci1 = contract_2e(g2e, ci0, norb, nelec, (link, link)) - - ci0 = ci0.get() - ref = direct_spin1.contract_2e(g2e, ci0, norb, nelec, (link, link)) - assert abs(ci1 - ref).max() < 1e-12 +from gpu4pyscf.fci.direct_spin1 import make_rdm12 +from pyscf import ao2mo, mcscf, scf +from pyscf.fci import direct_spin1 + + +class KnownValues(unittest.TestCase): + def test_cpu_gpu(self): + norb = 6 + nelec = (3, 3) + mol = pyscf.M( + atom='N 0 0 0; N 0 0 1.1', basis='sto-3g', verbose=0) + mf = scf.RHF(mol).run(conv_tol=1e-12) + mc = mcscf.CASCI(mf, norb, sum(nelec)) + mc.canonicalization = False + e_ref, _, ci = mc.kernel()[:3] + h1e, ecore = mc.get_h1eff() + eri = ao2mo.restore(4, mc.get_h2eff(), norb) + + link = direct_spin1._unpack(norb, nelec, None) + ref = direct_spin1.contract_2e(eri, ci, norb, nelec, link) + out = contract_2e(eri, cp.asarray(ci), norb, nelec, link) + self.assertLess(abs(out.get() - ref).max(), 1e-10) + + solver = FCISolver() + e_gpu, ci_gpu = solver.kernel(h1e, eri, norb, nelec, ecore=ecore) + + self.assertTrue(solver.converged) + self.assertIsInstance(ci_gpu, cp.ndarray) + self.assertLess(abs(e_gpu - e_ref), 1e-8) + + dm1, dm2 = make_rdm12(ci_gpu, norb, nelec) + ref1, ref2 = direct_spin1.make_rdm12( + cp.asnumpy(ci_gpu), norb, nelec) + self.assertLess(abs(dm1.get() - ref1).max(), 1e-10) + self.assertLess(abs(dm2.get() - ref2).max(), 1e-10) + + +if __name__ == '__main__': + unittest.main() diff --git a/gpu4pyscf/mcscf/__init__.py b/gpu4pyscf/mcscf/__init__.py new file mode 100644 index 000000000..fb9c1390a --- /dev/null +++ b/gpu4pyscf/mcscf/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 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. + +from .df import DFCASCI, DFCASSCF diff --git a/gpu4pyscf/mcscf/casci.py b/gpu4pyscf/mcscf/casci.py new file mode 100644 index 000000000..913d31cfa --- /dev/null +++ b/gpu4pyscf/mcscf/casci.py @@ -0,0 +1,152 @@ +# Copyright 2026 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. + +from functools import reduce + +import cupy +from pyscf import gto +from pyscf import scf as cpu_scf +from pyscf.mcscf import casci as cpu_casci + +from gpu4pyscf import scf +from gpu4pyscf.fci import direct_spin1 as gpu_direct_spin1 +from gpu4pyscf.lib import logger +from gpu4pyscf.lib import utils +from gpu4pyscf.scf import hf + + +def h1e_for_cas(casci, mo_coeff=None, ncas=None, ncore=None, hcore=None, + return_corevhf=False): + wall0 = logger.perf_counter() + if mo_coeff is None: + mo_coeff = casci.mo_coeff + if ncas is None: + ncas = casci.ncas + if ncore is None: + ncore = casci.ncore + + mo_coeff = cupy.asarray(mo_coeff) + mo_core = mo_coeff[:, :ncore] + mo_cas = mo_coeff[:, ncore:ncore+ncas] + + if hcore is None: + hcore = cupy.asarray(casci.get_hcore()) + else: + hcore = cupy.asarray(hcore) + energy_core = casci.energy_nuc() + if mo_core.size == 0: + corevhf = 0 + else: + core_dm = cupy.dot(mo_core, mo_core.conj().T) * 2 + corevhf = cupy.asarray(casci.get_veff(casci.mol, core_dm)) + energy_core += float(cupy.einsum('ij,ji', core_dm, hcore).real.get()) + energy_core += float(cupy.einsum('ij,ji', core_dm, corevhf).real.get()) * .5 + + h1eff = reduce(cupy.dot, (mo_cas.conj().T, hcore+corevhf, mo_cas)) + out = h1eff.get(), float(energy_core) + if return_corevhf: + out += (corevhf,) + timing = getattr(casci, 'timing', None) + if isinstance(timing, dict): + timing['h1e_wall'] = (timing.get('h1e_wall', 0.) + + logger.perf_counter() - wall0) + return out + + +class _CASCI(cpu_casci.CASCI): + _keys = cpu_casci.CASCI._keys.union({'timing'}) + canonicalization = False + + get_h1eff = h1e_for_cas + h1e_for_cas = h1e_for_cas + + to_cpu = utils.to_cpu + to_gpu = utils.to_gpu + device = utils.device + + def __init__(self, mf_or_mol, ncas=0, nelecas=0, ncore=None): + if isinstance(mf_or_mol, gto.MoleBase): + mf_or_mol = scf.RHF(mf_or_mol) + elif (hasattr(mf_or_mol, 'istype') and + any(mf_or_mol.istype(x) for x in ('UHF', 'ROHF', 'GHF'))): + raise NotImplementedError( + 'GPU CASCI supports restricted HF/KS objects only') + elif not getattr(mf_or_mol, '__module__', '').startswith('gpu4pyscf'): + if isinstance(mf_or_mol, cpu_scf.hf.RHF): + logger.warn( + mf_or_mol, + 'CPU restricted SCF object converted to GPU for CASCI') + mf_or_mol = mf_or_mol.to_gpu() + else: + raise NotImplementedError( + 'GPU CASCI supports restricted HF/KS objects only') + + if not isinstance(mf_or_mol, hf.RHF): + raise NotImplementedError( + 'GPU CASCI supports restricted HF/KS objects only') + + super().__init__(mf_or_mol, ncas, nelecas, ncore) + fcisolver = gpu_direct_spin1.FCISolver(self.mol) + fcisolver.__dict__.update(self.fcisolver.__dict__) + self.fcisolver = fcisolver + self.canonicalization = False + + def energy_nuc(self): + return self._scf.energy_nuc() + + def get_hcore(self, mol=None): + return self._scf.get_hcore(mol) + + def get_jk(self, mol, dm, hermi=1, with_j=True, with_k=True, omega=None): + return self._scf.get_jk( + mol, dm, hermi, with_j=with_j, with_k=with_k, omega=omega) + + def get_veff(self, mol=None, dm=None, hermi=1): + if mol is None: + mol = self.mol + if dm is None: + mo_core = cupy.asarray(self.mo_coeff[:, :self.ncore]) + dm = mo_core @ mo_core.conj().T * 2 + vj, vk = self.get_jk(mol, dm, hermi) + return vj - vk * .5 + + def get_h2eff(self, mo_coeff=None): + raise NotImplementedError('CASCI integral backend is not configured') + + def kernel(self, mo_coeff=None, ci0=None, verbose=None): + if self.canonicalization: + raise NotImplementedError('GPU CASCI canonicalization is not implemented') + if self.natorb: + raise NotImplementedError('GPU CASCI natural orbitals are not implemented') + self.timing = {} + wall0 = logger.perf_counter() + out = super().kernel(mo_coeff, ci0, verbose) + total_wall = logger.perf_counter() - wall0 + fci_timing = dict(getattr(self.fcisolver, 'timing', {})) + ao2mo_wall = self.timing.get('ao2mo_wall', 0.) + h1e_wall = self.timing.get('h1e_wall', 0.) + fci_wall = fci_timing.get('total_wall', 0.) + postprocess_wall = total_wall - ao2mo_wall - h1e_wall - fci_wall + self.timing.update({ + 'total_wall': total_wall, + 'ao2mo_wall': ao2mo_wall, + 'h1e_wall': h1e_wall, + 'fci': fci_timing, + 'postprocess_wall': postprocess_wall, + }) + log = logger.new_logger(self, verbose) + log.debug('CASCI timing: total %.3f s; AO2MO %.3f s; h1e %.3f s; ' + 'FCI %.3f s; postprocess %.3f s', total_wall, ao2mo_wall, + h1e_wall, fci_wall, postprocess_wall) + return out diff --git a/gpu4pyscf/mcscf/casscf.py b/gpu4pyscf/mcscf/casscf.py new file mode 100644 index 000000000..233484296 --- /dev/null +++ b/gpu4pyscf/mcscf/casscf.py @@ -0,0 +1,336 @@ +# Copyright 2026 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. + +from functools import reduce +import math + +import cupy +from cupyx.scipy.linalg import expm +from pyscf import gto +from pyscf import lib +from pyscf import scf as cpu_scf +from pyscf.mcscf import mc1step as cpu_mc1step + +from gpu4pyscf import scf +from gpu4pyscf.fci import direct_spin1 as gpu_direct_spin1 +from gpu4pyscf.lib import logger +from gpu4pyscf.lib import utils +from gpu4pyscf.lib.cupy_helper import contract +from gpu4pyscf.mcscf.casci import h1e_for_cas +from gpu4pyscf.scf import hf + + +class _ERIS: + def __init__(self, casscf, mo, paaa, hcore=None): + ncore = casscf.ncore + nocc = ncore + casscf.ncas + self.paaa = cupy.asarray(paaa) + self.aaaa = self.paaa[ncore:nocc].copy() + if hcore is None: + hcore = cupy.asarray(casscf.get_hcore()) + self.hcore = hcore + self.h1eff, self.ecore, self.vhf_c = casscf.h1e_for_cas( + mo, hcore=hcore, return_corevhf=True) + + +def gen_g_hdiag(casscf, mo, casdm1, casdm2, eris): + ncore = casscf.ncore + ncas = casscf.ncas + nocc = ncore + ncas + nmo = mo.shape[1] + + core = slice(0, ncore) + act = slice(ncore, nocc) + vir = slice(nocc, nmo) + mo_act = mo[:, act] + + fcore = reduce(cupy.dot, (mo.T, eris.hcore + eris.vhf_c, mo)) + + dm_act = reduce(cupy.dot, (mo_act, casdm1, mo_act.T)) + jact, kact = casscf.get_jk(casscf.mol, dm_act) + fact = reduce(cupy.dot, (mo.T, 2*jact - kact, mo)) + + y = fcore[:, act] @ casdm1 + z = contract('puvw,tuvw->pt', eris.paaa, casdm2, alpha=.5) + + g_ia = 4 * fcore[core, vir] + 2 * fact[core, vir] + g_ta = 2 * y[vir, :].T + 4 * z[vir, :].T + g_it = 4 * fcore[core, act] + 2 * fact[core, act] + g_it += -2 * y[core, :] - 4 * z[core, :] + + f_diag = (4 * fcore + 2 * fact).diagonal() + fcore_diag = fcore.diagonal() + fact_diag = fact.diagonal() + gamma_diag = casdm1.diagonal() + y_diag = y[act].diagonal() + z_diag = z[act].diagonal() + + d_ia = f_diag[vir][None, :] - f_diag[core][:, None] + d_ta = 2 * gamma_diag[:, None] * fcore_diag[vir][None, :] + d_ta += gamma_diag[:, None] * fact_diag[vir][None, :] + d_ta += (-2 * y_diag - 4 * z_diag)[:, None] + d_it = f_diag[act][None, :] - f_diag[core][:, None] + d_it += 2 * gamma_diag[None, :] * fcore_diag[core][:, None] + d_it += gamma_diag[None, :] * fact_diag[core][:, None] + d_it += (-2 * y_diag - 4 * z_diag)[None, :] + + return g_ia, g_ta, g_it, d_ia, d_ta, d_it + + +def build_rotation_matrix(casscf, mo, terms, denom_floor=1e-8, + max_abs_step=.03): + ncore = casscf.ncore + nocc = ncore + casscf.ncas + nmo = mo.shape[1] + core = slice(0, ncore) + act = slice(ncore, nocc) + vir = slice(nocc, nmo) + g_ia, g_ta, g_it, d_ia, d_ta, d_it = terms + + if max_abs_step is not None and max_abs_step <= 0: + raise ValueError('max_abs_step must be positive') + + # Apply one level shift to keep the diagonal model positive. The orbital + # step limit is applied afterwards and does not make this shift more + # conservative. + level_shift = cupy.asarray(0., dtype=mo.dtype) + gd_pairs = ((g_ia, d_ia), (g_ta, d_ta), (g_it, d_it)) + for _, hdiag in gd_pairs: + if hdiag.size: + level_shift = cupy.maximum(level_shift, cupy.max(denom_floor - hdiag)) + + s_ia = g_ia / (d_ia + level_shift) + s_ta = g_ta / (d_ta + level_shift) + s_it = g_it / (d_it + level_shift) + if max_abs_step is not None: + s_ia = cupy.clip(s_ia, -max_abs_step, max_abs_step) + s_ta = cupy.clip(s_ta, -max_abs_step, max_abs_step) + s_it = cupy.clip(s_it, -max_abs_step, max_abs_step) + + s = cupy.zeros((nmo, nmo), dtype=mo.dtype) + s[core, vir] = s_ia + s[vir, core] = -s_ia.T + s[act, vir] = s_ta + s[vir, act] = -s_ta.T + s[core, act] = s_it + s[act, core] = -s_it.T + return s + + +def kernel(casscf, mo_coeff, tol=1e-7, conv_tol_grad=None, ci0=None, + callback=None, verbose=logger.NOTE, dump_chk=True): + log = logger.new_logger(casscf, verbose) + cput0 = (logger.process_clock(), logger.perf_counter()) + if callback is None: + callback = casscf.callback + if ci0 is None: + ci0 = casscf.ci + + mo = cupy.asarray(mo_coeff) + if conv_tol_grad is None: + conv_tol_grad = math.sqrt(tol) + logger.info(casscf, 'Set conv_tol_grad to %g', conv_tol_grad) + + conv = False + e_last = None + e_tot = e_cas = fcivec = eris = casdm1 = None + denom_floor = casscf.denom_floor + max_stepsize = casscf.max_stepsize + timing = casscf.timing = { + 'macro_cycles': 0, + 'hcore_wall': 0., + 'ao2mo_wall': 0., + 'h1e_wall': 0., + 'fci_wall': 0., + 'fci_setup_wall': 0., + 'fci_davidson_wall': 0., + 'fci_iterations': 0, + 'rdm_wall': 0., + 'orbital_derivatives_wall': 0., + 'orbital_rotation_wall': 0., + } + t0 = log.init_timer() + hcore = cupy.asarray(casscf.get_hcore()) + timing['hcore_wall'] = log.timer_silent(*t0)[2] * 1e-3 + + for istep in range(1, casscf.max_cycle_macro + 1): + t0 = log.init_timer() + eris = casscf.ao2mo(mo, hcore=hcore) + timing['ao2mo_wall'] += log.timer_silent(*t0)[2] * 1e-3 + + t0 = log.init_timer() + e_tot, e_cas, fcivec = casscf.casci( + mo, ci0, eris, log, locals()) + timing['fci_wall'] += log.timer_silent(*t0)[2] * 1e-3 + fci_timing = getattr(casscf.fcisolver, 'timing', {}) + timing['fci_setup_wall'] += fci_timing.get('setup_wall', 0.) + timing['fci_davidson_wall'] += fci_timing.get('davidson_wall', 0.) + timing['fci_iterations'] += fci_timing.get('davidson_iterations', 0) + + t0 = log.init_timer() + casdm1, casdm2 = casscf.fcisolver.make_rdm12(fcivec, casscf.ncas, + casscf.nelecas) + casdm1 = cupy.asarray(casdm1) + casdm2 = cupy.asarray(casdm2) + timing['rdm_wall'] += log.timer_silent(*t0)[2] * 1e-3 + + t0 = log.init_timer() + terms = gen_g_hdiag(casscf, mo, casdm1, casdm2, eris) + g_norm = max((float(cupy.abs(x).max().get()) + for x in terms[:3] if x.size), default=0.) + timing['orbital_derivatives_wall'] += ( + log.timer_silent(*t0)[2] * 1e-3) + timing['macro_cycles'] = istep + de = e_tot - e_last if e_last is not None else e_tot + log.info('cycle %3d E = %#.15g de = %.6g |g| = %.6g', + istep, e_tot, de, g_norm) + if max_stepsize is not None: + max_stepsize = casscf.max_stepsize_scheduler(locals()) + if callable(callback): + callback(locals()) + + if e_last is not None and abs(de) < tol and g_norm < conv_tol_grad: + conv = True + break + e_last = e_tot + ci0 = fcivec + if istep == casscf.max_cycle_macro: + break + + t0 = log.init_timer() + s = build_rotation_matrix(casscf, mo, terms, denom_floor, max_stepsize) + mo = mo @ expm(s) + timing['orbital_rotation_wall'] += log.timer_silent(*t0)[2] * 1e-3 + if dump_chk and casscf.chkfile: + chk_env = locals().copy() + chk_env['mo'] = cupy.asnumpy(mo) + chk_env['casdm1'] = cupy.asnumpy(casdm1) + if casscf.chk_ci: + chk_env['fcivec'] = cupy.asnumpy(fcivec) + casscf.dump_chk(chk_env) + + if conv: + log.info('Diagonal-Hessian CASSCF converged in %3d steps', istep) + else: + log.info('Diagonal-Hessian CASSCF not converged in %3d steps', istep) + if dump_chk and casscf.chkfile: + chk_env = locals().copy() + chk_env['mo'] = cupy.asnumpy(mo) + chk_env['casdm1'] = cupy.asnumpy(casdm1) + if casscf.chk_ci: + chk_env['fcivec'] = cupy.asnumpy(fcivec) + casscf.dump_chk(chk_env) + log.timer('Diagonal-Hessian CASSCF', *cput0) + return conv, e_tot, e_cas, fcivec, mo, None + + +class _CASSCF(cpu_mc1step.CASSCF): + _keys = cpu_mc1step.CASSCF._keys.union({'denom_floor', 'timing'}) + canonicalization = False + denom_floor = 1e-8 + max_stepsize = .04 + + get_h1eff = h1e_for_cas + h1e_for_cas = h1e_for_cas + to_cpu = utils.to_cpu + to_gpu = utils.to_gpu + device = utils.device + + def __init__(self, mf_or_mol, ncas=0, nelecas=0, ncore=None, frozen=None): + if frozen is not None: + raise NotImplementedError('GPU CASSCF frozen orbitals are not implemented') + if isinstance(mf_or_mol, gto.MoleBase): + mf_or_mol = scf.RHF(mf_or_mol) + elif (hasattr(mf_or_mol, 'istype') and + any(mf_or_mol.istype(x) for x in ('UHF', 'ROHF', 'GHF'))): + raise NotImplementedError( + 'GPU CASSCF supports restricted HF/KS objects only') + elif not getattr(mf_or_mol, '__module__', '').startswith('gpu4pyscf'): + if not isinstance(mf_or_mol, cpu_scf.hf.RHF): + raise NotImplementedError( + 'GPU CASSCF supports restricted HF/KS objects only') + mf_or_mol = mf_or_mol.to_gpu() + + if not isinstance(mf_or_mol, hf.RHF): + raise NotImplementedError( + 'GPU CASSCF supports restricted HF/KS objects only') + + super().__init__(mf_or_mol, ncas, nelecas, ncore, frozen) + fcisolver = gpu_direct_spin1.FCISolver(self.mol) + fcisolver.__dict__.update(self.fcisolver.__dict__) + self.fcisolver = fcisolver + self.canonicalization = False + + def ao2mo(self, mo_coeff=None, hcore=None): + raise NotImplementedError('CASSCF integral backend is not configured') + + def casci(self, mo_coeff, ci0=None, eris=None, verbose=None, envs=None): + if eris is None: + eris = self.ao2mo(mo_coeff) + max_memory = max(400, self.max_memory - lib.current_memory()[0]) + e_tot, fcivec = self.fcisolver.kernel( + eris.h1eff, eris.aaaa, self.ncas, self.nelecas, + ci0=ci0, verbose=verbose, max_memory=max_memory, + ecore=eris.ecore) + return e_tot, e_tot - eris.ecore, fcivec + + def kernel(self, mo_coeff=None, ci0=None, callback=None): + if self.canonicalization: + raise NotImplementedError('GPU CASSCF canonicalization is not implemented') + if self.natorb: + raise NotImplementedError('GPU CASSCF natural orbitals are not implemented') + self.timing = {} + wall0 = logger.perf_counter() + + if mo_coeff is None: + if self.mo_coeff is None and self._scf.mol.nelectron > 0: + self._scf.run() + self.mo_coeff = self._scf.mo_coeff + mo_coeff = self.mo_coeff + else: + self.mo_coeff = mo_coeff + if ci0 is None: + ci0 = self.ci + if callback is None: + callback = self.callback + + self.check_sanity() + self.dump_flags() + (self.converged, self.e_tot, self.e_cas, self.ci, + self.mo_coeff, self.mo_energy) = kernel( + self, mo_coeff, tol=self.conv_tol, + conv_tol_grad=self.conv_tol_grad, ci0=ci0, + callback=callback, verbose=self.verbose) + logger.note(self, 'CASSCF energy = %#.15g', self.e_tot) + self._finalize() + + total_wall = logger.perf_counter() - wall0 + accounted_wall = sum(self.timing[key] for key in ( + 'hcore_wall', 'ao2mo_wall', 'fci_wall', 'rdm_wall', + 'orbital_derivatives_wall', 'orbital_rotation_wall')) + self.timing['total_wall'] = total_wall + self.timing['other_wall'] = max(0., total_wall - accounted_wall) + log = logger.new_logger(self) + log.debug( + 'CASSCF timing: total %.3f s in %d cycles; hcore %.3f s; ' + 'AO2MO %.3f s; FCI %.3f s; RDM %.3f s; orbital derivatives ' + '%.3f s; orbital rotation %.3f s; other %.3f s', + total_wall, self.timing['macro_cycles'], + self.timing['hcore_wall'], + self.timing['ao2mo_wall'], self.timing['fci_wall'], + self.timing['rdm_wall'], + self.timing['orbital_derivatives_wall'], + self.timing['orbital_rotation_wall'], self.timing['other_wall']) + return (self.e_tot, self.e_cas, self.ci, self.mo_coeff, + self.mo_energy) diff --git a/gpu4pyscf/mcscf/df.py b/gpu4pyscf/mcscf/df.py new file mode 100644 index 000000000..0674a6b8b --- /dev/null +++ b/gpu4pyscf/mcscf/df.py @@ -0,0 +1,120 @@ +# Copyright 2026 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. + +import cupy +import numpy +from pyscf import mcscf as cpu_mcscf +from pyscf.mcscf import mc1step as cpu_mc1step + +from gpu4pyscf.df import df as gpu_df +from gpu4pyscf.lib import logger +from gpu4pyscf.lib import utils +from gpu4pyscf.mcscf.casci import _CASCI +from gpu4pyscf.mcscf.casscf import _CASSCF, _ERIS + + +def _get_with_df(mc, auxbasis=None, with_df=None): + if with_df is None: + scf_df = getattr(mc._scf, 'with_df', None) + if (scf_df is not None and + (auxbasis is None or auxbasis == scf_df.auxbasis)): + return scf_df + + if auxbasis is None and isinstance(mc.mol.basis, str): + from pyscf.df.addons import predefined_auxbasis + auxbasis = predefined_auxbasis(mc.mol, mc.mol.basis, xc='HF') + with_df = gpu_df.DF(mc.mol, auxbasis) + with_df.max_memory = mc.max_memory + with_df.stdout = mc.stdout + with_df.verbose = mc.verbose + return with_df + + +class _DFCAS: + _keys = {'with_df'} + + def reset(self, mol=None): + if self.with_df is not getattr(self._scf, 'with_df', None): + self.with_df.reset(mol) + return super().reset(mol) + + def get_jk(self, mol, dm, hermi=1, with_j=True, with_k=True, omega=None): + return self.with_df.get_jk( + dm, hermi, with_j=with_j, with_k=with_k, omega=omega) + + def get_h2eff(self, mo_coeff=None): + wall0 = logger.perf_counter() + ncore = self.ncore + nocc = ncore + self.ncas + if mo_coeff is None: + mo_coeff = self.mo_coeff[:, ncore:nocc] + elif mo_coeff.shape[1] != self.ncas: + mo_coeff = mo_coeff[:, ncore:nocc] + eri = self.with_df.ao2mo(mo_coeff) + out = eri.get() if isinstance(eri, cupy.ndarray) else eri + timing = getattr(self, 'timing', None) + if isinstance(timing, dict): + timing['ao2mo_wall'] = (timing.get('ao2mo_wall', 0.) + + logger.perf_counter() - wall0) + return out + + def to_cpu(self): + if isinstance(self, _CASSCF): + out = cpu_mcscf.DFCASSCF( + self._scf.to_cpu(), self.ncas, self.nelecas, + auxbasis=self.with_df.auxbasis, ncore=self.ncore) + else: + out = cpu_mcscf.DFCASCI( + self._scf.to_cpu(), self.ncas, self.nelecas, + auxbasis=self.with_df.auxbasis, ncore=self.ncore) + return utils.to_cpu(self, out=out) + + +class DFCASCI(_DFCAS, _CASCI): + def __init__(self, mf_or_mol, ncas, nelecas, auxbasis=None, ncore=None, + with_df=None): + _CASCI.__init__(self, mf_or_mol, ncas, nelecas, ncore) + self.with_df = _get_with_df(self, auxbasis, with_df) + + +class DFCASSCF(_DFCAS, _CASSCF): + def __init__(self, mf_or_mol, ncas, nelecas, auxbasis=None, ncore=None, + frozen=None, with_df=None): + _CASSCF.__init__(self, mf_or_mol, ncas, nelecas, ncore, frozen) + self.with_df = _get_with_df(self, auxbasis, with_df) + + def ao2mo(self, mo_coeff=None, hcore=None): + if mo_coeff is None: + mo_coeff = self.mo_coeff + ncore = self.ncore + nocc = ncore + self.ncas + mo_cas = mo_coeff[:, ncore:nocc] + paaa = self.with_df.ao2mo( + (mo_coeff, mo_cas, mo_cas, mo_cas), compact=False) + paaa = paaa.reshape(mo_coeff.shape[1], self.ncas, + self.ncas, self.ncas) + return _ERIS(self, mo_coeff, paaa, hcore) + + +def from_cpu(mc): + cls = DFCASSCF if isinstance(mc, cpu_mc1step.CASSCF) else DFCASCI + out = cls(mc._scf, mc.ncas, mc.nelecas, + auxbasis=mc.with_df.auxbasis, ncore=mc.ncore) + for key, value in mc.__dict__.items(): + if key not in ('_scf', 'with_df', 'fcisolver'): + if isinstance(value, numpy.ndarray): + value = cupy.asarray(value) + out.__dict__[key] = value + out.fcisolver.__dict__.update(mc.fcisolver.__dict__) + return out diff --git a/gpu4pyscf/mcscf/tests/test_casscf.py b/gpu4pyscf/mcscf/tests/test_casscf.py new file mode 100644 index 000000000..0fe82939e --- /dev/null +++ b/gpu4pyscf/mcscf/tests/test_casscf.py @@ -0,0 +1,79 @@ +# Copyright 2026 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. + +import unittest + +import cupy +import pyscf +from pyscf import dft, mcscf as cpu_mcscf, scf + +from gpu4pyscf import mcscf +from gpu4pyscf.fci.direct_spin1 import FCISolver +from gpu4pyscf.mcscf import df as gpu_mcscf_df + + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.mol = pyscf.M( + atom='N 0 0 -0.7; N 0 0 0.7', basis='sto-3g', + verbose=0, output='/dev/null') + cls.mf_cpu = scf.RHF(cls.mol).density_fit(auxbasis='weigend') + cls.mf_cpu.kernel() + cls.mf_gpu = cls.mf_cpu.to_gpu() + cls.mf_rks_cpu = dft.RKS(cls.mol).density_fit(auxbasis='weigend') + cls.mf_rks_cpu.xc = 'pbe' + cls.mf_rks_cpu.kernel() + cls.mf_rks_gpu = cls.mf_rks_cpu.to_gpu() + + @classmethod + def tearDownClass(cls): + cls.mol.stdout.close() + + def test_df_casscf(self): + mc = mcscf.DFCASSCF(self.mf_gpu, 4, 4) + mc.max_cycle_macro = 20 + mc.conv_tol = 1e-8 + mc.conv_tol_grad = 1e-5 + e_tot = mc.kernel()[0] + + self.assertIsInstance(mc, gpu_mcscf_df.DFCASSCF) + self.assertIsInstance(mc.fcisolver, FCISolver) + self.assertIsInstance(mc.ci, cupy.ndarray) + self.assertIsInstance(mc.mo_coeff, cupy.ndarray) + self.assertTrue(mc.converged) + self.assertLess(abs(e_tot - -107.5445420582518), 2e-8) + self.assertGreater(mc.timing['macro_cycles'], 1) + self.assertLessEqual(mc.timing['macro_cycles'], mc.max_cycle_macro) + self.assertGreater(mc.timing['ao2mo_wall'], 0.) + self.assertGreater(mc.timing['fci_wall'], 0.) + self.assertGreater(mc.timing['orbital_derivatives_wall'], 0.) + self.assertGreater(mc.timing['total_wall'], 0.) + + def test_rks_reference(self): + ref = cpu_mcscf.DFCASCI( + self.mf_rks_cpu, 4, 4, auxbasis='weigend') + ref.canonicalization = False + mc = mcscf.DFCASSCF(self.mf_rks_gpu, 4, 4) + mc.max_cycle_macro = 1 + e_tot = mc.kernel()[0] + + self.assertIsInstance(mc, gpu_mcscf_df.DFCASSCF) + self.assertIs(mc._scf, self.mf_rks_gpu) + self.assertIs(mc.with_df, self.mf_rks_gpu.with_df) + self.assertLess(abs(e_tot - ref.kernel()[0]), 1e-8) + + +if __name__ == '__main__': + unittest.main() diff --git a/gpu4pyscf/mcscf/tests/test_df_casci.py b/gpu4pyscf/mcscf/tests/test_df_casci.py new file mode 100644 index 000000000..c3a0aa000 --- /dev/null +++ b/gpu4pyscf/mcscf/tests/test_df_casci.py @@ -0,0 +1,102 @@ +# Copyright 2026 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. + +import unittest + +import pyscf +from pyscf import ao2mo +from pyscf import dft as cpu_dft +from pyscf import mcscf as cpu_mcscf +from pyscf import scf as cpu_scf + +from gpu4pyscf import mcscf +from gpu4pyscf.mcscf import df as gpu_mcscf_df + + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.mol = pyscf.M( + atom='N 0 0 -0.7; N 0 0 0.7', basis='sto-3g', + verbose=0, output='/dev/null') + cls.mf_cpu = cpu_scf.RHF(cls.mol).density_fit(auxbasis='weigend') + cls.mf_cpu.conv_tol = 1e-12 + cls.mf_cpu.kernel() + cls.mf_gpu = cls.mf_cpu.to_gpu() + cls.mf_rks_cpu = cpu_dft.RKS(cls.mol).density_fit(auxbasis='weigend') + cls.mf_rks_cpu.xc = 'pbe' + cls.mf_rks_cpu.conv_tol = 1e-12 + cls.mf_rks_cpu.kernel() + cls.mf_rks_gpu = cls.mf_rks_cpu.to_gpu() + + @classmethod + def tearDownClass(cls): + cls.mol.stdout.close() + + def test_df_casci(self): + mc_cpu = cpu_mcscf.DFCASCI(self.mf_cpu, 4, 4, auxbasis='weigend') + mc_cpu.canonicalization = False + mc_gpu = mcscf.DFCASCI(self.mf_gpu, 4, 4) + + self.assertIsInstance(mc_gpu, gpu_mcscf_df.DFCASCI) + h1_cpu, ecore_cpu = mc_cpu.get_h1eff() + h1_gpu, ecore_gpu = mc_gpu.get_h1eff() + self.assertLess(abs(h1_gpu - h1_cpu).max(), 1e-8) + self.assertLess(abs(ecore_gpu - ecore_cpu), 1e-8) + + eri_cpu = ao2mo.restore(1, mc_cpu.get_h2eff(), mc_cpu.ncas) + eri_gpu = ao2mo.restore(1, mc_gpu.get_h2eff(), mc_gpu.ncas) + self.assertLess(abs(eri_gpu - eri_cpu).max(), 1e-8) + + e_cpu = mc_cpu.kernel()[0] + e_gpu = mc_gpu.kernel()[0] + self.assertLess(abs(e_gpu - e_cpu), 1e-8) + self.assertGreater(mc_gpu.timing['ao2mo_wall'], 0) + self.assertGreater(mc_gpu.timing['h1e_wall'], 0) + self.assertGreater(mc_gpu.timing['fci']['contract_2e_calls'], 0) + + def test_rks_reference(self): + mc_cpu = cpu_mcscf.DFCASCI( + self.mf_rks_cpu, 4, 4, auxbasis='weigend') + mc_cpu.canonicalization = False + mc_gpu = mcscf.DFCASCI(self.mf_rks_gpu, 4, 4) + + self.assertIsInstance(mc_gpu, gpu_mcscf_df.DFCASCI) + self.assertIs(mc_gpu._scf, self.mf_rks_gpu) + self.assertIs(mc_gpu.with_df, self.mf_rks_gpu.with_df) + self.assertLess(abs(mc_gpu.kernel()[0] - mc_cpu.kernel()[0]), 1e-8) + + def test_conversions(self): + ref = mcscf.DFCASCI(self.mf_gpu, 4, 4).kernel()[0] + + mc_gpu = mcscf.DFCASCI(self.mf_cpu, 4, 4, auxbasis='weigend') + self.assertLess(abs(mc_gpu.kernel()[0] - ref), 1e-9) + + mc_cpu = cpu_mcscf.DFCASCI( + self.mf_cpu, 4, 4, auxbasis='weigend') + mc_cpu.canonicalization = False + mc_from_cpu = mc_cpu.to_gpu() + self.assertIsInstance(mc_from_cpu, gpu_mcscf_df.DFCASCI) + self.assertLess(abs(mc_from_cpu.kernel()[0] - ref), 1e-8) + + mc_to_cpu = mc_gpu.to_cpu() + self.assertIsInstance(mc_to_cpu, cpu_mcscf.casci.CASCI) + self.assertIsNotNone(mc_to_cpu.with_df) + mc_to_cpu.canonicalization = False + self.assertLess(abs(mc_to_cpu.kernel()[0] - ref), 1e-8) + + +if __name__ == '__main__': + print('Full tests for GPU DFCASCI') + unittest.main()