pyfock 0.2.1__tar.gz → 0.2.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyfock
-Version: 0.2.1
+Version: 0.2.3
 Summary: A simplistic and efficient pure-python quantum chemistry library from Phys Whiz.
 Home-page: https://github.com/manassharma07/pyfock
 Download-URL: 
@@ -204,6 +204,7 @@ Dynamic: summary
 - ✅ **ASE Calculator**: Optional ASE interface (geometry optimization and the wider ASE ecosystem), using analytical forces by default
 - ✅ **XC integration grids**: Treutler-Ahlrichs radial and Lebedev angular grids with region-wise angular pruning and Becke partitioning (levels 0-9), built by a parallel Numba kernel or on the GPU (~6x faster, same grid); energies agree with PySCF to the SCF convergence threshold in our benchmarks; numgrid grids remain available as an alternative scheme
 - ✅ **SANO initial guess**: superposition of atomic natural-orbital densities (ANO-RCC-MB natural orbitals projected onto the calculation basis, no atomic SCF needed) as the default SCF starting point; it roughly halves the number of SCF iterations relative to the core-Hamiltonian guess and prunes the XC grid accurately
+- ✅ **Extrapolated starting densities in geometry optimizations**: `density_guess='extrapolate'` in the ASE calculator carries the converged densities of the last steps along with the atoms and extrapolates them to the new geometry (adapting the density projection/extrapolation code contributed by Prof. Vincenzo Barone), for 16-18 % fewer SCF iterations along the same optimization path
 - ✅ **Cross-Platform**: Works on Linux, macOS, and Windows
 
 ## Installation
@@ -404,6 +405,11 @@ Available names are `skala-1.1` (recommended), `skala-1.1-rev1`, `skala-1.1-rev0
 See [`examples/ex44_Skala_neural_functional.py`](examples/ex44_Skala_neural_functional.py) for a runnable
 version, with and without the dispersion correction.
 
+Two Kaggle notebooks walk through the integration and can be run in the browser:
+[Skala in PyFock vs. PySCF](https://www.kaggle.com/code/ducktape07/skala-pyfock-vs-pyscf) runs Skala on the CPU, checks PyFock's results against PySCF,
+compares the timings (PyFock is considerably faster) and finishes with a geometry optimization;
+[PyFock + Skala on a Tesla T4](https://www.kaggle.com/code/ducktape07/pyfock-skala-on-gpu-tesla-t4) shows how to run Skala on the GPU and benchmarks it.
+
 #### Getting the model
 
 `pip install "pyfock[skala]"` is all that is needed (see
@@ -439,10 +445,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
@@ -451,8 +457,23 @@ Consequences worth knowing:
   grid-translation term); its features are integrals over each atomic grid, so a frozen grid would cost
   ~1e-2 Ha/Bohr rather than the ~1e-4 it costs a meta-GGA.
 - **Dispersion is off by default** — pass `dispersion=True` to include it, see below.
+- **Do not use Skala with ghost atoms on the default grid** (counterpoise corrections). Every atom that
+  owns grid points is part of Skala's functional — its non-local features are anchored at the atoms of
+  the grid — and PyFock gives ghost atoms grid points like any other atom. A ghost atom has no density of
+  its own, which is far outside what the model was trained on, and its potential can bind electrons on
+  the ghost basis functions: for CH4@(H2O)20 in def2-TZVPPD the SCF of the water cage with a ghost CH4
+  inside diverged, its first HOMO-LUMO gap 0.06 eV against 10.2 eV for the complex. Conventional
+  functionals are unaffected. What does work is evaluating each fragment on the grid of its real atoms
+  only, so that the ghost atoms contribute basis functions and nothing else — at grid level 4 or finer,
+  `grids=Grids(real_atoms_mol, level=4)`, since ghost functions in a cavity are then sampled only by the
+  outer shells of their neighbours' grids, and at level 3 that was too coarse for Skala in the (H2O)24
+  cage (a spurious bound ghost state again). Use the same level for the complex.
+- **Memory grows with the grid points per model call**: about 26 kB per point on the CPU, autograd graph
+  included, so the default 250 000 points per call take about 10 GB. Lower
+  `dftObj.skala_max_points_per_chunk` (also used by `DFT_Grad`) on a machine with little memory; only
+  the noise described next changes.
 - **The model carries about 1e-9 Ha of its own numerical noise.** Presenting the network with differently
-  shaped batches (a different `max_points_per_chunk`) shifts the energy at that level. Within one
+  shaped batches (a different `skala_max_points_per_chunk`) shifts the energy at that level. Within one
   calculation the chunking is fixed, so the shift is systematic rather than random and SCF convergence to
   `1e-8` is unaffected — but do not expect two runs with different chunk sizes to agree bit for bit.
 
@@ -466,7 +487,7 @@ conformer ranking.
 
 D3 depends only on the atomic numbers, the coordinates and those damping parameters — not on the
 density. It never enters the Kohn-Sham matrix and cannot change the SCF, so PyFock evaluates it once
-after convergence and adds it to the total energy:
+after convergence and adds it to the total energy — and `DFT_Grad` adds its gradient to the forces:
 
 ```bash
 pip install dftd3
@@ -476,24 +497,29 @@ pip install dftd3
 dftObj = DFT(mol, basis, auxbasis, xc='skala-1.1', dispersion=True)
 energy, dmat = dftObj.scf()      # dispersion-corrected total energy
 print(dftObj.Edisp)              # the correction on its own, in Hartree
+forces = DFT_Grad(dftObj).calculate()['forces']   # dispersion-corrected too
 ```
 
 `dispersion=True` uses whatever the functional declares. For any other functional, name the
 parametrisation yourself — `dispersion='pbe'`, `dispersion='b3lyp'` — and tune the damping through
 `dftObj.dispersion_version` (default `'d3bj'`) and `dftObj.dispersion_atm` (default `False`).
 [`pyfock.Dispersion`](pyfock/Dispersion.py) also exposes `d3_energy` and `d3_energy_and_gradient`
-directly, the latter giving `dE_disp/dR` for forces. See
+directly, the latter giving `dE_disp/dR` — which `DFT_Grad` forces already contain, so do not add it to
+them a second time. See
 [`examples/ex43_D3_dispersion_correction.py`](examples/ex43_D3_dispersion_correction.py).
 
-The ASE calculator uses the same backend by default:
+The ASE calculator uses the same backend by default, and like `DFT` it takes Skala's parametrisation
+from the checkpoint, so Skala needs no `dispersion_kwargs`; any other functional must name its own:
 
 ```python
+PyFockCalculator(functional='skala-1.1', dispersion=True)   # D3(BJ) with b3lyp5, as Skala declares
 PyFockCalculator(functional='PBE', dispersion=True, dispersion_kwargs={'xc': 'pbe'})
 ```
 
 `torch-dftd` (`pip install pyfock[dispersion-gpu]`) remains available there for GPU runs, where it
 evaluates the correction on the device. Select it with `backend='torch-dftd'`, or implicitly by asking
-for a non-CPU device:
+for a non-CPU device (with Skala, `xc` can be left out here too; it becomes torch-dftd's `'b3-lyp'` with
+BJ damping):
 
 ```python
 dispersion_kwargs={'xc': 'pbe', 'backend': 'torch-dftd', 'device': 'cuda'}
@@ -565,27 +591,40 @@ 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).
+
+The [PyFock + Skala on a Tesla T4](https://www.kaggle.com/code/ducktape07/pyfock-skala-on-gpu-tesla-t4) Kaggle notebook shows the GPU setup end to end and
+benchmarks the PyFock + Skala integration on Kaggle's free T4.
 
 ### Initial Guess for the SCF
 
@@ -694,7 +733,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 +745,64 @@ forces = result["forces"]      # (natoms, 3) in Ha/Bohr
 gradient = result["gradient"]  # = -forces
 ```
 
+If the SCF applied a D3 correction (`dispersion=...`), the gradient includes its term, so it is always
+the gradient of the energy `scf()` returned, which is also `result["energy"]`.
+
+#### 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):
@@ -738,6 +835,17 @@ reads its thread count from the environment when numpy is first imported, so wit
 set `OMP_NUM_THREADS` at the top of your script rather than relying on `ncores` alone. See
 [`examples/ex45_ASE_geometry_optimization_with_Skala.py`](examples/ex45_ASE_geometry_optimization_with_Skala.py).
 
+#### Starting densities along an optimization path
+
+Every SCF of an optimization starts from the converged density of the previous step. With
+`density_guess='extrapolate'` the densities of the last five steps are instead carried along with
+the atoms — their occupied natural orbitals re-orthonormalized in the new overlap metric — and
+extrapolated to the new geometry with weights fitted to the geometries, then made N-representable.
+
+```python
+PyFockCalculator(functional="PBE", basis="def2-SVP", density_guess="extrapolate")   # default: 'previous'
+```
+
 #### Grid response
 
 By default the XC gradient treats the quadrature grid as fixed — its dependence on the nuclear
@@ -809,6 +917,8 @@ streamlit run app.py
 - [x] [Kaggle Notebook: Benchmarking PyFock against PySCF](https://www.kaggle.com/code/ducktape07/pyfock-vs-pyscf#10.-Timing-summary-table)
 - [x] [Kaggle Notebook: PyFock GPU Dynamic Precision](https://www.kaggle.com/code/ducktape07/pyfock-gpu-dynamic-precision)
 - [x] [Kaggle Notebook: PySCF vs. PyFock GPU Benchmark](https://www.kaggle.com/code/ducktape07/pyscf-vs-pyfock-gpu-benchmark)
+- [x] [Kaggle Notebook: Skala in PyFock vs. PySCF (CPU): accuracy, speed and geometry optimization](https://www.kaggle.com/code/ducktape07/skala-pyfock-vs-pyscf)
+- [x] [Kaggle Notebook: PyFock + Skala on the GPU (Tesla T4): setup and benchmark](https://www.kaggle.com/code/ducktape07/pyfock-skala-on-gpu-tesla-t4)
 
 
 ## Documentation
@@ -832,7 +942,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.3}/README.md

@@ -141,6 +141,7 @@
 - ✅ **ASE Calculator**: Optional ASE interface (geometry optimization and the wider ASE ecosystem), using analytical forces by default
 - ✅ **XC integration grids**: Treutler-Ahlrichs radial and Lebedev angular grids with region-wise angular pruning and Becke partitioning (levels 0-9), built by a parallel Numba kernel or on the GPU (~6x faster, same grid); energies agree with PySCF to the SCF convergence threshold in our benchmarks; numgrid grids remain available as an alternative scheme
 - ✅ **SANO initial guess**: superposition of atomic natural-orbital densities (ANO-RCC-MB natural orbitals projected onto the calculation basis, no atomic SCF needed) as the default SCF starting point; it roughly halves the number of SCF iterations relative to the core-Hamiltonian guess and prunes the XC grid accurately
+- ✅ **Extrapolated starting densities in geometry optimizations**: `density_guess='extrapolate'` in the ASE calculator carries the converged densities of the last steps along with the atoms and extrapolates them to the new geometry (adapting the density projection/extrapolation code contributed by Prof. Vincenzo Barone), for 16-18 % fewer SCF iterations along the same optimization path
 - ✅ **Cross-Platform**: Works on Linux, macOS, and Windows
 
 ## Installation
@@ -341,6 +342,11 @@ Available names are `skala-1.1` (recommended), `skala-1.1-rev1`, `skala-1.1-rev0
 See [`examples/ex44_Skala_neural_functional.py`](examples/ex44_Skala_neural_functional.py) for a runnable
 version, with and without the dispersion correction.
 
+Two Kaggle notebooks walk through the integration and can be run in the browser:
+[Skala in PyFock vs. PySCF](https://www.kaggle.com/code/ducktape07/skala-pyfock-vs-pyscf) runs Skala on the CPU, checks PyFock's results against PySCF,
+compares the timings (PyFock is considerably faster) and finishes with a geometry optimization;
+[PyFock + Skala on a Tesla T4](https://www.kaggle.com/code/ducktape07/pyfock-skala-on-gpu-tesla-t4) shows how to run Skala on the GPU and benchmarks it.
+
 #### Getting the model
 
 `pip install "pyfock[skala]"` is all that is needed (see
@@ -376,10 +382,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
@@ -388,8 +394,23 @@ Consequences worth knowing:
   grid-translation term); its features are integrals over each atomic grid, so a frozen grid would cost
   ~1e-2 Ha/Bohr rather than the ~1e-4 it costs a meta-GGA.
 - **Dispersion is off by default** — pass `dispersion=True` to include it, see below.
+- **Do not use Skala with ghost atoms on the default grid** (counterpoise corrections). Every atom that
+  owns grid points is part of Skala's functional — its non-local features are anchored at the atoms of
+  the grid — and PyFock gives ghost atoms grid points like any other atom. A ghost atom has no density of
+  its own, which is far outside what the model was trained on, and its potential can bind electrons on
+  the ghost basis functions: for CH4@(H2O)20 in def2-TZVPPD the SCF of the water cage with a ghost CH4
+  inside diverged, its first HOMO-LUMO gap 0.06 eV against 10.2 eV for the complex. Conventional
+  functionals are unaffected. What does work is evaluating each fragment on the grid of its real atoms
+  only, so that the ghost atoms contribute basis functions and nothing else — at grid level 4 or finer,
+  `grids=Grids(real_atoms_mol, level=4)`, since ghost functions in a cavity are then sampled only by the
+  outer shells of their neighbours' grids, and at level 3 that was too coarse for Skala in the (H2O)24
+  cage (a spurious bound ghost state again). Use the same level for the complex.
+- **Memory grows with the grid points per model call**: about 26 kB per point on the CPU, autograd graph
+  included, so the default 250 000 points per call take about 10 GB. Lower
+  `dftObj.skala_max_points_per_chunk` (also used by `DFT_Grad`) on a machine with little memory; only
+  the noise described next changes.
 - **The model carries about 1e-9 Ha of its own numerical noise.** Presenting the network with differently
-  shaped batches (a different `max_points_per_chunk`) shifts the energy at that level. Within one
+  shaped batches (a different `skala_max_points_per_chunk`) shifts the energy at that level. Within one
   calculation the chunking is fixed, so the shift is systematic rather than random and SCF convergence to
   `1e-8` is unaffected — but do not expect two runs with different chunk sizes to agree bit for bit.
 
@@ -403,7 +424,7 @@ conformer ranking.
 
 D3 depends only on the atomic numbers, the coordinates and those damping parameters — not on the
 density. It never enters the Kohn-Sham matrix and cannot change the SCF, so PyFock evaluates it once
-after convergence and adds it to the total energy:
+after convergence and adds it to the total energy — and `DFT_Grad` adds its gradient to the forces:
 
 ```bash
 pip install dftd3
@@ -413,24 +434,29 @@ pip install dftd3
 dftObj = DFT(mol, basis, auxbasis, xc='skala-1.1', dispersion=True)
 energy, dmat = dftObj.scf()      # dispersion-corrected total energy
 print(dftObj.Edisp)              # the correction on its own, in Hartree
+forces = DFT_Grad(dftObj).calculate()['forces']   # dispersion-corrected too
 ```
 
 `dispersion=True` uses whatever the functional declares. For any other functional, name the
 parametrisation yourself — `dispersion='pbe'`, `dispersion='b3lyp'` — and tune the damping through
 `dftObj.dispersion_version` (default `'d3bj'`) and `dftObj.dispersion_atm` (default `False`).
 [`pyfock.Dispersion`](pyfock/Dispersion.py) also exposes `d3_energy` and `d3_energy_and_gradient`
-directly, the latter giving `dE_disp/dR` for forces. See
+directly, the latter giving `dE_disp/dR` — which `DFT_Grad` forces already contain, so do not add it to
+them a second time. See
 [`examples/ex43_D3_dispersion_correction.py`](examples/ex43_D3_dispersion_correction.py).
 
-The ASE calculator uses the same backend by default:
+The ASE calculator uses the same backend by default, and like `DFT` it takes Skala's parametrisation
+from the checkpoint, so Skala needs no `dispersion_kwargs`; any other functional must name its own:
 
 ```python
+PyFockCalculator(functional='skala-1.1', dispersion=True)   # D3(BJ) with b3lyp5, as Skala declares
 PyFockCalculator(functional='PBE', dispersion=True, dispersion_kwargs={'xc': 'pbe'})
 ```
 
 `torch-dftd` (`pip install pyfock[dispersion-gpu]`) remains available there for GPU runs, where it
 evaluates the correction on the device. Select it with `backend='torch-dftd'`, or implicitly by asking
-for a non-CPU device:
+for a non-CPU device (with Skala, `xc` can be left out here too; it becomes torch-dftd's `'b3-lyp'` with
+BJ damping):
 
 ```python
 dispersion_kwargs={'xc': 'pbe', 'backend': 'torch-dftd', 'device': 'cuda'}
@@ -502,27 +528,40 @@ 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).
+
+The [PyFock + Skala on a Tesla T4](https://www.kaggle.com/code/ducktape07/pyfock-skala-on-gpu-tesla-t4) Kaggle notebook shows the GPU setup end to end and
+benchmarks the PyFock + Skala integration on Kaggle's free T4.
 
 ### Initial Guess for the SCF
 
@@ -631,7 +670,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 +682,64 @@ forces = result["forces"]      # (natoms, 3) in Ha/Bohr
 gradient = result["gradient"]  # = -forces
 ```
 
+If the SCF applied a D3 correction (`dispersion=...`), the gradient includes its term, so it is always
+the gradient of the energy `scf()` returned, which is also `result["energy"]`.
+
+#### 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):
@@ -675,6 +772,17 @@ reads its thread count from the environment when numpy is first imported, so wit
 set `OMP_NUM_THREADS` at the top of your script rather than relying on `ncores` alone. See
 [`examples/ex45_ASE_geometry_optimization_with_Skala.py`](examples/ex45_ASE_geometry_optimization_with_Skala.py).
 
+#### Starting densities along an optimization path
+
+Every SCF of an optimization starts from the converged density of the previous step. With
+`density_guess='extrapolate'` the densities of the last five steps are instead carried along with
+the atoms — their occupied natural orbitals re-orthonormalized in the new overlap metric — and
+extrapolated to the new geometry with weights fitted to the geometries, then made N-representable.
+
+```python
+PyFockCalculator(functional="PBE", basis="def2-SVP", density_guess="extrapolate")   # default: 'previous'
+```
+
 #### Grid response
 
 By default the XC gradient treats the quadrature grid as fixed — its dependence on the nuclear
@@ -746,6 +854,8 @@ streamlit run app.py
 - [x] [Kaggle Notebook: Benchmarking PyFock against PySCF](https://www.kaggle.com/code/ducktape07/pyfock-vs-pyscf#10.-Timing-summary-table)
 - [x] [Kaggle Notebook: PyFock GPU Dynamic Precision](https://www.kaggle.com/code/ducktape07/pyfock-gpu-dynamic-precision)
 - [x] [Kaggle Notebook: PySCF vs. PyFock GPU Benchmark](https://www.kaggle.com/code/ducktape07/pyscf-vs-pyfock-gpu-benchmark)
+- [x] [Kaggle Notebook: Skala in PyFock vs. PySCF (CPU): accuracy, speed and geometry optimization](https://www.kaggle.com/code/ducktape07/skala-pyfock-vs-pyscf)
+- [x] [Kaggle Notebook: PyFock + Skala on the GPU (Tesla T4): setup and benchmark](https://www.kaggle.com/code/ducktape07/pyfock-skala-on-gpu-tesla-t4)
 
 
 ## Documentation
@@ -769,7 +879,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