Improve Vxc Screening by JensWehner · Pull Request #100 · microsoft/skala

Jens (JensWehner) · 2026-08-04T07:10:45Z

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.

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.
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.
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.
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.
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.
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.
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.
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.
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(). The vjp for the pullback from model outputs to echange correlation potential now uses an analytical form instead of pytorch autograd.
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.

Comparison of main branch and this branch vs fully dense evaluation
runtime:

max memory:

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.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

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,

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 36 out of 36 changed files in this pull request and generated no new comments.

Suppressed comments (7)

tests/test_xc_integrator.py:4

tests/utils.py is not importable as a top-level utils module (there is no utils.py at the repository root), so this import will fail under pytest. Import the helper from the tests namespace instead.
tests/test_ao_screening.py:7
tests/utils.py is not importable as a top-level utils module (there is no utils.py at the repository root), so this import will fail under pytest. Import the helper from the tests namespace instead.

from utils import QuadraticFunctional, patch_ao_screening

tests/test_gpu4pyscf_ao_screening.py:26

tests/utils.py is not importable as a top-level utils module (there is no utils.py at the repository root), so this import will fail under pytest. Import the helper from the tests namespace instead.

from utils import QuadraticFunctional, patch_ao_screening  # noqa: E402

tests/test_ao_screening_benchmark.py:17

tests/utils.py is not importable as a top-level utils module (there is no utils.py at the repository root), so this import will fail under pytest. Import the helper from the tests namespace instead.

from utils import patch_ao_screening

tests/test_gpu4pyscf_gradients.py:26

tests/utils.py is not importable as a top-level utils module (there is no utils.py at the repository root), so this import will fail under pytest. Import the helper from the tests namespace instead.
src/skala/functional/density.py:20
IMMUTABLES is typed as frozenset[str] but it contains Feature enum values. This will confuse type-checking (and IDEs) and can hide real key-type mistakes when working with FeatureMap.

IMMUTABLES: frozenset[str] = frozenset([Feature.GRID_COORDS, Feature.GRID_WEIGHTS])

src/skala/ase/init.py:11

This import is only used to re-export Skala via __all__, so Ruff/Pyflakes will likely flag it as F401 (imported but unused) without an explicit ignore (consistent with src/skala/foundry/__init__.py).

from skala.ase.calculator import Skala

…atrix

Jens (JensWehner) · 2026-08-10T07:13:44Z

I would ignore these files, in this folder, I am not sure if they should be committed.

Sebastian Ehlert (awvwgk) · 2026-08-10T08:43:41Z

-    - name: Run mypy
-      run: mypy .
-      shell: micromamba-shell {0}


Can we keep mypy here to run it in our production environment rather than the incomplete pre-commit env?

yes, the precommit mypy at the moment uses the local mypy version so it should use the full environment.

Sebastian Ehlert (awvwgk) · 2026-08-10T08:44:11Z

Do we want to keep this large notebooks as part of the repo?

probably not, I am not sure how much of the benchmarking code we should keep.

Sebastian Ehlert (awvwgk) · 2026-08-10T08:56:49Z

+        spatial_grid_layout = cast(
+            SpatialGridLayout | None,
+            grid_state.get("_skala_spatial_grid_layout"),
+        )
+        if spatial_grid_layout is not None:
+            return spatial_grid_layout


Is this caching save in case we use a GradientScanner as in the ASE calculator, where we reuse the Kohn-Sham object and the Grids object but just reset them with new molecule geometries?

Good catch, I did not think of grid object reuse, so I would have to hook that into the grid_object reset. That is annoying. Why would you ever reset a grid object, sigh.

then I guess I have to create two classes as well which inherit from the respective Gridobjects.

Sebastian Ehlert (awvwgk) · 2026-08-10T08:59:46Z

+        dm_eval = dm.double()
+        electron_count = torch.zeros(2, device=self.device, dtype=dm_eval.dtype)
+        energy = torch.tensor(0.0, device=self.device, dtype=dm_eval.dtype)
+        feature_function = feature_math.MGGAFeatureFunction(self.feature_spec)


Should we add update the feature spec here as in _integrate_dense with set(self.feature_spec.names) | {Feature.DENSITY, Feature.GRID_WEIGHTS}?

yes, good point.

Sebastian Ehlert (awvwgk) · 2026-08-10T09:40:35Z

+    blksize = _resolve_ao_block_size(mol, feature_function, block_size, max_memory, gpu)
+
+    if blksize is not None and blksize >= grids.weights.shape[0]:
+        features = evaluate_full_grid(dm.double(), mol, grids.coords, feature_function)


Should this pass through the gpu option?

Suggested change

features = evaluate_full_grid(dm.double(), mol, grids.coords, feature_function)

features = evaluate_full_grid(dm.double(), mol, grids.coords, feature_function, gpu)

Thanks, ohh yes.

Sebastian Ehlert (awvwgk)

Left some comments, overall looks good to me.

One note, the gradient calculation is currently not affected by this update is it? We would still use the dense path there, so the ASE calculator which goes via the GradientScanner would see the same peek memory / compute usage as before?

Derk Kooi (DerkKooi)

Further simplifications are possible on the custom autodiff side because of the assumption of features linear in the density matrix. I have described the modifications in some detail in this review.

Derk Kooi (DerkKooi) · 2026-08-10T11:13:24Z

+    return loop_type(dm, mol, grids, feature_function, blksize)
+
+
+class ChunkEvalForward(Function):


If we are assuming that our features are only linear in the DM then we can (and I think should) go one step further in the simplification of the two custom autograd functions (both ChunkEvalForward and ChunkEvalBackward). I've outlined the simplifications in separate comments.

Derk Kooi (DerkKooi) · 2026-08-10T11:15:48Z

+            context.blksize,
+            context.compile_feature_function,
+            context.gpu,
+            *vectors_jvp,


vectors_jvp can be done away with everywhere, since any jvp of a linear function of dm (func(dm)) is just that same linear function evaluated on the v of the jvp (func(v)).

Derk Kooi (DerkKooi) · 2026-08-10T11:16:51Z

+                device=ctx.dm.device,
+                dtype=ctx.dm.dtype,
+            )
+        vector_tangent = grad_inputs[7] if ctx.vectors_jvp else grad_inputs[0]


this just becomes vector_tangent=grad_inputs[0]

Derk Kooi (DerkKooi) · 2026-08-10T11:17:22Z

+                dtype=ctx.dm.dtype,
+            )
+        return ChunkEvalForward.apply(
+            ctx.dm,


and here replace ctx.dm by vector_tangent.

Derk Kooi (DerkKooi) · 2026-08-10T11:18:36Z

+        if ctx.vectors_jvp:
+            dm_grad = ctx.dm * 0
+        else:
+            dm_grad = ChunkEvalBackward.apply(


ChunkEvalBackward doesn't need dm, so feature_cotangent can be added as the first argument instead.

Derk Kooi (DerkKooi) · 2026-08-10T11:26:34Z

+            *vectors_jvp,
+        ) = inputs
+        context.vectors_jvp = tuple(vectors_jvp)
+        ctx.save_for_backward(context.dm)


Linear function means we don't have to save the dm for the backward.

Derk Kooi (DerkKooi) · 2026-08-10T11:27:02Z

+            context.gpu,
+            _feature_cotangent,
+        ) = inputs
+        ctx.save_for_backward(context.dm)


Don't need to save the dm or the _feature_cotangent for the backward.

Derk Kooi (DerkKooi) · 2026-08-10T11:27:22Z

+    def jvp(
+        ctx: _ChunkEvalBackwardContext, *grad_inputs: torch.Tensor | None
+    ) -> torch.Tensor:
+        feature_cotangent_tangent = grad_inputs[7]


feature_cotangent_tangent = grad_inputs[0]

Derk Kooi (DerkKooi) · 2026-08-10T11:28:11Z

+        if feature_cotangent_tangent is None:
+            return torch.zeros_like(ctx.dm)
+        return ChunkEvalBackward.apply(
+            ctx.dm,


feature_cotangent_tangent as first argument.

Derk Kooi (DerkKooi) · 2026-08-10T11:29:43Z

+        grads: list[Tensor | None] = [ctx.dm * 0]
+        grads += [None] * 6
+        grads.append(
+            ChunkEvalForward.apply(


First argument is grad_outputs[0], check again ctx.requires_grad[0].

Jens (JensWehner) · 2026-08-10T13:08:01Z

Left some comments, overall looks good to me.

One note, the gradient calculation is currently not affected by this update is it? We would still use the dense path there, so the ASE calculator which goes via the GradientScanner would see the same peek memory / compute usage as before?

Yes, correct. I did not look at gradients at all so far.

Jens (JensWehner) · 2026-08-12T09:15:56Z

So talking to Thijs Vogels (@tvogels) to make the review much easier I split this up into multiple PRs.

#102 #103 #105 #106

Jens (JensWehner) requested a lite review from Copilot August 4, 2026 07:10

Copilot started reviewing on behalf of Jens (JensWehner) August 4, 2026 07:11 View session

Copilot AI reviewed Aug 4, 2026

View reviewed changes

Comment thread tests/test_ao_screening.py Outdated

Comment thread tests/test_gpu4pyscf_ao_screening.py Outdated

Comment thread src/skala/pyscf/numint.py Outdated

Jens (JensWehner) changed the title ~~Imropve Vxc Screening~~ Improve Vxc Screening Aug 4, 2026

Jens (JensWehner) and others added 19 commits August 4, 2026 10:03

implement screening

766234a

add docstrings

658c264

add failing test for consistency

eaf5d1c

add docstring and use dense eval

a4dad08

work on shell screening

3fe4371

add nuclear gradient test

abee701

add memory benchmark

9428b67

remove chunked case

8ad6f33

Apply suggestions from code review

0343e37

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

refactor tests

9885e84

fix consistent screened response memory policy

a02c507

refactor remove unsupported kinetic tensor features

0fc36be

refactor extract chunk evaluation helpers

a425a10

refactor centralize feature evaluation policy

a6676cf

refactor isolate feature mathematics

5783c7a

refactor isolate ao autograd evaluation

d299389

refactor isolate screened feature evaluation

683e350

refactor make screened grid cache explicit

92ef053

refactor separate xc integration from pyscf adapter

60385d5

Jens (JensWehner) force-pushed the screening branch from 9aca3f7 to 60385d5 Compare August 5, 2026 14:35

make functions private

1307736

Jens (JensWehner) requested a lite review from Copilot August 6, 2026 15:04

Copilot started reviewing on behalf of Jens (JensWehner) August 6, 2026 15:05 View session

Copilot AI reviewed Aug 6, 2026

View reviewed changes

Comment thread src/skala/pyscf/xc_integrator.py

Jens (JensWehner) and others added 2 commits August 6, 2026 17:16

Potential fix for pull request finding

428c393

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

fix screening in the cpu case

5c0d8db

Jens (JensWehner) added 3 commits August 7, 2026 13:44

fix enum properly

55fe903

make tests nicer

f5fa1f2

make tests nicer

28cf7ef

Jens (JensWehner) requested a lite review from Copilot August 7, 2026 11:59

Copilot started reviewing on behalf of Jens (JensWehner) August 7, 2026 12:01 View session

Copilot AI reviewed Aug 7, 2026

View reviewed changes

Jens (JensWehner) added 9 commits August 7, 2026 14:18

fix doc tests

8a3b114

clean notebooks

740085d

commit notebooks for analysis

c0a4a9d

improve model chunking

5114d3b

use analytic vjp for backpropagation of ml model outputs to density m…

26fa246

…atrix

small optimisation

af88a1c

use newer pyscf version in test

59b020b

pin versions

01a42db

raise error bound

130bced

Jens (JensWehner) commented Aug 10, 2026

View reviewed changes

Sebastian Ehlert (awvwgk) reviewed Aug 10, 2026

View reviewed changes

Sebastian Ehlert (awvwgk) previously approved these changes Aug 10, 2026

View reviewed changes

Derk Kooi (DerkKooi) reviewed Aug 10, 2026

View reviewed changes

Jens (JensWehner) added 2 commits August 10, 2026 16:47

Make linear model simpler and add test

0308eb8

require special skala grids

c56a14a

Jens (JensWehner) dismissed Sebastian Ehlert (awvwgk)’s stale review via c56a14a August 10, 2026 15:37

add pyscf 2.9 fixes

eddd046

Jens (JensWehner) closed this Aug 12, 2026

	features = evaluate_full_grid(dm.double(), mol, grids.coords, feature_function)
	features = evaluate_full_grid(dm.double(), mol, grids.coords, feature_function, gpu)

		return loop_type(dm, mol, grids, feature_function, blksize)


		class ChunkEvalForward(Function):

Conversation

Jens (JensWehner) commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Branch changes compared with main

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Sebastian Ehlert (awvwgk) left a comment

Choose a reason for hiding this comment

Uh oh!

Derk Kooi (DerkKooi) left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jens (JensWehner) commented Aug 10, 2026

Uh oh!

Jens (JensWehner) commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Jens (JensWehner) commented Aug 4, 2026 •

edited

Loading

Jens (JensWehner) commented Aug 12, 2026 •

edited

Loading