diff --git a/gpu4pyscf/df/tests/test_df_hessian.py b/gpu4pyscf/df/tests/test_df_hessian.py index a5a365cf7..8061b876a 100644 --- a/gpu4pyscf/df/tests/test_df_hessian.py +++ b/gpu4pyscf/df/tests/test_df_hessian.py @@ -358,7 +358,7 @@ def test_unstable_j2c(self): mo_occ = mf.mo_occ test_hessian_round2 = hobj.partial_hess_elec(mo_energy, mo_coeff, mo_occ) - assert np.max(np.abs(test_hessian_round1 - test_hessian_round2)) < 2e-7 + assert abs(test_hessian_round1 - test_hessian_round2).max().item() < 2e-7 if __name__ == "__main__": print("Full Tests for DF Hessian") diff --git a/gpu4pyscf/df/tests/test_df_rks_grad.py b/gpu4pyscf/df/tests/test_df_rks_grad.py index c69a0b63a..9a4b26d2b 100644 --- a/gpu4pyscf/df/tests/test_df_rks_grad.py +++ b/gpu4pyscf/df/tests/test_df_rks_grad.py @@ -94,8 +94,8 @@ def _check_grad(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-6): grad_fd = np.array(grad_fd).reshape(-1,3) print('finite difference gradient:') print(grad_fd) - print('difference between analytical and finite difference gradient:', cupy.linalg.norm(g_analy - grad_fd)) - assert(cupy.linalg.norm(g_analy - grad_fd) < tol) + print('difference between analytical and finite difference gradient:', np.linalg.norm(g_analy - grad_fd)) + assert(np.linalg.norm(g_analy - grad_fd) < tol) def _vs_cpu(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-9): mf = rks.RKS(mol, xc=xc).density_fit(auxbasis=auxbasis0) diff --git a/gpu4pyscf/df/tests/test_df_uhf.py b/gpu4pyscf/df/tests/test_df_uhf.py index cdb3dd2c1..7032f31a7 100644 --- a/gpu4pyscf/df/tests/test_df_uhf.py +++ b/gpu4pyscf/df/tests/test_df_uhf.py @@ -81,8 +81,8 @@ def _check_grad(mol, tol=1e-5, disp=None): grad_fd = np.array(grad_fd).reshape(-1,3) print('finite difference gradient:') print(grad_fd) - print('difference between analytical and finite difference gradient:', cupy.linalg.norm(g_analy - grad_fd)) - assert(cupy.linalg.norm(g_analy - grad_fd) < tol) + print('difference between analytical and finite difference gradient:', np.linalg.norm(g_analy - grad_fd)) + assert(np.linalg.norm(g_analy - grad_fd) < tol) class KnownValues(unittest.TestCase): ''' diff --git a/gpu4pyscf/dft/tests/test_numint.py b/gpu4pyscf/dft/tests/test_numint.py index de51ba2ef..f2a0f77f1 100644 --- a/gpu4pyscf/dft/tests/test_numint.py +++ b/gpu4pyscf/dft/tests/test_numint.py @@ -272,7 +272,8 @@ def test_sparse_index(self): i1 = min(i0+numint.MIN_BLK_SIZE, ngrids) ref = numint._sparse_index( opt._sorted_mol, grids.coords[i0:i1], opt.l_ctr_offsets, ao_loc, opt) - assert all(np.array_equal(r, x) for r, x in zip(ref[1:], dat[i][1:])) + assert all(r.shape == x.shape and bool((r == x).all()) + for r, x in zip(ref[1:], dat[i][1:])) def test_scale_ao(self): ao = cupy.random.rand(1, 3, 256) diff --git a/gpu4pyscf/lib/cupy_helper.py b/gpu4pyscf/lib/cupy_helper.py index 7239603f5..b1234b857 100644 --- a/gpu4pyscf/lib/cupy_helper.py +++ b/gpu4pyscf/lib/cupy_helper.py @@ -1265,7 +1265,7 @@ def batched_vec3_norm2(batched_vec3): ''' einsum('gx,gx->g', vec3, vec3) for the (N,3)-array vec3 ''' - assert type(batched_vec3) is cupy.ndarray + assert isinstance(batched_vec3, cupy.ndarray) assert batched_vec3.dtype == cupy.float64 assert batched_vec3.ndim == 2 assert batched_vec3.shape[0] == 3 or batched_vec3.shape[1] == 3 diff --git a/gpu4pyscf/pbc/df/ft_ao.py b/gpu4pyscf/pbc/df/ft_ao.py index e494b479b..ca67143c5 100644 --- a/gpu4pyscf/pbc/df/ft_ao.py +++ b/gpu4pyscf/pbc/df/ft_ao.py @@ -361,7 +361,7 @@ def ft_evaluator(self, batch_size=None, compressing=True, cart=None, dims, tmp = np.empty_like(dims), dims dims[cell.sorted_idx] = tmp ao_loc = cp.asarray(np.append(0, np.cumsum(dims.ravel()))) - ao_loc = np.append(ao_loc[cell.sorted_idx], nao) + ao_loc = cp.append(ao_loc[cell.sorted_idx], nao) ao_loc = cp.asarray(ao_loc, dtype=np.int32) if batch_size is None: diff --git a/gpu4pyscf/pbc/df/int3c2e.py b/gpu4pyscf/pbc/df/int3c2e.py index c1cc7d11d..e84df4645 100644 --- a/gpu4pyscf/pbc/df/int3c2e.py +++ b/gpu4pyscf/pbc/df/int3c2e.py @@ -128,7 +128,7 @@ def sr_aux_e2(cell, auxcell, omega, kpts=None, bvk_kmesh=None, j_only=False): out = cp.empty((nkpts,nkpts,naux,nao,nao), dtype=np.complex128) kk_conserv = double_translation_indices(int3c2e_opt.bvk_kmesh) for k in range(nkpts): - ki_idx, kj_idx = np.where(kk_conserv == k) + ki_idx, kj_idx = cp.where(kk_conserv == k) out[k] = _unpack_cderi_v2(j3c[k], pair_address, kj_idx, conj_mapping, expLk, nao, axis) j3c = None diff --git a/gpu4pyscf/pbc/df/rsdf_builder.py b/gpu4pyscf/pbc/df/rsdf_builder.py index 209ba5a67..c8e94dab1 100644 --- a/gpu4pyscf/pbc/df/rsdf_builder.py +++ b/gpu4pyscf/pbc/df/rsdf_builder.py @@ -404,8 +404,14 @@ def proc(): ctypes.c_int(naux), ctypes.c_int(nao_pairs), ctypes.c_int(p0), ctypes.c_int(p1)) j3c = None - if num_devices > 1: - stream.synchronize() + # store_col_segment enqueues its device->host write asynchronously + # and does not itself block. The next reader of `cderi` is a + # plain host-side read (mydf.loop()) with no ordering relationship + # to this stream, so this segment must be drained here every + # time -- not only when num_devices > 1. Skipping this on a + # single device happened to be safe on this stream/allocator + # combination but is not guaranteed in general. + stream.synchronize() t1 = log.timer_debug1(f'store int3c2e on Device {device_id}', *t1) multi_gpu.run(proc, non_blocking=True) @@ -572,8 +578,9 @@ def proc(): if err != 0: raise RuntimeError('store_col_segment kernel failed') j3c = None - if num_devices > 1: - stream.synchronize() + # See the matching comment in compressed_cderi_j_only: this must + # run unconditionally, not only when num_devices > 1. + stream.synchronize() t1 = log.timer_debug1(f'store int3c2e on Device {device_id}', *t1) multi_gpu.run(proc, non_blocking=True) diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py b/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py index c3f16b267..4eecbd771 100644 --- a/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py +++ b/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py @@ -325,7 +325,7 @@ def test_contract_dm_kpts(): np.random.seed(9) auxvec = np.random.rand(auxcell.nao) vj = opt.contract_auxvec(opt.auxcell.apply_C_dot(auxvec), kpts=kpts) - ref = cp.einsum('kpqr,r->kpq', j3c, auxvec) + ref = cp.einsum('kpqr,r->kpq', j3c, cp.asarray(auxvec)) # auxvec is host data assert abs(vj - ref).max() < 1e-10 def test_int3c2e_batch_evaluation(): diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py b/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py index 5f4a84e26..6c4dfaf32 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py @@ -396,7 +396,7 @@ def test_uniform_grid_division_mode(self): test_coords = [] for frag_grids in grids.loop_grids(): test_coords.append(frag_grids.coords) - assert np.max(np.abs(frag_grids.weights - ref_weight)) < 1e-14 + assert cp.max(cp.abs(frag_grids.weights - ref_weight)) < 1e-14 test_coords = cp.vstack(test_coords).get() ref_coords = ref_coords[np.lexsort((ref_coords[:, 2], ref_coords[:, 1], ref_coords[:, 0])), :] diff --git a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py index cada001a1..c7c49b57b 100644 --- a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py +++ b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py @@ -40,7 +40,7 @@ def test_krhf_smearing(self): mf = cell.KRHF(kpts=cell.make_kpts([2,1,1])).to_gpu() mf = mf.smearing(0.1, 'fermi') nkpts = len(mf.kpts) - mo_energy_kpts = cp.array([cp.arange(nao)*.2+cp.cos(i+.5)*.1 for i in range(nkpts)]) + mo_energy_kpts = cp.array([cp.arange(nao)*.2+np.cos(i+.5)*.1 for i in range(nkpts)]) mf.get_occ(mo_energy_kpts) self.assertAlmostEqual(mf.entropy, 6.1656394960533021/2, 9) @@ -56,7 +56,7 @@ def test_kuhf_smearing(self): mf = cell.KUHF(kpts=cell.make_kpts([2,1,1])).to_gpu() mf = mf.smearing(0.1, 'fermi') nkpts = len(mf.kpts) - mo_energy_kpts = cp.array([cp.arange(nao)*.2+cp.cos(i+.5)*.1 for i in range(nkpts)]) + mo_energy_kpts = cp.array([cp.arange(nao)*.2+np.cos(i+.5)*.1 for i in range(nkpts)]) mo_energy_kpts = cp.array([mo_energy_kpts, mo_energy_kpts+cp.cos(mo_energy_kpts)*.02]) mf.get_occ(mo_energy_kpts) self.assertAlmostEqual(mf.entropy, 6.1803390081500869/2, 9) diff --git a/gpu4pyscf/pbc/tools/k2gamma.py b/gpu4pyscf/pbc/tools/k2gamma.py index f27685bfa..5d0215b34 100644 --- a/gpu4pyscf/pbc/tools/k2gamma.py +++ b/gpu4pyscf/pbc/tools/k2gamma.py @@ -82,7 +82,8 @@ def double_translation_indices(kmesh): tz = cp.array(translation_map(kmesh[2]), dtype=np.int32) idx = cp.ravel_multi_index([tx[:,None,None,:,None,None], ty[None,:,None,None,:,None], - tz[None,None,:,None,None,:]], kmesh) + tz[None,None,:,None,None,:]], + tuple(int(n) for n in kmesh)) nk = np.prod(kmesh) return idx.reshape(nk, nk) diff --git a/gpu4pyscf/qmmm/external_field.py b/gpu4pyscf/qmmm/external_field.py index 96c85a865..642d62713 100644 --- a/gpu4pyscf/qmmm/external_field.py +++ b/gpu4pyscf/qmmm/external_field.py @@ -53,7 +53,7 @@ def __init__(self, method, electric_field=None, origin=None): if electric_field is not None: electric_field = cp.asarray(electric_field) - assert type(electric_field) is cp.ndarray + assert isinstance(electric_field, cp.ndarray) assert electric_field.shape == (3,) self.electric_field = electric_field @@ -61,7 +61,7 @@ def __init__(self, method, electric_field=None, origin=None): origin = np.zeros(3) else: origin = cp.asarray(origin).get() - assert type(origin) is np.ndarray + assert isinstance(origin, np.ndarray) assert origin.shape == (3,) self.origin = origin