multigrid_v3._partition_ke_for_fft - prevent shell pairs whose Ecut estimate was clamped to ke_cutoff from falling outside FFT buckets by clamping the estimate to the largest f32 not above ke_cutoff

gpu4pyscf/pbc/dft/multigrid_v3.py

@@ -500,15 +500,29 @@ 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'''
     cell = ni.sorted_cell
     # 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')

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)