use model chunking on all paths

[Jens (JensWehner)]

This PR extends memory-aware model chunking from screened first-order XC evaluation to every relevant model execution path:

Applies model chunking to dense and screened XC energy/potential evaluation.
Applies chunking to Hessian-response model evaluations.
Applies chunking to CPU and GPU explicit nuclear-gradient VJPs.
Retains full-grid evaluation for functionals that do not support spatial decomposition.
It also simplifies the feature and integrator architecture:

Introduces ModelFeaturePlan to separate packed AO requirements from model-visible features.
Uses FeatureSpec.supports_spatial_decomposition as the single chunking capability signal.
Consolidates dense and screened integration and response paths.

[Copilot]

🟡 Changes recommended

A critical gradient regression and two moderate compatibility/feature-handling issues remain unresolved.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Extends memory-aware model chunking across XC integration, response, and CPU/GPU nuclear-gradient paths.

Changes:

• Unifies dense and screened feature evaluation.
• Adds feature planning and chunked gradient assembly.
• Updates tests and examples for the new APIs.

File summaries

File	Review
skala/tests/test_xc_integrator.py	Updates integrator mocks and expectations.
skala/tests/test_pyscf_gradients.py	Migrates CPU gradient tests.
skala/tests/test_model_chunking.py	Tests chunked feature gradients.
skala/tests/test_gpu4pyscf_gradients.py	Migrates GPU gradient tests.
skala/tests/test_evaluation.py	Tests feature-spec unions.
skala/tests/test_ao_screening.py	Expands dense and screened coverage.
skala/src/skala/pyscf/xc_integrator.py	Unifies chunked integration and response paths.
skala/src/skala/pyscf/model_chunking.py	Adds feature plans and gradient chunking. Critical: disconnected or constant outputs can now fail instead of returning zero gradients. Moderate: bound-shape-only specifications can raise KeyError.
skala/src/skala/pyscf/gradients.py	Uses chunked CPU gradient evaluation.
skala/src/skala/pyscf/features.py	Removes legacy feature generation. Moderate: removing public generate_features breaks downstream imports without compatibility handling.
skala/src/skala/pyscf/evaluation.py	Adds FeatureSpec unions.
skala/src/skala/pyscf/ao_evaluation.py	Returns packed raw AO features.
skala/src/skala/gpu4pyscf/gradients.py	Uses chunked GPU gradient evaluation.
model/examples/cpp/cpp_integration/prepare_inputs.py	Migrates example input preparation.

Review details

Suppressed comments (1)

skala/src/skala/pyscf/model_chunking.py:420

• The chunked branch has the same regression for constant energies or unused declared inputs: torch.autograd.grad raises rather than returning the zero feature derivatives that the old feature_derivatives path guaranteed. This is especially relevant when nuc_grad_feats is an explicit superset. Add the constant guard plus allow_unused=True and zero-fill disconnected inputs.

        energy_chunk = functional.get_exc(chunk_features)
        local_gradients = torch.autograd.grad(energy_chunk, tuple(chunk_inputs))

• Files reviewed: 14/14 changed files
• Comments generated: 3
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

The CPU nuclear-gradient path does not honor configured memory limits, allowing chunks to exceed the requested budget.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 19/19 changed files
• Comments generated: 1
• Review effort level: Balanced