pyfock 0.2.1__tar.gz → 0.2.2__tar.gz

{pyfock-0.2.1/pyfock.egg-info → pyfock-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyfock
-Version: 0.2.1
+Version: 0.2.2
 Summary: A simplistic and efficient pure-python quantum chemistry library from Phys Whiz.
 Home-page: https://github.com/manassharma07/pyfock
 Download-URL: 
@@ -439,10 +439,10 @@ Consequences worth knowing:
   weights and the raw single-atom weights.
 - **Do not density-prune the grid.** The model's non-local features are integrals over each atom's full
   grid.
-- **The GPU runs the model, the CPU runs the rest.** Pass `skala_gpu=True` and the neural functional is
-  evaluated on the device while the integrals, the Coulomb term and the assembly of Vxc stay on the CPU
-  — see [Skala on the GPU](#skala-on-the-gpu). PyFock's *full* GPU SCF path (`use_gpu=True`) does not
-  support Skala and raises a clear error; it would need a CuPy↔Torch bridge that is not wired up yet.
+- **The GPU is supported**, either way round: `use_gpu=True` runs the whole SCF on the device, and
+  `skala_gpu=True` moves only the neural functional there and leaves the rest on the CPU. See
+  [Skala on the GPU](#skala-on-the-gpu). Forces follow the SCF, so `use_gpu=True` puts both halves of
+  a geometry step on the device.
 - **Closed-shell (restricted) only**, like the rest of PyFock's DFT.
 - **Analytical nuclear gradients are supported**, so Skala works for geometry optimization like any
   other functional — see
@@ -565,27 +565,37 @@ while the semilocal functionals also carry AO work that grows with the basis.
 
 #### Skala on the GPU
 
-Skala is a PyTorch model, so it can be evaluated on the GPU without PyFock's own CuPy machinery. Pass
-`skala_gpu=True` and the model runs on the device while everything else — the integrals, the Coulomb
-term and the assembly of Vxc from the model's cotangents — stays on the CPU:
+There are two ways to use a GPU, and they compose with the rest of PyFock exactly as they do for any
+other functional. `use_gpu=True` runs the whole SCF on the device — integrals, Coulomb term, AO values,
+density, the model and the Vxc assembly:
 
 ```python
-dftObj = DFT(mol, basis, auxbasis, xc='skala-1.1', dispersion=True, skala_gpu=True)
+dftObj = DFT(mol, basis, auxbasis, xc='skala-1.1', dispersion=True, use_gpu=True)
 ```
 
-It works the same way through the ASE calculator, and forces come along with it:
+`skala_gpu=True` instead moves only the neural functional and leaves the rest of the SCF on the CPU.
+That is useful when the GPU path is unavailable or unwanted, since the model alone is most of Skala's
+cost:
 
 ```python
-atoms.calc = PyFockCalculator(functional='skala-1.1', skala_gpu=True, dispersion=True)
+dftObj = DFT(mol, basis, auxbasis, xc='skala-1.1', dispersion=True, skala_gpu=True)
+atoms.calc = PyFockCalculator(functional='skala-1.1', skala_gpu=True, dispersion=True)  # same in ASE
 ```
 
-All this needs is a CUDA build of PyTorch (`pip install torch --index-url
+Both need a CUDA build of PyTorch (`pip install torch --index-url
 https://download.pytorch.org/whl/cu128` for current cards); PyFock downloads the CUDA checkpoint, which
-is a separate file from the CPU one. `use_gpu=True`, the flag for PyFock's *full* GPU SCF path, is a
-different thing and still raises a clear error for Skala.
+is a separate file from the CPU one. `use_gpu=True` additionally needs CuPy, as it does for every
+functional. Skala is a float64 model, so `dynamic_precision` is switched off for it automatically.
+
+Non-locality is the only thing that changes on the device: Skala cannot be folded into the per-block
+loop of `eval_xc_3_cupy`, where each block calls the functional on its own points, so it gets the same
+three-pass driver as on the CPU (`Integrals.eval_xc_skala_cupy`). Everything inside those passes is the
+semilocal GPU code.
 
 `benchmarks_tests/benchmark_skala_gpu.py` measures the speedup and the CPU/GPU agreement, for
-single points and for geometry optimizations.
+single points and for geometry optimizations. Forces run on the device too -- `DFT_Grad` follows the
+SCF -- so a `use_gpu=True` geometry optimization never returns to the host for a step; see
+[Forces on the GPU](#forces-on-the-gpu).
 
 ### Initial Guess for the SCF
 
@@ -694,7 +704,7 @@ benchmark commands, memory accounting, and limitations of the surrounding driver
 ### Analytical Forces & Geometry Optimization
 
 After a converged DFT calculation, analytical nuclear gradients (and forces)
-are available directly via `DFT_Grad` (density fitting; LDA/GGA/meta-GGA and Skala; CPU):
+are available directly via `DFT_Grad` (density fitting; LDA/GGA/meta-GGA and Skala; CPU or GPU):
 
 ```python
 from pyfock import DFT_Grad
@@ -706,6 +716,61 @@ forces = result["forces"]      # (natoms, 3) in Ha/Bohr
 gradient = result["gradient"]  # = -forces
 ```
 
+#### Forces on the GPU
+
+`DFT_Grad` follows the SCF: converge with `use_gpu=True` and the gradient runs on the device too, so
+a geometry step never returns to the host. Pass `use_gpu` explicitly to override that in either
+direction:
+
+```python
+DFT_Grad(dftObj).calculate()                  # follows dftObj.use_gpu
+DFT_Grad(dftObj, use_gpu=True).calculate()    # device gradient after a CPU SCF
+DFT_Grad(dftObj, use_gpu=False).calculate()   # host gradient after a GPU SCF
+```
+
+Every term has a device implementation except the ECP one, which stays on the CPU. The 3c2e
+derivative, the fitting coefficients and the nuclear attraction are shell-triplet CUDA kernels; the
+XC term (semilocal and Skala, with or without the grid response) is CuPy over the grid blocks.
+
+Gradient wall time only, from the same converged SCF (RTX 5070 against 8 CPU cores, PBE,
+`def2-universal-jfit`, level-3 grid):
+
+| molecule | atoms | def2-SVP CPU | def2-SVP GPU | speed-up | def2-TZVP speed-up |
+|---|---|---|---|---|---|
+| H2O | 3 | 0.08 s | 0.07 s | 1.0x | 1.5x |
+| Benzene | 12 | 1.09 s | 0.23 s | 4.7x | 6.7x |
+| Caffeine | 24 | 4.57 s | 0.67 s | 6.8x | 9.5x |
+| Serotonin | 25 | 3.96 s | 0.64 s | 6.2x | 8.8x |
+| Adenine-Thymine | 30 | 7.56 s | 0.94 s | 8.1x | |
+| Cholesterol | 74 | 32.0 s | 3.54 s | 9.0x | |
+| C60 | 60 | 92.2 s | 10.7 s | 8.6x | |
+
+Small molecules do not fill the device -- water breaks even -- and the gain arrives once the shell
+triplets outnumber the cores by enough, then grows with the basis. The device and host gradients
+agree to 1e-12 - 1e-11 Ha/Bohr on forces of order 1e-1, which is round-off in the contraction, and
+against finite differences (grid response on, 1e-3 Bohr step) the GPU gradient is right to 1.0e-6
+Ha/Bohr for water and 2.6e-7 for benzene -- the truncation floor of the reference, not of the
+gradient.
+
+Two functionals sit apart from that picture, for reasons that are properties of the functional
+rather than of the gradient.
+
+*r2SCAN* reaches only 2.5 - 2.7x. Its native implementation has no analytic potential: it
+finite-differences its own energy expression, seven evaluations of ~160 elementwise operations. That
+is launch-bound on the device -- the same call costs 24 ms on a 20k-point block and 21 ms on a
+120k-point one -- so it speeds up ~3.6x while everything around it speeds up ~7x, and Amdahl does the
+rest. The same finite differencing puts the device/host agreement at ~5e-9 rather than ~1e-12.
+
+*Skala* reaches ~5x on caffeine, but its per-molecule ratios are worth reading as +-50%: the model
+call is a PyTorch forward/backward whose cost varies by an order of magnitude between identical
+calls, and with ~7 GB of the card reserved by PyTorch the CuPy pool takes several gradients to reach
+a steady state. Serotonin measures 3.0x inside the benchmark suite and 4.6x on its own. Time it with
+`--repeats`, and treat the first two gradients in a process as warm-up.
+
+`benchmarks_tests/benchmark_DFT_gradients_gpu.py` reproduces all of the above, including the
+finite-difference check; `tests/test_grad_gpu.py` checks each device routine against its CPU
+counterpart.
+
 For geometry optimization, use the ASE calculator (requires `ase`). It uses
 the analytical forces by default and falls back to finite differences only for
 configurations the analytical gradients do not yet cover (e.g. HF, no DF):
@@ -832,7 +897,7 @@ streamlit run app.py
 - [ ] Electron dynamics & Excited state calculations (RT-TDDFT)
 - [ ] Periodic boundary conditions
 - [x] Hybrid functionals with exact exchange (native B3LYP/PBE0 and LibXC hybrids, RI-K via DF_algo=11; CPU)
-- [x] Skala neural exchange-correlation functional, with analytical gradients (CPU; model on GPU via `skala_gpu=True`)
+- [x] Skala neural exchange-correlation functional, on CPU and GPU, with analytical gradients (CPU)
 - [x] Grid-response XC gradients (Becke weight derivatives, `Grids.becke_weight_gradient`)
 - [ ] Multi-GPU parallelization
 - [ ] Basis set optimization tools

{pyfock-0.2.1 → pyfock-0.2.2}/README.md

@@ -376,10 +376,10 @@ Consequences worth knowing:
   weights and the raw single-atom weights.
 - **Do not density-prune the grid.** The model's non-local features are integrals over each atom's full
   grid.
-- **The GPU runs the model, the CPU runs the rest.** Pass `skala_gpu=True` and the neural functional is
-  evaluated on the device while the integrals, the Coulomb term and the assembly of Vxc stay on the CPU
-  — see [Skala on the GPU](#skala-on-the-gpu). PyFock's *full* GPU SCF path (`use_gpu=True`) does not
-  support Skala and raises a clear error; it would need a CuPy↔Torch bridge that is not wired up yet.
+- **The GPU is supported**, either way round: `use_gpu=True` runs the whole SCF on the device, and
+  `skala_gpu=True` moves only the neural functional there and leaves the rest on the CPU. See
+  [Skala on the GPU](#skala-on-the-gpu). Forces follow the SCF, so `use_gpu=True` puts both halves of
+  a geometry step on the device.
 - **Closed-shell (restricted) only**, like the rest of PyFock's DFT.
 - **Analytical nuclear gradients are supported**, so Skala works for geometry optimization like any
   other functional — see
@@ -502,27 +502,37 @@ while the semilocal functionals also carry AO work that grows with the basis.
 
 #### Skala on the GPU
 
-Skala is a PyTorch model, so it can be evaluated on the GPU without PyFock's own CuPy machinery. Pass
-`skala_gpu=True` and the model runs on the device while everything else — the integrals, the Coulomb
-term and the assembly of Vxc from the model's cotangents — stays on the CPU:
+There are two ways to use a GPU, and they compose with the rest of PyFock exactly as they do for any
+other functional. `use_gpu=True` runs the whole SCF on the device — integrals, Coulomb term, AO values,
+density, the model and the Vxc assembly:
 
 ```python
-dftObj = DFT(mol, basis, auxbasis, xc='skala-1.1', dispersion=True, skala_gpu=True)
+dftObj = DFT(mol, basis, auxbasis, xc='skala-1.1', dispersion=True, use_gpu=True)
 ```
 
-It works the same way through the ASE calculator, and forces come along with it:
+`skala_gpu=True` instead moves only the neural functional and leaves the rest of the SCF on the CPU.
+That is useful when the GPU path is unavailable or unwanted, since the model alone is most of Skala's
+cost:
 
 ```python
-atoms.calc = PyFockCalculator(functional='skala-1.1', skala_gpu=True, dispersion=True)
+dftObj = DFT(mol, basis, auxbasis, xc='skala-1.1', dispersion=True, skala_gpu=True)
+atoms.calc = PyFockCalculator(functional='skala-1.1', skala_gpu=True, dispersion=True)  # same in ASE
 ```
 
-All this needs is a CUDA build of PyTorch (`pip install torch --index-url
+Both need a CUDA build of PyTorch (`pip install torch --index-url
 https://download.pytorch.org/whl/cu128` for current cards); PyFock downloads the CUDA checkpoint, which
-is a separate file from the CPU one. `use_gpu=True`, the flag for PyFock's *full* GPU SCF path, is a
-different thing and still raises a clear error for Skala.
+is a separate file from the CPU one. `use_gpu=True` additionally needs CuPy, as it does for every
+functional. Skala is a float64 model, so `dynamic_precision` is switched off for it automatically.
+
+Non-locality is the only thing that changes on the device: Skala cannot be folded into the per-block
+loop of `eval_xc_3_cupy`, where each block calls the functional on its own points, so it gets the same
+three-pass driver as on the CPU (`Integrals.eval_xc_skala_cupy`). Everything inside those passes is the
+semilocal GPU code.
 
 `benchmarks_tests/benchmark_skala_gpu.py` measures the speedup and the CPU/GPU agreement, for
-single points and for geometry optimizations.
+single points and for geometry optimizations. Forces run on the device too -- `DFT_Grad` follows the
+SCF -- so a `use_gpu=True` geometry optimization never returns to the host for a step; see
+[Forces on the GPU](#forces-on-the-gpu).
 
 ### Initial Guess for the SCF
 
@@ -631,7 +641,7 @@ benchmark commands, memory accounting, and limitations of the surrounding driver
 ### Analytical Forces & Geometry Optimization
 
 After a converged DFT calculation, analytical nuclear gradients (and forces)
-are available directly via `DFT_Grad` (density fitting; LDA/GGA/meta-GGA and Skala; CPU):
+are available directly via `DFT_Grad` (density fitting; LDA/GGA/meta-GGA and Skala; CPU or GPU):
 
 ```python
 from pyfock import DFT_Grad
@@ -643,6 +653,61 @@ forces = result["forces"]      # (natoms, 3) in Ha/Bohr
 gradient = result["gradient"]  # = -forces
 ```
 
+#### Forces on the GPU
+
+`DFT_Grad` follows the SCF: converge with `use_gpu=True` and the gradient runs on the device too, so
+a geometry step never returns to the host. Pass `use_gpu` explicitly to override that in either
+direction:
+
+```python
+DFT_Grad(dftObj).calculate()                  # follows dftObj.use_gpu
+DFT_Grad(dftObj, use_gpu=True).calculate()    # device gradient after a CPU SCF
+DFT_Grad(dftObj, use_gpu=False).calculate()   # host gradient after a GPU SCF
+```
+
+Every term has a device implementation except the ECP one, which stays on the CPU. The 3c2e
+derivative, the fitting coefficients and the nuclear attraction are shell-triplet CUDA kernels; the
+XC term (semilocal and Skala, with or without the grid response) is CuPy over the grid blocks.
+
+Gradient wall time only, from the same converged SCF (RTX 5070 against 8 CPU cores, PBE,
+`def2-universal-jfit`, level-3 grid):
+
+| molecule | atoms | def2-SVP CPU | def2-SVP GPU | speed-up | def2-TZVP speed-up |
+|---|---|---|---|---|---|
+| H2O | 3 | 0.08 s | 0.07 s | 1.0x | 1.5x |
+| Benzene | 12 | 1.09 s | 0.23 s | 4.7x | 6.7x |
+| Caffeine | 24 | 4.57 s | 0.67 s | 6.8x | 9.5x |
+| Serotonin | 25 | 3.96 s | 0.64 s | 6.2x | 8.8x |
+| Adenine-Thymine | 30 | 7.56 s | 0.94 s | 8.1x | |
+| Cholesterol | 74 | 32.0 s | 3.54 s | 9.0x | |
+| C60 | 60 | 92.2 s | 10.7 s | 8.6x | |
+
+Small molecules do not fill the device -- water breaks even -- and the gain arrives once the shell
+triplets outnumber the cores by enough, then grows with the basis. The device and host gradients
+agree to 1e-12 - 1e-11 Ha/Bohr on forces of order 1e-1, which is round-off in the contraction, and
+against finite differences (grid response on, 1e-3 Bohr step) the GPU gradient is right to 1.0e-6
+Ha/Bohr for water and 2.6e-7 for benzene -- the truncation floor of the reference, not of the
+gradient.
+
+Two functionals sit apart from that picture, for reasons that are properties of the functional
+rather than of the gradient.
+
+*r2SCAN* reaches only 2.5 - 2.7x. Its native implementation has no analytic potential: it
+finite-differences its own energy expression, seven evaluations of ~160 elementwise operations. That
+is launch-bound on the device -- the same call costs 24 ms on a 20k-point block and 21 ms on a
+120k-point one -- so it speeds up ~3.6x while everything around it speeds up ~7x, and Amdahl does the
+rest. The same finite differencing puts the device/host agreement at ~5e-9 rather than ~1e-12.
+
+*Skala* reaches ~5x on caffeine, but its per-molecule ratios are worth reading as +-50%: the model
+call is a PyTorch forward/backward whose cost varies by an order of magnitude between identical
+calls, and with ~7 GB of the card reserved by PyTorch the CuPy pool takes several gradients to reach
+a steady state. Serotonin measures 3.0x inside the benchmark suite and 4.6x on its own. Time it with
+`--repeats`, and treat the first two gradients in a process as warm-up.
+
+`benchmarks_tests/benchmark_DFT_gradients_gpu.py` reproduces all of the above, including the
+finite-difference check; `tests/test_grad_gpu.py` checks each device routine against its CPU
+counterpart.
+
 For geometry optimization, use the ASE calculator (requires `ase`). It uses
 the analytical forces by default and falls back to finite differences only for
 configurations the analytical gradients do not yet cover (e.g. HF, no DF):
@@ -769,7 +834,7 @@ streamlit run app.py
 - [ ] Electron dynamics & Excited state calculations (RT-TDDFT)
 - [ ] Periodic boundary conditions
 - [x] Hybrid functionals with exact exchange (native B3LYP/PBE0 and LibXC hybrids, RI-K via DF_algo=11; CPU)
-- [x] Skala neural exchange-correlation functional, with analytical gradients (CPU; model on GPU via `skala_gpu=True`)
+- [x] Skala neural exchange-correlation functional, on CPU and GPU, with analytical gradients (CPU)
 - [x] Grid-response XC gradients (Becke weight derivatives, `Grids.becke_weight_gradient`)
 - [ ] Multi-GPU parallelization
 - [ ] Basis set optimization tools

{pyfock-0.2.1 → pyfock-0.2.2}/pyfock/DFT.py

@@ -1221,15 +1221,14 @@ class DFT:
         # `xc` therefore stays the name string and the LibXC resolution below is skipped.
         skala = None
         if XC.is_skala(xc):
-            if self.use_gpu:
-                raise NotImplementedError(
-                    "PyFock's full GPU SCF path does not support Skala yet (the Coulomb and Vxc "
-                    'assembly would need a CuPy<->Torch bridge). Use use_gpu=False with '
-                    'skala_gpu=True to run the neural functional itself on the GPU, which is where '
-                    'nearly all of its cost is.')
             xc = XC.canonical_skala_name(xc)
-            skala = XC.load_skala(xc, use_gpu=bool(self.skala_gpu))
-            if self.skala_gpu:
+            # On a GPU run the model goes to the device along with everything else; on a CPU run
+            # skala_gpu moves just the model, which is where nearly all of Skala's cost sits.
+            skala = XC.load_skala(xc, use_gpu=bool(self.use_gpu) or bool(self.skala_gpu))
+            if self.use_gpu:
+                print('\nSkala will be evaluated on the GPU, like the rest of the SCF '
+                      '(Integrals.eval_xc_skala_cupy).', flush=True)
+            elif self.skala_gpu:
                 print('\nSkala will be evaluated on the GPU (skala_gpu=True); the rest of the SCF '
                       'runs on the CPU.', flush=True)
         self.skala = skala
@@ -1372,6 +1371,10 @@ class DFT:
                 # (H2O 8 -> 19, decane 9 -> 19 iterations, ~2x wall time), with energies that vary by 1e-9 Ha
                 # between runs. The default therefore switches itself off for them.
                 dynamic_precision = self.dynamic_precision
+                if dynamic_precision and skala is not None:
+                    print('\n\nSkala is a float64 model, so the XC term cannot be evaluated in single')
+                    print('precision; dynamic precision is disabled.', flush=True)
+                    dynamic_precision = False
                 if dynamic_precision and xc != 'HF' and XC.xc_semilocal_family(xc, self.use_libxc) == 4:
                     print('\n\nDynamic precision is not used for meta-GGA functionals: single precision resolves')
                     print('tau too poorly and the SCF then needs about twice as many iterations.', flush=True)
@@ -2036,14 +2039,24 @@ class DFT:
                 print('\n\n\n\nXC algo', XC_algo)
                 if skala is not None:
                     # Skala is non-local over each atomic grid, so it cannot be folded into the blocked
-                    # semilocal loop; it gets its own three-pass driver (see Integrals.eval_xc_skala).
-                    Exc, Vxc = Integrals.eval_xc_skala(basis, dmat, grids, skala, ncores=ncores,
-                                                       blocksize=blocksize,
-                                                       list_nonzero_indices=list_nonzero_indices,
-                                                       count_nonzero_indices=count_nonzero_indices,
-                                                       list_ao_values=list_ao_values,
-                                                       list_ao_grad_values=list_ao_grad_values,
-                                                       debug=debug)
+                    # semilocal loop, where every block calls the functional on its own points; it gets
+                    # its own three-pass driver instead.
+                    if self.use_gpu:
+                        Exc, Vxc = Integrals.eval_xc_skala_cupy(
+                            basis, dmat_cp, grids, skala, blocksize=blocksize,
+                            list_nonzero_indices=list_nonzero_indices,
+                            count_nonzero_indices=count_nonzero_indices,
+                            list_ao_values=list_ao_values, list_ao_grad_values=list_ao_grad_values,
+                            debug=debug, threads_per_block=threads_per_block)
+                        Vxc = cp.asarray(Vxc, dtype=cp.float64)
+                    else:
+                        Exc, Vxc = Integrals.eval_xc_skala(basis, dmat, grids, skala, ncores=ncores,
+                                                           blocksize=blocksize,
+                                                           list_nonzero_indices=list_nonzero_indices,
+                                                           count_nonzero_indices=count_nonzero_indices,
+                                                           list_ao_values=list_ao_values,
+                                                           list_ao_grad_values=list_ao_grad_values,
+                                                           debug=debug)
                 elif not self.use_gpu:
                     if XC_algo==1:
                         # Much slower than JOBLIB version