Improve Vxc Screening

[Jens (JensWehner)]

Branch changes compared with main

This summary compares the current screening branch with origin/main. The local
main branch is one commit behind origin/main; using it literally would also include
an unrelated upstream allocator commit. Uncommitted files are excluded.

1. Automatic AO screening for large systems

   XC evaluation now switches from dense AO contractions to screened contractions when
   the functional supports atom-aligned grids and mol.nao_nr() exceeds PySCF's
   SWITCH_SIZE. Smaller systems and unsupported functionals retain the dense path.

2. Spatially localized integration grids

   Grid points are recursively grouped into spatial blocks so each block touches fewer
   basis functions. Forward and inverse permutations restore the atom-major ordering
   expected by Skala.

3. Reusable, backend-aware screening metadata

   A SpatialGridLayout is cached on the original grid and shared across SkalaNumInt
   instances. CPU evaluation rebuilds PySCF shell masks, while GPU evaluation invalidates
   GPU4PySCF's AO-index cache for the reordered grid.

4. Sparse blockwise AO autograd

   New custom PyTorch autograd functions evaluate only active AO submatrices and scatter
   derivatives back into full density-matrix shapes. They handle CPU shell masks, GPU AO
   sorting, empty blocks, VJPs, JVPs, and higher-order differentiation.

5. Screened energy and potential integration

   The screened first-order path evaluates raw AO features globally once, processes the
   neural functional in memory-sized atom chunks, collects feature cotangents, and
   performs one final sparse VJP to obtain the XC potential.

6. Screened response and Hessian-vector products

   Response evaluation computes a screened feature JVP, chunk-local model Hessian
   actions, and a final sparse AO VJP. XCIntegrator returns the XC-only response, while
   SkalaNumInt adds the Coulomb response required by PySCF.

7. AO traversal and model chunking are decoupled

   AO screening traverses the spatial grid independently of neural-model batching. Model
   chunks preserve complete atomic grids and are sized using available memory after
   reserving space for global feature and cotangent buffers.

8. PySCF integration is split into focused modules

   The former large features.py module is now a thin public feature generator. Feature
   policy, feature mathematics, AO evaluation, screening, model chunking, and XC
   integration have separate owners; SkalaNumInt primarily adapts NumPy and CuPy data
   to Torch.

9. Feature support is deliberately narrowed

   The unsupported nonlinear kinetic-tensor features ked_var and ked_det were
   removed. Screened raw features are restricted to density-matrix-linear meta-GGA
   quantities: density, grad, kin, and lapl. Direct second_order=True calls now
   raise an error in favor of gen_response().

10. Correctness and performance coverage is substantially expanded

    New tests compare dense and screened CPU/GPU results for RKS, UKS, the XC potential,
    response/HVP, sparse and empty AO blocks, and nuclear gradients. Benchmark
    infrastructure measures runtime and peak memory, with pytest-benchmark,
    pytest-timeout, Memray support, and dedicated benchmark, GPU, and profiling markers.

[Copilot]

Pull request overview

This PR introduces global spatial AO screening for Skala’s PySCF and GPU4PySCF integrations, evaluating AO-derived features in spatially grouped blocks (with atom-major restoration) and using chunk-local autograd plus blockwise VJP/JVP to bound memory while preserving a dense fallback.

Changes:

• Add a global screened feature-evaluation path (spatial grid sorting + blockwise AO screening) for first- and second-order XC / response computations.
• Update memory estimation to account for full-grid raw-feature/cotangent buffers and add targeted unit tests for the estimators.
• Expand CPU/GPU test coverage for screening decisions, ordering/permutations, dense-vs-screened equivalence, and profiling/benchmark harnesses (with new pytest markers/deps).

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_memory_estimators.py	Adds unit tests for the new global raw-feature buffer memory estimator and reserved-memory chunk sizing.
tests/test_gpu4pyscf_gradients.py	Extends GPU gradient tests to compare CPU/GPU and dense/screened routes against stored references.
tests/test_gpu4pyscf_classes.py	Marks GPU4PySCF class tests with a GPU marker and adjusts import ordering for CUDA gating.
tests/test_gpu4pyscf_ao_screening.py	Adds GPU-focused AO screening tests covering sorting, response equivalence, masking/scatter, and CPU-vs-GPU validation.
tests/test_ao_screening.py	Adds CPU AO screening tests for screening decisions, permutations, masking behavior, and screened response correctness.
tests/test_ao_screening_benchmark.py	Adds benchmark/profiling tests to compare screened vs dense runtime and peak memory (CPU + CUDA).
src/skala/pyscf/numint.py	Switches atom-chunked path to global screened evaluation, adds _should_screen_aos, and updates response/HVP logic accordingly.
src/skala/pyscf/memory_estimators.py	Adds reserved-memory support and a new estimator for global raw-feature buffer storage.
src/skala/pyscf/gradients.py	Minor dict-comprehension cleanup while building gradient feature dictionaries.
src/skala/pyscf/features.py	Adds spatial grid ordering, caching, AO screening block loop helpers, and global screened feature evaluation machinery.
src/skala/pyscf/backend.py	Reorders/cleans __all__ exports.
src/skala/gpu4pyscf/gradients.py	Minor dict-comprehension cleanup mirroring the CPU gradients module.
src/skala/functional/utils/irreps.py	Reorders __slots__ entries.
src/skala/functional/load.py	Simplifies all(...) checks (removes unnecessary list creation).
src/skala/functional/init.py	Reorders __all__ exports.
src/skala/ase/init.py	Removes a now-redundant # noqa: F401 from the Skala export.
pyproject.toml	Adds pytest timeout/addopts/markers and new dev dependencies (benchmark/timeout/memray).
environment-gpu.yml	Adds pytest-benchmark and pytest-timeout to the GPU development/test environment.
environment-cpu.yml	Adds pytest-benchmark and pytest-timeout to the CPU development/test environment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 30 out of 30 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

src/skala/pyscf/xc_integrator.py:292

• sorted_raw_features (and the model-chunk raw features derived from it) are computed from dm0.double(), but screened_feature_jvp() is called with dm0/dm1 directly. If dm0/dm1 are not float64 this will produce a tangent with a different dtype than local_raw_features, and the subsequent torch.autograd.grad(..., grad_outputs=atom_major_tangent[...]) will error due to dtype mismatch.

            atom_major_tangent = screened_feature_jvp(
                dm0,
                dm1,
                mol,
                spatial_grid_layout,