diff --git a/gpu4pyscf/pbc/dft/multigrid_v3.py b/gpu4pyscf/pbc/dft/multigrid_v3.py index c9aab23e6..efa330bea 100644 --- a/gpu4pyscf/pbc/dft/multigrid_v3.py +++ b/gpu4pyscf/pbc/dft/multigrid_v3.py @@ -500,6 +500,19 @@ def _get_L_bases(nimgs, a): L_bases = cp.array(np.hstack([Tx, Ty, Tz])) return L_bases +def _ke_max_f32(ke_max): + '''Largest float32 value not above ke_max. + + pair_ke and Ecut estimates are float32 while the bucket bounds derived + from ke_max are float64. Clamping with float32(ke_max) can round above + ke_max; pairs clamped that way fall outside the last bucket when its + upper bound equals ke_max exactly and are silently dropped. + ''' + v = np.float32(ke_max) + if float(v) > float(ke_max): + v = np.nextafter(v, np.float32(0)) + return v + def _estimate_Ecut_and_grid_ranges(ni, bas_ij_idx, ke_max, precision, xctype): '''Estimate the FFT energy cutoff and the spread of each orbital pair in real space''' @@ -507,8 +520,9 @@ def _estimate_Ecut_and_grid_ranges(ni, bas_ij_idx, ke_max, precision, xctype): # Some orbitals may require high Ecut, sometimes higher than ke_max. # Use ke_max to limit the highest Ecut. This ensures that these orbital # pairs are included in the last bucket in _partition_ke_for_fft. + ke_max32 = _ke_max_f32(ke_max) Ecut_by_shell = _estimate_fft_Ecut_per_shell(cell, precision) - Ecut_by_shell[Ecut_by_shell > ke_max] = ke_max + Ecut_by_shell[Ecut_by_shell > ke_max32] = ke_max32 Ecut_by_shell = cp.asarray(Ecut_by_shell, dtype=np.float32) npairs = len(bas_ij_idx) @@ -535,7 +549,7 @@ def _estimate_Ecut_and_grid_ranges(ni, bas_ij_idx, ke_max, precision, xctype): ctypes.c_int(li_inc), ctypes.c_int(lj_inc), ctypes.c_float(math.log(precision)), ctypes.c_float(undressed_threshold), - ctypes.c_float(ke_max)) + ctypes.c_float(ke_max32)) if err != 0: raise RuntimeError('grid range kernel failed') return pair_ke, grid_frac_ranges @@ -701,6 +715,14 @@ def _partition_ke_for_fft(ni, pair_idx, init_ke, ke_max, precision, xctype, log) # increase the mesh, causing the loop stuck mesh[mesh < 8] = 8 ke_lower, ke_upper = ke_upper, mesh_to_ke(a, mesh).min() + + # Every pair with a positive Ecut estimate must be in exactly one bucket. + # A pair left out is silently missing from the density and the gradients. + n_bucketed = sum(len(b) for bucket in buckets for b in bucket['bas_ij_cache']) + n_pairs = int((pair_ke > 0).sum()) + if n_bucketed != n_pairs: + raise RuntimeError(f'FFT bucket partition incomplete: {n_bucketed} of ' + f'{n_pairs} shell pairs assigned (ke_max={ke_max})') return buckets def _non_trivial_bvk_pairs(ni, precision): @@ -774,7 +796,7 @@ def _aft_Ecut_estimation(ni, bas_ij_idx, ke_max, precision, xctype='LDA'): ctypes.c_float(math.log(precision)), # Set the upper limit of Ecut. This ensures all high-Ecut orbital pairs # are handled by the last bucket in fft_buckets - ctypes.c_float(ke_max), + ctypes.c_float(_ke_max_f32(ke_max)), ctypes.c_int(is_mgga)) if err != 0: raise RuntimeError('Ecut kernel failed') diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py index 64b505f30..a98ea268e 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py @@ -1117,3 +1117,46 @@ def test_ne_derivatives(self): if __name__ == '__main__': print("Full Tests for multigrid v3") unittest.main() + + +class BucketPartition(unittest.TestCase): + def test_fft_bucket_completeness(self): + # Every shell pair with a positive Ecut estimate must land in exactly + # one FFT bucket. pair_ke is float32 and the bucket bounds are float64: + # when the ke ladder reaches the final mesh exactly, the last bucket's + # upper bound equals ke_cutoff, and pairs clamped to a float32 value + # rounded above ke_cutoff were left out of every bucket. For this + # cubic 5-bohr cell the affected cubic meshes include 26, 31 and 44. + from gpu4pyscf.pbc.dft.multigrid_v3 import ( + _non_trivial_bvk_pairs, _bvk_pairs_to_supmol_pairs, + _estimate_Ecut_and_grid_ranges) + cell = pyscf.M(atom='He 0 0 0', basis=[[0, (1., 1.)], [0, (60., 1.)]], + unit='B', precision=1e-10, a=np.eye(3)*5) + for n in (26, 31, 44): + ni = multigrid.MultiGridNumInt(cell) + ni.enable_aft = False + ni.mesh = [n] * 3 + ni.build(xctype='LDA') + scell = ni.sorted_cell + rad = scell.rcut / ni.bvkcell.vol**(1./3) + 1 + precision = scell.precision / (4*np.pi * rad**2 * 2) + bas_ij_idx = _non_trivial_bvk_pairs(ni, precision) + supmol_pairs = _bvk_pairs_to_supmol_pairs(ni, bas_ij_idx, precision, 'LDA') + pair_ke = _estimate_Ecut_and_grid_ranges( + ni, supmol_pairs, ni.ke_cutoff, precision, 'LDA')[0] + n_clamped = int((pair_ke >= np.float32(ni.ke_cutoff * (1 - 1e-6))).sum()) + self.assertGreater(n_clamped, 0, f'mesh {n}: no clamped pairs, test is void') + expected = int((pair_ke > 0).sum()) + got = sum(len(b) for bucket in ni.fft_buckets + for b in bucket['bas_ij_cache']) + self.assertEqual(got, expected, f'mesh {n}: {got} of {expected} pairs bucketed') + + def test_clamped_pairs_in_last_bucket(self): + # The clamped Ecut value must not exceed ke_cutoff in float64, so the + # clamped pairs satisfy pair_ke <= ke_upper for a last bucket whose + # upper bound equals ke_cutoff. + from gpu4pyscf.pbc.dft.multigrid_v3 import _ke_max_f32 + for ke_max in (218.12470933, 308.46562254, 124.91583061, 0.1, 1333.96245263): + v = _ke_max_f32(ke_max) + self.assertLessEqual(float(v), ke_max) + self.assertLess(ke_max - float(v), 1e-4 * ke_max)